[KYUUBI #7348] Separate the principal and keytab used for ZooKeeper authentication between the service and the engine

kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/HighAvailabilityConf.scala

@@ -106,12 +106,27 @@ object HighAvailabilityConf {
       .version("1.3.2")
       .fallbackConf(KyuubiConf.SERVER_PRINCIPAL)
 
+  val HA_ZK_ENGINE_AUTH_PRINCIPAL: ConfigEntry[Option[String]] =
+    buildConf("kyuubi.ha.zookeeper.engine.auth.principal")
+      .doc("Kerberos principal name that is used for the engine's ZooKeeper authentication. " +
+        "It fallback to `kyuubi.ha.zookeeper.auth.principal`")
+      .version("1.12.0")
+      .fallbackConf(HighAvailabilityConf.HA_ZK_AUTH_PRINCIPAL)
+
   val HA_ZK_AUTH_KEYTAB: ConfigEntry[Option[String]] =
     buildConf("kyuubi.ha.zookeeper.auth.keytab")
       .doc("Location of the Kyuubi server's keytab that is used for ZooKeeper authentication.")
       .version("1.3.2")
       .fallbackConf(KyuubiConf.SERVER_KEYTAB)
 
+  val HA_ZK_ENGINE_AUTH_KEYTAB: ConfigEntry[Option[String]] =
+    buildConf("kyuubi.ha.zookeeper.engine.auth.keytab")
+      .doc("Location of the Kyuubi server's keytab that is used for " +
+        "the engine's ZooKeeper authentication. " +
+        "It fallback to `kyuubi.ha.zookeeper.auth.keytab`")
+      .version("1.12.0")
+      .fallbackConf(HighAvailabilityConf.HA_ZK_AUTH_KEYTAB)
+
   val HA_ZK_AUTH_DIGEST: OptionalConfigEntry[String] =
     buildConf("kyuubi.ha.zookeeper.auth.digest")
       .doc("The digest auth string is used for ZooKeeper authentication, like: username:password.")

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

@@ -28,7 +28,7 @@ import com.google.common.annotations.VisibleForTesting
 import org.apache.hadoop.security.UserGroupInformation
 
 import org.apache.kyuubi.Logging
-import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.{ConfigEntry, KyuubiConf}
 import org.apache.kyuubi.ha.HighAvailabilityConf._
 import org.apache.kyuubi.ha.client.{AuthTypes, RetryPolicies}
 import org.apache.kyuubi.ha.client.RetryPolicies._
@@ -110,53 +110,56 @@ object ZookeeperClientProvider extends Logging {
    */
   @throws[Exception]
   def setUpZooKeeperAuth(conf: KyuubiConf): Unit = {
-    def setupZkAuth(): Unit = (conf.get(HA_ZK_AUTH_PRINCIPAL), getKeyTabFile(conf)) match {
-      case (Some(principal), Some(keytab)) if UserGroupInformation.isSecurityEnabled =>
-        if (!new File(keytab).exists()) {
-          throw new IOException(s"${HA_ZK_AUTH_KEYTAB.key}: $keytab does not exists")
-        }
-        System.setProperty("zookeeper.sasl.clientconfig", "KyuubiZooKeeperClient")
-        conf.get(HA_ZK_AUTH_SERVER_PRINCIPAL).foreach { zkServerPrincipal =>
-          // ZOOKEEPER-1467 allows configuring SPN in client
-          System.setProperty("zookeeper.server.principal", zkServerPrincipal)
-        }
-        val zkClientPrincipal = KyuubiHadoopUtils.getServerPrincipal(principal)
-        val jaasConf = jaasConfigurationCache.computeIfAbsent(
-          (principal, keytab),
-          _ => {
-            // HDFS-16591 makes breaking change on JaasConfiguration
-            DynConstructors.builder()
-              .impl( // Hadoop 3.3.5 and above
-                "org.apache.hadoop.security.authentication.util.JaasConfiguration",
-                classOf[String],
-                classOf[String],
-                classOf[String])
-              .impl( // Hadoop 3.3.4 and previous
-                // scalastyle:off
-                "org.apache.hadoop.security.token.delegation.ZKDelegationTokenSecretManager$JaasConfiguration",
-                // scalastyle:on
-                classOf[String],
-                classOf[String],
-                classOf[String])
-              .build[Configuration]()
-              .newInstance("KyuubiZooKeeperClient", zkClientPrincipal, keytab)
-          })
-        Configuration.setConfiguration(jaasConf)
-      case _ =>
-    }
+    def setupZkAuth(
+        principalConfKey: ConfigEntry[Option[String]],
+        keytabConfKey: ConfigEntry[Option[String]]): Unit =
+      (conf.get(principalConfKey), getKeyTabFile(conf, keytabConfKey)) match {
+        case (Some(principal), Some(keytab)) if UserGroupInformation.isSecurityEnabled =>
+          if (!new File(keytab).exists()) {
+            throw new IOException(s"${keytabConfKey.key}: $keytab does not exists")
+          }
+          System.setProperty("zookeeper.sasl.clientconfig", "KyuubiZooKeeperClient")
+          conf.get(HA_ZK_AUTH_SERVER_PRINCIPAL).foreach { zkServerPrincipal =>
+            // ZOOKEEPER-1467 allows configuring SPN in client
+            System.setProperty("zookeeper.server.principal", zkServerPrincipal)
+          }
+          val zkClientPrincipal = KyuubiHadoopUtils.getServerPrincipal(principal)
+          val jaasConf = jaasConfigurationCache.computeIfAbsent(
+            (principal, keytab),
+            _ => {
+              // HDFS-16591 makes breaking change on JaasConfiguration
+              DynConstructors.builder()
+                .impl( // Hadoop 3.3.5 and above
+                  "org.apache.hadoop.security.authentication.util.JaasConfiguration",
+                  classOf[String],
+                  classOf[String],
+                  classOf[String])
+                .impl( // Hadoop 3.3.4 and previous
+                  // scalastyle:off
+                  "org.apache.hadoop.security.token.delegation.ZKDelegationTokenSecretManager$JaasConfiguration",
+                  // scalastyle:on
+                  classOf[String],
+                  classOf[String],
+                  classOf[String])
+                .build[Configuration]()
+                .newInstance("KyuubiZooKeeperClient", zkClientPrincipal, keytab)
+            })
+          Configuration.setConfiguration(jaasConf)
+        case _ =>
+      }
 
     if (conf.get(HA_ENGINE_REF_ID).isEmpty &&
       AuthTypes.withName(conf.get(HA_ZK_AUTH_TYPE)) == AuthTypes.KERBEROS) {
-      setupZkAuth()
+      setupZkAuth(HA_ZK_AUTH_PRINCIPAL, HA_ZK_AUTH_KEYTAB)
     } else if (conf.get(HA_ENGINE_REF_ID).nonEmpty &&
       AuthTypes.withName(conf.get(HA_ZK_ENGINE_AUTH_TYPE)) == AuthTypes.KERBEROS) {
-      setupZkAuth()
+      setupZkAuth(HA_ZK_ENGINE_AUTH_PRINCIPAL, HA_ZK_ENGINE_AUTH_KEYTAB)
     }
   }
 
   @VisibleForTesting
-  def getKeyTabFile(conf: KyuubiConf): Option[String] = {
-    conf.get(HA_ZK_AUTH_KEYTAB).map { fullPath =>
+  def getKeyTabFile(conf: KyuubiConf, key: ConfigEntry[Option[String]]): Option[String] = {
+    conf.get(key).map { fullPath =>
       val filename = new File(fullPath).getName
       if (new File(filename).exists()) filename else fullPath
     }

kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClientSuite.scala

@@ -113,6 +113,7 @@ abstract class ZookeeperDiscoveryClientSuite extends DiscoveryClientTests
 
   test("set up zookeeper auth") {
     tryWithSecurityEnabled {
+      val loginConf = javax.security.auth.login.Configuration.getConfiguration
       val keytab = File.createTempFile("kentyao", ".keytab")
       val principal = "kentyao/_HOST@apache.org"
 
@@ -133,8 +134,39 @@ abstract class ZookeeperDiscoveryClientSuite extends DiscoveryClientTests
 
       conf.set(HA_ZK_AUTH_KEYTAB.key, s"${keytab.getName}")
       val e = intercept[IOException](setUpZooKeeperAuth(conf))
-      assert(
-        e.getMessage === s"${HA_ZK_AUTH_KEYTAB.key}: ${getKeyTabFile(conf).get} does not exists")
+      val keytabPath = getKeyTabFile(conf, HA_ZK_AUTH_KEYTAB).get
+      assert(e.getMessage === s"${HA_ZK_AUTH_KEYTAB.key}: ${keytabPath} does not exists")
+      javax.security.auth.login.Configuration.setConfiguration(loginConf)
+    }
+  }
+
+  test("set up zookeeper auth for engine") {
+    tryWithSecurityEnabled {
+      val loginConf = javax.security.auth.login.Configuration.getConfiguration
+      val keytab = File.createTempFile("engine", ".keytab")
+      val principal = "engine/_HOST@apache.org"
+
+      conf.set(HA_ZK_ENGINE_AUTH_KEYTAB.key, keytab.getCanonicalPath)
+      conf.set(HA_ZK_ENGINE_AUTH_PRINCIPAL.key, principal)
+      conf.set(HA_ENGINE_REF_ID, "ref")
+      conf.set(HA_ZK_ENGINE_AUTH_TYPE.key, AuthTypes.KERBEROS.toString)
+
+      setUpZooKeeperAuth(conf)
+      val configuration = Configuration.getConfiguration
+      val entries = configuration.getAppConfigurationEntry("KyuubiZooKeeperClient")
+
+      assert(entries.head.getLoginModuleName === "com.sun.security.auth.module.Krb5LoginModule")
+      val options = entries.head.getOptions.asScala.toMap
+
+      val hostname = StringUtils.toLowerCase(InetAddress.getLocalHost.getCanonicalHostName)
+      assert(options("principal") === s"engine/$hostname@apache.org")
+      assert(options("useKeyTab").toString.toBoolean)
+
+      conf.set(HA_ZK_ENGINE_AUTH_KEYTAB.key, s"${keytab.getName}")
+      val e = intercept[IOException](setUpZooKeeperAuth(conf))
+      val keytabPath = getKeyTabFile(conf, HA_ZK_ENGINE_AUTH_KEYTAB).get
+      assert(e.getMessage === s"${HA_ZK_ENGINE_AUTH_KEYTAB.key}: ${keytabPath} does not exists")
+      javax.security.auth.login.Configuration.setConfiguration(loginConf)
     }
   }
 

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilder.scala

@@ -209,7 +209,7 @@ class SparkProcessBuilder(
   }
 
   private def zkAuthKeytabFileConf(sparkConf: Map[String, String]): Map[String, String] = {
-    val zkAuthKeytab = conf.get(HighAvailabilityConf.HA_ZK_AUTH_KEYTAB)
+    val zkAuthKeytab = conf.get(HighAvailabilityConf.HA_ZK_ENGINE_AUTH_KEYTAB)
     if (zkAuthKeytab.isDefined) {
       sparkConf.get(SPARK_FILES) match {
         case Some(files) =>

kyuubi-server/src/test/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilderSuite.scala

@@ -267,8 +267,8 @@ class SparkProcessBuilderSuite extends KerberizedTestHelper with MockitoSugar {
   test("zookeeper kerberos authentication") {
     val conf = KyuubiConf()
     conf.set(HighAvailabilityConf.HA_ZK_ENGINE_AUTH_TYPE.key, AuthTypes.KERBEROS.toString)
-    conf.set(HighAvailabilityConf.HA_ZK_AUTH_KEYTAB.key, testKeytab)
-    conf.set(HighAvailabilityConf.HA_ZK_AUTH_PRINCIPAL.key, testPrincipal)
+    conf.set(HighAvailabilityConf.HA_ZK_ENGINE_AUTH_KEYTAB.key, testKeytab)
+    conf.set(HighAvailabilityConf.HA_ZK_ENGINE_AUTH_PRINCIPAL.key, testPrincipal)
 
     val b1 = new SparkProcessBuilder("test", true, conf)
     assert(b1.toString.contains(s"--conf spark.files=$testKeytab"))