From 7d4d81fff178dc8025ef72003606e4df860c75f5 Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Mon, 28 Jan 2019 00:34:51 -0500
Subject: [PATCH] Improve HandshakeSessionHandler#cleanVhost()

---
 .../connection/client/HandshakeSessionHandler.java   |  2 +-
 .../client/HandshakeSessionHandlerTest.java          | 12 ++++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

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 27004bc2d..fdc16fb87 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
@@ -139,7 +139,7 @@ public class HandshakeSessionHandler implements MinecraftSessionHandler {
 
     // If we connect through an SRV record, there will be a period at the end (DNS usually elides
     // this ending octet).
-    if (cleaned.endsWith(".")) {
+    if (!cleaned.isEmpty() && cleaned.charAt(cleaned.length() - 1) == '.') {
       cleaned = cleaned.substring(0, cleaned.length() - 1);
     }
     return cleaned;
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
index b57e344e1..c89fc9b87 100644
--- a/proxy/src/test/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandlerTest.java
+++ b/proxy/src/test/java/com/velocitypowered/proxy/connection/client/HandshakeSessionHandlerTest.java
@@ -11,20 +11,32 @@ class HandshakeSessionHandlerTest {
   @Test
   void cleanVhostHandlesGoodHostname() {
     assertEquals("localhost", cleanVhost("localhost"));
+    assertEquals("mc.example.com", cleanVhost("mc.example.com"));
   }
 
   @Test
   void cleanVhostHandlesTrailingOctet() {
     assertEquals("localhost", cleanVhost("localhost."));
+    assertEquals("mc.example.com", cleanVhost("mc.example.com."));
   }
 
   @Test
   void cleanVhostHandlesForge() {
     assertEquals("localhost", cleanVhost("localhost" + HANDSHAKE_HOSTNAME_TOKEN));
+    assertEquals("mc.example.com", cleanVhost("mc.example.com" + HANDSHAKE_HOSTNAME_TOKEN));
   }
 
   @Test
   void cleanVhostHandlesOctetsAndForge() {
     assertEquals("localhost", cleanVhost("localhost." + HANDSHAKE_HOSTNAME_TOKEN));
+    assertEquals("mc.example.com", cleanVhost("mc.example.com." + HANDSHAKE_HOSTNAME_TOKEN));
+  }
+
+  @Test
+  void cleanVhostHandlesEmptyHostnames() {
+    assertEquals("", cleanVhost(""));
+    assertEquals("", cleanVhost(HANDSHAKE_HOSTNAME_TOKEN));
+    assertEquals("", cleanVhost("."));
+    assertEquals("", cleanVhost("." + HANDSHAKE_HOSTNAME_TOKEN));
   }
 }
\ No newline at end of file