[WIP] Add custom truststore support for LDAP SSL

kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala

@@ -1230,6 +1230,42 @@ object KyuubiConf {
       .stringConf
       .createOptional
 
+  val AUTHENTICATION_LDAP_SSL_ENABLE: ConfigEntry[Boolean] =
+    buildConf("kyuubi.authentication.ldap.ssl.enable")
+      .doc("Set this to true for using SSL encryption when connecting to LDAP servers.")
+      .version("1.13.0")
+      .audience(SERVER)
+      .immutable
+      .booleanConf
+      .createWithDefault(false)
+
+  val AUTHENTICATION_LDAP_SSL_TRUSTSTORE_PATH: OptionalConfigEntry[String] =
+    buildConf("kyuubi.authentication.ldap.ssl.truststore.path")
+      .doc("The truststore path used for SSL connections to LDAP servers.")
+      .version("1.13.0")
+      .audience(SERVER)
+      .immutable
+      .stringConf
+      .createOptional
+
+  val AUTHENTICATION_LDAP_SSL_TRUSTSTORE_PASSWORD: OptionalConfigEntry[String] =
+    buildConf("kyuubi.authentication.ldap.ssl.truststore.password")
+      .doc("The truststore password used for SSL connections to LDAP servers.")
+      .version("1.13.0")
+      .audience(SERVER)
+      .immutable
+      .stringConf
+      .createOptional
+
+  val AUTHENTICATION_LDAP_SSL_TRUSTSTORE_TYPE: OptionalConfigEntry[String] =
+    buildConf("kyuubi.authentication.ldap.ssl.truststore.type")
+      .doc("The truststore type used for SSL connections to LDAP servers.")
+      .version("1.13.0")
+      .audience(SERVER)
+      .immutable
+      .stringConf
+      .createOptional
+
   val AUTHENTICATION_JDBC_DRIVER: OptionalConfigEntry[String] =
     buildConf("kyuubi.authentication.jdbc.driver.class")
       .doc("Driver class name for JDBC Authentication Provider.")

kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/ldap/LdapSSLSocketFactory.scala

@@ -0,0 +1,76 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.service.authentication.ldap
+
+import java.io.IOException
+import java.net.{InetAddress, Socket}
+import javax.net.SocketFactory
+import javax.net.ssl.SSLContext
+
+final class LdapSSLSocketFactory private (socketFactory: SocketFactory) extends SocketFactory {
+
+  @throws[IOException]
+  override def createSocket(): Socket = socketFactory.createSocket()
+
+  @throws[IOException]
+  override def createSocket(host: String, port: Int): Socket =
+    socketFactory.createSocket(host, port)
+
+  @throws[IOException]
+  override def createSocket(
+      host: String,
+      port: Int,
+      localHost: InetAddress,
+      localPort: Int): Socket = {
+    socketFactory.createSocket(host, port, localHost, localPort)
+  }
+
+  @throws[IOException]
+  override def createSocket(host: InetAddress, port: Int): Socket =
+    socketFactory.createSocket(host, port)
+
+  @throws[IOException]
+  override def createSocket(
+      address: InetAddress,
+      port: Int,
+      localAddress: InetAddress,
+      localPort: Int): Socket = {
+    socketFactory.createSocket(address, port, localAddress, localPort)
+  }
+}
+
+object LdapSSLSocketFactory {
+
+  private val sslContext = new ThreadLocal[SSLContext]
+
+  def getDefault(): SocketFactory = {
+    val context = sslContext.get()
+    if (context == null) {
+      throw new IllegalStateException("SSLContext was not set for LDAP SSL connection")
+    }
+    new LdapSSLSocketFactory(context.getSocketFactory)
+  }
+
+  def setSSLContextForCurrentThread(context: SSLContext): Unit = {
+    sslContext.set(context)
+  }
+
+  def clearSslContextForCurrentThread(): Unit = {
+    sslContext.remove()
+  }
+}

kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/ldap/LdapSSLUtils.scala

@@ -0,0 +1,119 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.service.authentication.ldap
+
+import java.io.{FileInputStream, IOException}
+import java.security.{GeneralSecurityException, KeyStore}
+import java.security.cert.{CertificateFactory, X509Certificate}
+import javax.net.ssl.{SSLContext, TrustManagerFactory}
+
+import scala.collection.JavaConverters._
+
+private[ldap] object LdapSSLUtils {
+
+  @throws[GeneralSecurityException]
+  @throws[IOException]
+  def createSSLContext(
+      trustStorePath: String,
+      trustStorePassword: String,
+      trustStoreType: String): SSLContext = {
+    val trustStore =
+      if (Option(trustStorePath).exists(_.trim.nonEmpty)) {
+        loadTrustStore(trustStorePath, trustStorePassword, trustStoreType)
+      } else {
+        null
+      }
+
+    val trustManagerFactory =
+      TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm)
+    trustManagerFactory.init(trustStore)
+
+    val sslContext = SSLContext.getInstance("SSL")
+    sslContext.init(null, trustManagerFactory.getTrustManagers, null)
+    sslContext
+  }
+
+  @throws[GeneralSecurityException]
+  @throws[IOException]
+  private def loadTrustStore(
+      trustStorePath: String,
+      trustStorePassword: String,
+      trustStoreType: String): KeyStore = {
+    val certificatesKeyStore =
+      try {
+        loadCertificates(trustStorePath)
+      } catch {
+        case _: GeneralSecurityException | _: IOException => None
+      }
+    certificatesKeyStore
+      .getOrElse(loadKeyStore(trustStorePath, trustStorePassword, trustStoreType))
+  }
+
+  @throws[GeneralSecurityException]
+  @throws[IOException]
+  private def loadKeyStore(
+      trustStorePath: String,
+      trustStorePassword: String,
+      trustStoreType: String): KeyStore = {
+    val trustStore =
+      KeyStore.getInstance(
+        Option(trustStoreType)
+          .filter(_.trim.nonEmpty)
+          .getOrElse(KeyStore.getDefaultType))
+    val in = new FileInputStream(trustStorePath)
+    try {
+      trustStore.load(in, toCharArray(trustStorePassword))
+    } finally {
+      in.close()
+    }
+    trustStore
+  }
+
+  @throws[GeneralSecurityException]
+  @throws[IOException]
+  private def loadCertificates(trustStorePath: String): Option[KeyStore] = {
+    val certificateFactory = CertificateFactory.getInstance("X.509")
+    val certificateChain = {
+      val in = new FileInputStream(trustStorePath)
+      try {
+        val certificates = certificateFactory.generateCertificates(in)
+        certificates.asScala.map(_.asInstanceOf[X509Certificate]).toSeq
+      } finally {
+        in.close()
+      }
+    }
+
+    if (certificateChain.isEmpty) {
+      None
+    } else {
+      val trustStore = KeyStore.getInstance(KeyStore.getDefaultType)
+      trustStore.load(null, null)
+      var index = 1
+      certificateChain.foreach { certificate =>
+        val certificateAlias = s"Certificate_$index";
+        trustStore.setCertificateEntry(certificateAlias, certificate)
+        index += 1
+      }
+      Some(trustStore)
+    }
+  }
+
+  private def toCharArray(password: String): Array[Char] = {
+    if (password == null) null else password.toCharArray
+  }
+}

kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/ldap/LdapSearch.scala

@@ -30,7 +30,10 @@ import org.apache.kyuubi.config.KyuubiConf
  * @param conf Kyuubi configuration
  * @param ctx  Directory service that will be used for the queries.
  */
-class LdapSearch(conf: KyuubiConf, ctx: DirContext) extends DirSearch with Logging {
+class LdapSearch(
+    conf: KyuubiConf,
+    ctx: DirContext,
+    clearSslContextOnClose: Boolean = false) extends DirSearch with Logging {
 
   final private val baseDn = conf.get(KyuubiConf.AUTHENTICATION_LDAP_BASE_DN).orNull
   final private val groupBases: Array[String] =
@@ -51,6 +54,10 @@ class LdapSearch(conf: KyuubiConf, ctx: DirContext) extends DirSearch with Loggi
     catch {
       case e: NamingException =>
         warn("Exception when closing LDAP context:", e)
+    } finally {
+      if (clearSslContextOnClose) {
+        LdapSSLSocketFactory.clearSslContextForCurrentThread()
+      }
     }
   }
 

kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/ldap/LdapSearchFactory.scala

@@ -17,20 +17,24 @@
 
 package org.apache.kyuubi.service.authentication.ldap
 
+import java.io.IOException
+import java.security.{GeneralSecurityException, KeyStore}
 import java.util
 import javax.naming.{Context, NamingException}
 import javax.naming.directory.{DirContext, InitialDirContext}
+import javax.net.ssl.SSLContext
 import javax.security.sasl.AuthenticationException
 
 import org.apache.kyuubi.Logging
 import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf._
 
 class LdapSearchFactory extends DirSearchFactory with Logging {
   @throws[AuthenticationException]
   override def getInstance(conf: KyuubiConf, principal: String, password: String): DirSearch = {
     try {
       val ctx = createDirContext(conf, principal, password)
-      new LdapSearch(conf, ctx)
+      new LdapSearch(conf, ctx, conf.get(AUTHENTICATION_LDAP_SSL_ENABLE))
     } catch {
       case e: NamingException =>
         debug(s"Could not connect to the LDAP Server: Authentication failed for $principal")
@@ -43,14 +47,79 @@ class LdapSearchFactory extends DirSearchFactory with Logging {
       conf: KyuubiConf,
       principal: String,
       password: String): DirContext = {
+    val env = createDirContextEnvironment(conf, principal, password)
+    try {
+      new InitialDirContext(env)
+    } catch {
+      case e: NamingException =>
+        if (env.containsKey(LdapSearchFactory.LDAP_SOCKET_FACTORY)) {
+          LdapSSLSocketFactory.clearSslContextForCurrentThread()
+        }
+        throw e
+    }
+  }
+
+  private[ldap] def createDirContextEnvironment(
+      conf: KyuubiConf,
+      principal: String,
+      password: String): util.Hashtable[String, AnyRef] = {
     val ldapUrl = conf.get(KyuubiConf.AUTHENTICATION_LDAP_URL)
     val env = new util.Hashtable[String, AnyRef]
     ldapUrl.foreach(env.put(Context.PROVIDER_URL, _))
     env.put(Context.INITIAL_CONTEXT_FACTORY, "com.sun.jndi.ldap.LdapCtxFactory")
     env.put(Context.SECURITY_AUTHENTICATION, "simple")
     env.put(Context.SECURITY_PRINCIPAL, principal)
     env.put(Context.SECURITY_CREDENTIALS, password)
+    if (conf.get(AUTHENTICATION_LDAP_SSL_ENABLE)) {
+      configureSSLSocketFactory(conf, env)
+    }
     debug(s"Connecting using principal $principal to ldap server: ${ldapUrl.orNull}")
-    new InitialDirContext(env)
+    env
+  }
+
+  @throws[NamingException]
+  private def configureSSLSocketFactory(
+      conf: KyuubiConf,
+      env: util.Hashtable[String, AnyRef]): Unit = {
+    val trustStorePath = conf.get(AUTHENTICATION_LDAP_SSL_TRUSTSTORE_PATH)
+    val trustStorePassword = conf.get(AUTHENTICATION_LDAP_SSL_TRUSTSTORE_PASSWORD)
+    val trustStoreType = conf.get(AUTHENTICATION_LDAP_SSL_TRUSTSTORE_TYPE)
+
+    if (trustStorePath.isEmpty) {
+      throw new IllegalArgumentException(
+        s"${AUTHENTICATION_LDAP_SSL_TRUSTSTORE_PATH.key} not configured for SSL connection")
+    }
+    if (trustStorePassword.isEmpty) {
+      throw new IllegalArgumentException(
+        s"${AUTHENTICATION_LDAP_SSL_TRUSTSTORE_PASSWORD.key} not configured for SSL connection")
+    }
+
+    try {
+      val sslContext = getSSLContext(
+        trustStorePath.get,
+        trustStorePassword.get,
+        trustStoreType.getOrElse(KeyStore.getDefaultType))
+      LdapSSLSocketFactory.setSSLContextForCurrentThread(sslContext)
+      env.put(LdapSearchFactory.LDAP_SOCKET_FACTORY, classOf[LdapSSLSocketFactory].getName)
+    } catch {
+      case e @ (_: GeneralSecurityException | _: IOException) =>
+        val namingException = new NamingException("Failed to configure LDAP SSL context")
+        namingException.initCause(e)
+        throw namingException
+    }
+  }
+
+  private def getSSLContext(
+      trustStorePath: String,
+      trustStorePassword: String,
+      trustStoreType: String): SSLContext = {
+    LdapSSLUtils.createSSLContext(
+      trustStorePath,
+      trustStorePassword,
+      trustStoreType)
   }
 }
+
+object LdapSearchFactory {
+  private val LDAP_SOCKET_FACTORY = "java.naming.ldap.factory.socket"
+}