[KYUUBI #7712] Support IPv6 host:port parsing in Kyuubi service discovery

kyuubi-common/src/main/scala/org/apache/kyuubi/engine/deploy/yarn/ApplicationMaster.scala

@@ -32,7 +32,7 @@ import org.apache.hadoop.yarn.security.AMRMTokenIdentifier
 import org.apache.kyuubi.{KyuubiException, Logging, Utils}
 import org.apache.kyuubi.config.{KyuubiConf, KyuubiReservedKeys}
 import org.apache.kyuubi.service.Serverable
-import org.apache.kyuubi.util.KyuubiHadoopUtils
+import org.apache.kyuubi.util.{IPStackUtils, KyuubiHadoopUtils}
 import org.apache.kyuubi.util.command.CommandLineUtils.confKeyValues
 import org.apache.kyuubi.util.reflect.{DynFields, DynMethods}
 
@@ -184,8 +184,8 @@ object ApplicationMaster extends Logging {
   }
 
   private def resolveHostAndPort(connectionUrl: String): (String, Int) = {
-    val strings = connectionUrl.split(":")
-    (strings(0), strings(1).toInt)
+    val hostPort = IPStackUtils.getHostAndPort(connectionUrl)
+    (hostPort.getHostname, hostPort.getPort)
   }
 
   private def cleanupStagingDir(): Unit = {

kyuubi-common/src/main/scala/org/apache/kyuubi/service/TFrontendService.scala

@@ -36,7 +36,7 @@ import org.apache.kyuubi.shaded.hive.service.rpc.thrift._
 import org.apache.kyuubi.shaded.thrift.protocol.TProtocol
 import org.apache.kyuubi.shaded.thrift.server.{ServerContext, TServerEventHandler}
 import org.apache.kyuubi.shaded.thrift.transport.TTransport
-import org.apache.kyuubi.util.{JavaUtils, KyuubiHadoopUtils, NamedThreadFactory}
+import org.apache.kyuubi.util.{IPStackUtils, JavaUtils, KyuubiHadoopUtils, NamedThreadFactory}
 
 /**
  * Apache Thrift based hive-service-rpc base class
@@ -120,7 +120,7 @@ abstract class TFrontendService(name: String)
       case (None, None) => serverAddr.getHostAddress
     }
 
-    host + ":" + actualPort
+    IPStackUtils.concatHostPort(host, actualPort)
   }
 
   protected def getProxyUser(

kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/DiscoveryClient.scala

@@ -19,6 +19,7 @@ package org.apache.kyuubi.ha.client
 
 import org.apache.kyuubi.Logging
 import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.util.IPStackUtils
 
 /**
  * A collection of apis that discovery client need implement.
@@ -201,8 +202,8 @@ object DiscoveryClient {
         maybeInfos("hive.server2.thrift.bind.host"),
         maybeInfos("hive.server2.thrift.port").toInt)
     } else {
-      val strings = instance.split(":")
-      (strings(0), strings(1).toInt)
+      val hostPort = IPStackUtils.getHostAndPort(instance)
+      (hostPort.getHostname, hostPort.getPort)
     }
   }
 }

kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/ServiceNodeInfo.scala

@@ -17,6 +17,8 @@
 
 package org.apache.kyuubi.ha.client
 
+import org.apache.kyuubi.util.IPStackUtils
+
 case class ServiceNodeInfo(
     namespace: String,
     nodeName: String,
@@ -25,5 +27,5 @@ case class ServiceNodeInfo(
     version: Option[String],
     engineRefId: Option[String],
     attributes: Map[String, String] = Map.empty) {
-  def instance: String = s"$host:$port"
+  def instance: String = IPStackUtils.concatHostPort(host, port)
 }

kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClient.scala

@@ -44,6 +44,7 @@ import org.apache.kyuubi.shaded.curator.utils.ZKPaths
 import org.apache.kyuubi.shaded.zookeeper.{CreateMode, KeeperException, WatchedEvent, Watcher}
 import org.apache.kyuubi.shaded.zookeeper.CreateMode.PERSISTENT
 import org.apache.kyuubi.shaded.zookeeper.KeeperException.NodeExistsException
+import org.apache.kyuubi.util.IPStackUtils
 import org.apache.kyuubi.util.ThreadUtils
 
 class ZookeeperDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient {
@@ -316,15 +317,15 @@ class ZookeeperDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient {
     if (!instance.contains(":")) {
       return instance
     }
-    val hostPort = instance.split(":", 2)
+    val hostPort = IPStackUtils.getHostAndPort(instance)
     val confsToPublish = collection.mutable.Map[String, String]()
 
     // Hostname
-    confsToPublish += ("hive.server2.thrift.bind.host" -> hostPort(0))
+    confsToPublish += ("hive.server2.thrift.bind.host" -> hostPort.getHostname)
     // Transport mode
     confsToPublish += ("hive.server2.transport.mode" -> "binary")
     // Transport specific confs
-    confsToPublish += ("hive.server2.thrift.port" -> hostPort(1))
+    confsToPublish += ("hive.server2.thrift.port" -> hostPort.getPort.toString)
     confsToPublish += ("hive.server2.thrift.sasl.qop" -> conf.get(KyuubiConf.SASL_QOP))
     // Auth specific confs
     val authenticationMethod = conf.get(KyuubiConf.AUTHENTICATION_METHOD).mkString(",")

kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DiscoveryClientSuite.scala

@@ -19,7 +19,7 @@ package org.apache.kyuubi.ha.client
 
 import org.apache.kyuubi.KyuubiFunSuite
 
-trait DiscoveryClientSuite extends KyuubiFunSuite {
+class DiscoveryClientSuite extends KyuubiFunSuite {
 
   test("parse host and port from instance string") {
     val host = "127.0.0.1"
@@ -36,5 +36,16 @@ trait DiscoveryClientSuite extends KyuubiFunSuite {
     val (host2, port2) = DiscoveryClient.parseInstanceHostPort(instance2)
     assert(host === host2)
     assert(port === port2)
+
+    // IPv6 address with square brackets
+    val ipv6Host = "fc00:172::1"
+    val (host3, port3) = DiscoveryClient.parseInstanceHostPort(s"[$ipv6Host]:$port")
+    assert(ipv6Host === host3)
+    assert(port === port3)
+
+    // IPv6 address without square brackets
+    val (host4, port4) = DiscoveryClient.parseInstanceHostPort(s"$ipv6Host:$port")
+    assert(ipv6Host === host4)
+    assert(port === port4)
   }
 }

kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/hs2connection/KyuubiConfFileParser.java

@@ -26,6 +26,7 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.*;
+import org.apache.kyuubi.util.IPStackUtils;
 import org.apache.kyuubi.util.JavaUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -150,7 +151,7 @@ private void addHosts(Properties props) throws KyuubiConfFileParseException {
     }
 
     int portNum = getPortNum(thriftMode);
-    props.setProperty("hosts", host + ":" + portNum);
+    props.setProperty("hosts", IPStackUtils.concatHostPort(host, portNum));
   }
 
   private int getPortNum(THRIFT_MODE thriftMode) {

kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiConnection.java

@@ -72,6 +72,7 @@
 import org.apache.kyuubi.shaded.thrift.transport.THttpClient;
 import org.apache.kyuubi.shaded.thrift.transport.TTransport;
 import org.apache.kyuubi.shaded.thrift.transport.TTransportException;
+import org.apache.kyuubi.util.IPStackUtils;
 import org.apache.kyuubi.util.SubjectUtil;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -189,7 +190,7 @@
     for (int numRetries = 0; ; ) {
       try {
         // open the client transport
         openTransport();
         // set up the client
         TCLIService.Iface _client = new TCLIService.Client(new TBinaryProtocol(transport));
         // Wrap the client with a thread-safe proxy to serialize the RPC calls
@@ -198,7 +199,7 @@
         openSession();
         if (!isBeeLineMode) {
           showLaunchEngineLog();
           waitLaunchEngineToComplete();
           executeInitSql();
         }
         break;
@@ -405,7 +406,7 @@
   }
 
   private void openTransport() throws Exception {
     transport = isHttpTransportMode() ? createHttpTransport() : createBinaryTransport();
     if (!transport.isOpen()) {
       transport.open();
     }
@@ -428,14 +429,14 @@
     } else if (!httpPath.startsWith("/")) {
       httpPath = "/" + httpPath;
     }
-    return schemeName + "://" + host + ":" + port + httpPath;
+    return schemeName + "://" + IPStackUtils.concatHostPort(host, port) + httpPath;
   }
 
   private TTransport createHttpTransport() throws SQLException, TTransportException {
     CloseableHttpClient httpClient;
     boolean useSsl = isSslConnection();
     // Create an http client from the configs
     httpClient = getHttpClient(useSsl);
     int maxMessageSize = getMaxMessageSize();
     TConfiguration.Builder tConfBuilder = TConfiguration.custom();
     if (maxMessageSize > 0) {

kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/Utils.java

@@ -30,6 +30,7 @@
 import org.apache.commons.lang3.StringUtils;
 import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TStatus;
 import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TStatusCode;
+import org.apache.kyuubi.util.IPStackUtils;
 import org.apache.kyuubi.util.reflect.DynConstructors;
 import org.apache.kyuubi.util.reflect.DynMethods;
 import org.slf4j.Logger;
@@ -406,7 +407,7 @@ public static JdbcConnectionParams extractURLComponents(String uri, Properties i
         connParams.setPort(port);
       }
       // We check for invalid host, port while configuring connParams with configureConnParams()
-      authorityStr = connParams.getHost() + ":" + connParams.getPort();
+      authorityStr = IPStackUtils.concatHostPort(connParams.getHost(), connParams.getPort());
       LOG.debug("Resolved authority: " + authorityStr);
       uri = uri.replace(dummyAuthorityString, authorityStr);
     }
@@ -419,7 +420,7 @@ public static JdbcConnectionParams extractURLComponents(String uri, Properties i
   static void configureConnParamsFromZooKeeper(JdbcConnectionParams connParams)
       throws ZooKeeperHiveClientException, JdbcUriParseException {
     ZooKeeperHiveClientHelper.configureConnParams(connParams);
-    String authorityStr = connParams.getHost() + ":" + connParams.getPort();
+    String authorityStr = IPStackUtils.concatHostPort(connParams.getHost(), connParams.getPort());
     LOG.debug("Resolved authority: " + authorityStr);
     String jdbcUriString = connParams.getJdbcUriString();
     // Replace ZooKeeper ensemble from the authority component of the JDBC Uri provided by the
@@ -532,8 +533,8 @@ static boolean updateConnParamsFromZooKeeper(JdbcConnectionParams connParams) {
           connParams
               .getJdbcUriString()
               .replace(
-                  oldServerHost + ":" + oldServerPort,
-                  connParams.getHost() + ":" + connParams.getPort()));
+                  IPStackUtils.concatHostPort(oldServerHost, oldServerPort),
+                  IPStackUtils.concatHostPort(connParams.getHost(), connParams.getPort())));
       LOG.info("Selected HiveServer2 instance with uri: " + connParams.getJdbcUriString());
     } catch (ZooKeeperHiveClientException e) {
       LOG.error(e.getMessage());

kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/ZooKeeperHiveClientHelper.java

@@ -30,6 +30,7 @@
 import org.apache.kyuubi.shaded.curator.framework.CuratorFramework;
 import org.apache.kyuubi.shaded.curator.framework.CuratorFrameworkFactory;
 import org.apache.kyuubi.shaded.curator.retry.ExponentialBackoffRetry;
+import org.apache.kyuubi.util.IPStackUtils;
 
 class ZooKeeperHiveClientHelper {
   // Pattern for key1=value1;key2=value2
@@ -96,13 +97,14 @@ private static void updateParamsWithZKServerNode(
     // it must be the server uri added by an older version HS2
     Matcher matcher = kvPattern.matcher(dataStr);
     if (!matcher.find()) {
-      String[] split = dataStr.split(":");
-      if (split.length != 2) {
+      try {
+        IPStackUtils.HostPort hostPort = IPStackUtils.getHostAndPort(dataStr);
+        connParams.setHost(hostPort.getHostname());
+        connParams.setPort(hostPort.getPort());
+      } catch (IllegalArgumentException e) {
         throw new ZooKeeperHiveClientException(
-            "Unable to read HiveServer2 uri from ZooKeeper: " + dataStr);
+            "Unable to parse HiveServer2 uri from ZooKeeper: " + dataStr, e);
       }
-      connParams.setHost(split[0]);
-      connParams.setPort(Integer.parseInt(split[1]));
     } else {
       applyConfs(dataStr, connParams);
     }

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala

@@ -47,7 +47,7 @@ import org.apache.kyuubi.metrics.MetricsSystem
 import org.apache.kyuubi.operation.log.OperationLog
 import org.apache.kyuubi.plugin.GroupProvider
 import org.apache.kyuubi.service.authentication.{AuthTypes, AuthUtils}
-import org.apache.kyuubi.util.JavaUtils
+import org.apache.kyuubi.util.{IPStackUtils, JavaUtils}
 
 /**
  * The description and functionality of an engine at server side
@@ -464,7 +464,7 @@ private[kyuubi] object EngineRef {
       val host = conf.get(FRONTEND_ADVERTISED_HOST)
         .orElse(conf.get(FRONTEND_THRIFT_BINARY_BIND_HOST))
         .getOrElse(JavaUtils.findLocalInetAddress.getHostAddress)
-      s"jdbc:kyuubi://$host:$port/default"
+      s"jdbc:kyuubi://${IPStackUtils.concatHostPort(host, port)}/default"
     }
   }
 }

kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiRestFrontendService.scala

@@ -41,7 +41,7 @@ import org.apache.kyuubi.server.ui.{JettyServer, JettyUtils}
 import org.apache.kyuubi.service.{AbstractFrontendService, Serverable, Service, ServiceUtils}
 import org.apache.kyuubi.service.authentication.{AuthTypes, AuthUtils}
 import org.apache.kyuubi.session.{KyuubiBatchSession, KyuubiSessionManager, SessionHandle}
-import org.apache.kyuubi.util.{JavaUtils, ThreadUtils}
+import org.apache.kyuubi.util.{IPStackUtils, JavaUtils, ThreadUtils}
 import org.apache.kyuubi.util.ThreadUtils.scheduleTolerableRunnableWithFixedDelay
 
 /**
@@ -110,7 +110,7 @@ class KyuubiRestFrontendService(override val serverable: Serverable)
   override def connectionUrl: String = {
     checkInitialized()
     conf.get(FRONTEND_ADVERTISED_HOST) match {
-      case Some(advertisedHost) => s"$advertisedHost:$port"
+      case Some(advertisedHost) => IPStackUtils.concatHostPort(advertisedHost, port)
       case None => server.getServerUri
     }
   }

kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiTrinoFrontendService.scala

@@ -25,7 +25,7 @@ import org.apache.kyuubi.config.KyuubiConf._
 import org.apache.kyuubi.server.trino.api.v1.ApiRootResource
 import org.apache.kyuubi.server.ui.JettyServer
 import org.apache.kyuubi.service.{AbstractFrontendService, Serverable, Service}
-import org.apache.kyuubi.util.JavaUtils
+import org.apache.kyuubi.util.{IPStackUtils, JavaUtils}
 
 /**
  * A frontend service based on RESTful api via HTTP protocol.
@@ -64,7 +64,7 @@ class KyuubiTrinoFrontendService(override val serverable: Serverable)
   override def connectionUrl: String = {
     checkInitialized()
     conf.get(FRONTEND_ADVERTISED_HOST) match {
-      case Some(advertisedHost) => s"$advertisedHost:$port"
+      case Some(advertisedHost) => IPStackUtils.concatHostPort(advertisedHost, port)
       case None => server.getServerUri
     }
   }

kyuubi-server/src/main/scala/org/apache/kyuubi/server/ui/JettyServer.scala

@@ -23,6 +23,7 @@ import org.eclipse.jetty.util.component.LifeCycle
 import org.eclipse.jetty.util.thread.{QueuedThreadPool, ScheduledExecutorScheduler}
 
 import org.apache.kyuubi.Logging
+import org.apache.kyuubi.util.IPStackUtils
 import org.apache.kyuubi.util.JavaUtils
 
 private[kyuubi] class JettyServer(
@@ -37,7 +38,7 @@ private[kyuubi] class JettyServer(
       server.addConnector(connector)
       val localPort = connector.getLocalPort
       require(localPort > 0, "Jetty server port should be positive, but got " + localPort)
-      _serverUri = connector.getHost + ":" + localPort
+      _serverUri = IPStackUtils.concatHostPort(connector.getHost, localPort)
     } catch {
       case e: Exception =>
         stop()
@@ -56,7 +57,7 @@ private[kyuubi] class JettyServer(
 
   @volatile private var _serverUri: String = _
   def getServerUri: String = Option(_serverUri).getOrElse {
-    val uri = connector.getHost + ":" + connector.getLocalPort
+    val uri = IPStackUtils.concatHostPort(connector.getHost, connector.getLocalPort)
     warn("Jetty server is not started yet, returning " + uri)
     uri
   }