From 5bf936d1b96013920964b9048a62079a45b52060 Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Sun, 30 Dec 2018 03:57:05 -0500
Subject: [PATCH] Add unit tests

---
 .../client/HandshakeSessionHandler.java       | 11 +++++--
 .../client/HandshakeSessionHandlerTest.java   | 31 +++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)
 create mode 100644 proxy/src/test/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandlerTest.java

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandler.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandler.java
index b47e23590..67320c66a 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandler.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandler.java
@@ -138,8 +138,15 @@ public class HandshakeSessionHandler implements MinecraftSessionHandler {
     }
   }
 
-  private String cleanVhost(String hostname) {
-    // Clean out any anything after any zero byte
+  /**
+   * Cleans the specified virtual host hostname.
+   *
+   * @param hostname the host name to clean
+   * @return the cleaned hostname
+   */
+  static String cleanVhost(String hostname) {
+    // Clean out any anything after any zero bytes (this includes BungeeCord forwarding and the
+    // legacy Forge handshake indicator).
     String cleaned = hostname;
     int zeroIdx = cleaned.indexOf('\0');
     if (zeroIdx > -1) {
diff --git a/proxy/src/test/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandlerTest.java b/proxy/src/test/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandlerTest.java
new file mode 100644
index 000000000..b6f180c56
--- /dev/null
+++ b/proxy/src/test/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandlerTest.java
@@ -0,0 +1,31 @@
+package com.velocitypowered.proxy.connection.client;
+
+import static com.velocitypowered.proxy.connection.client.HandshakeSessionHandler.cleanVhost;
+import static com.velocitypowered.proxy.connection.forge.legacy.LegacyForgeConstants.*;
+import static org.junit.jupiter.api.Assertions.*;
+
+import com.velocitypowered.proxy.connection.forge.legacy.LegacyForgeConstants;
+import org.junit.jupiter.api.Test;
+
+class HandshakeSessionHandlerTest {
+
+  @Test
+  void cleanVhostHandlesGoodHostname() {
+    assertEquals("localhost", cleanVhost("localhost"));
+  }
+
+  @Test
+  void cleanVhostHandlesTrailingOctet() {
+    assertEquals("localhost", cleanVhost("localhost."));
+  }
+
+  @Test
+  void cleanVhostHandlesForge() {
+    assertEquals("localhost", cleanVhost("localhost" + HANDSHAKE_HOSTNAME_TOKEN));
+  }
+
+  @Test
+  void cleanVhostHandlesOctetsAndForge() {
+    assertEquals("localhost", cleanVhost("localhost." + HANDSHAKE_HOSTNAME_TOKEN));
+  }
+}
\ No newline at end of file