From 51f3db2ce672b82ab27b846d56b62af4de5fc41a Mon Sep 17 00:00:00 2001
From: soullkk <55041925+soullkk@users.noreply.github.com>
Date: Fri, 31 Mar 2023 11:48:10 +0800
Subject: [PATCH] Fix peon errors when executing tasks in ipv6(#13972) (#13995)

---
 .../apache/druid/rpc/ServiceClientImpl.java   | 16 +++++++++--
 .../druid/rpc/ServiceClientImplTest.java      | 28 +++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/server/src/main/java/org/apache/druid/rpc/ServiceClientImpl.java b/server/src/main/java/org/apache/druid/rpc/ServiceClientImpl.java
index 98191b3e132..2146ed5d5be 100644
--- a/server/src/main/java/org/apache/druid/rpc/ServiceClientImpl.java
+++ b/server/src/main/java/org/apache/druid/rpc/ServiceClientImpl.java
@@ -475,6 +475,18 @@ public class ServiceClientImpl implements ServiceClient
     );
   }
 
+  /**
+   * Sanitizes IPv6 address if it has brackets. Eg. host = "[1:2:3:4:5:6:7:8]" will be returned as "1:2:3:4:5:6:7:8"
+   * after this function
+   */
+  static String sanitizeHost(String host)
+  {
+    if (host.charAt(0) == '[') {
+      return host.substring(1, host.length() - 1);
+    }
+    return host;
+  }
+
   /**
    * Returns a {@link ServiceLocation} without a path component, based on a URI.
    */
@@ -488,13 +500,13 @@ public class ServiceClientImpl implements ServiceClient
 
     try {
       final URI uri = new URI(uriString);
-      final String host = uri.getHost();
 
-      if (host == null) {
+      if (uri.getHost() == null) {
         return null;
       }
 
       final String scheme = uri.getScheme();
+      final String host = sanitizeHost(uri.getHost());
 
       if ("http".equals(scheme)) {
         return new ServiceLocation(host, uri.getPort() < 0 ? 80 : uri.getPort(), -1, "");
diff --git a/server/src/test/java/org/apache/druid/rpc/ServiceClientImplTest.java b/server/src/test/java/org/apache/druid/rpc/ServiceClientImplTest.java
index 20487aeac8e..69cb12e423c 100644
--- a/server/src/test/java/org/apache/druid/rpc/ServiceClientImplTest.java
+++ b/server/src/test/java/org/apache/druid/rpc/ServiceClientImplTest.java
@@ -663,6 +663,34 @@ public class ServiceClientImplTest
         new ServiceLocation("1.2.3.4", -1, 443, ""),
         ServiceClientImpl.serviceLocationNoPathFromUri("https://1.2.3.4/foo")
     );
+
+    Assert.assertEquals(
+            new ServiceLocation("1:2:3:4:5:6:7:8", 9999, -1, ""),
+            ServiceClientImpl.serviceLocationNoPathFromUri("http://[1:2:3:4:5:6:7:8]:9999/foo")
+    );
+
+    Assert.assertEquals(
+            new ServiceLocation("1:2:3:4:5:6:7:8", 80, -1, ""),
+            ServiceClientImpl.serviceLocationNoPathFromUri("http://[1:2:3:4:5:6:7:8]/foo")
+    );
+
+    Assert.assertEquals(
+            new ServiceLocation("1:2:3:4:5:6:7:8", -1, 9999, ""),
+            ServiceClientImpl.serviceLocationNoPathFromUri("https://[1:2:3:4:5:6:7:8]:9999/foo")
+    );
+
+    Assert.assertEquals(
+            new ServiceLocation("1:2:3:4:5:6:7:8", -1, 443, ""),
+            ServiceClientImpl.serviceLocationNoPathFromUri("https://[1:2:3:4:5:6:7:8]/foo")
+    );
+  }
+
+  @Test
+  public void test_normalizeHost()
+  {
+    Assert.assertEquals("1:2:3:4:5:6:7:8", ServiceClientImpl.sanitizeHost("[1:2:3:4:5:6:7:8]"));
+    Assert.assertEquals("1:2:3:4:5:6:7:8", ServiceClientImpl.sanitizeHost("1:2:3:4:5:6:7:8"));
+    Assert.assertEquals("1.2.3.4", ServiceClientImpl.sanitizeHost("1.2.3.4"));
   }
 
   @Test