Support launch SparkProcessBuilder with keytab and principal (#269)

* Support launch SparkProcessBuilder with keytab and principal * fix tests
apache · Jan 13, 2021 · 33bdbbd · 33bdbbd
1 parent 8a99795
commit 33bdbbd
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 6 deletions.
diff --git a/...ql-engine/src/test/scala/org/apache/kyuubi/engine/spark/SparkSQLEngineListenerSuite.scala b/...ql-engine/src/test/scala/org/apache/kyuubi/engine/spark/SparkSQLEngineListenerSuite.scala
@@ -33,7 +33,7 @@ class SparkSQLEngineListenerSuite extends KyuubiFunSuite {
 
   test("application end") {
     val spark = SparkSession
-      .builder().master("local").getOrCreate()
+      .builder().master("local").config("spark.ui.port", "0").getOrCreate()
 
     val engine = new SparkSQLEngine(spark)
     engine.initialize(KyuubiConf())

diff --git a/kyuubi-main/pom.xml b/kyuubi-main/pom.xml
@@ -65,6 +65,16 @@
             <artifactId>hive-jdbc</artifactId>
             <scope>test</scope>
         </dependency>
+
+        <dependency>
+            <groupId>org.apache.hadoop</groupId>
+            <artifactId>hadoop-minikdc</artifactId>
+        </dependency>
+
+        <dependency>
+            <groupId>org.apache.directory.server</groupId>
+            <artifactId>apacheds-service</artifactId>
+        </dependency>
     </dependencies>
 
     <build>

diff --git a/kyuubi-main/src/main/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilder.scala b/kyuubi-main/src/main/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilder.scala
@@ -17,10 +17,13 @@
 
 package org.apache.kyuubi.engine.spark
 
+import java.io.IOException
 import java.nio.file.{Files, Path, Paths}
 
 import scala.collection.mutable.ArrayBuffer
 
+import org.apache.hadoop.security.UserGroupInformation
+
 import org.apache.kyuubi._
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf.ENGINE_SPARK_MAIN_RESOURCE
@@ -108,8 +111,11 @@ class SparkProcessBuilder(
       buffer += CONF
       buffer += s"$k=$v"
     }
-    buffer += PROXY_USER
-    buffer += proxyUser
+    // iff the keytab is specified, PROXY_USER is not supported
+    if (!useKeytab()) {
+      buffer += PROXY_USER
+      buffer += proxyUser
+    }
 
     mainResource.foreach { r => buffer += r }
 
@@ -119,11 +125,32 @@ class SparkProcessBuilder(
   override def toString: String = commands.mkString(" ")
 
   override protected def module: String = "kyuubi-spark-sql-engine"
+
+  private def useKeytab(): Boolean = {
+    val principal = conf.get(PRINCIPAL)
+    val keytab = conf.get(KEYTAB)
+    if (principal.isEmpty || keytab.isEmpty) {
+      false
+    } else {
+      try {
+        val ugi = UserGroupInformation
+          .loginUserFromKeytabAndReturnUGI(principal.get, keytab.get)
+        ugi.getShortUserName == proxyUser
+      } catch {
+        case e: IOException =>
+          error(s"Failed to login for ${principal.get}", e)
+          false
+      }
+    }
+  }
 }
 
 object SparkProcessBuilder {
+  final val APP_KEY = "spark.app.name"
+
   private final val CONF = "--conf"
   private final val CLASS = "--class"
   private final val PROXY_USER = "--proxy-user"
-  final val APP_KEY = "spark.app.name"
+  private final val PRINCIPAL = "spark.kerberos.principal"
+  private final val KEYTAB = "spark.kerberos.keytab"
 }
diff --git a/kyuubi-main/src/test/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilderSuite.scala b/kyuubi-main/src/test/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilderSuite.scala
@@ -20,11 +20,12 @@ package org.apache.kyuubi.engine.spark
 import java.nio.file.{Files, Paths}
 import java.util.concurrent.TimeUnit
 
-import org.apache.kyuubi.{KyuubiFunSuite, KyuubiSQLException}
+import org.apache.kyuubi.{KerberizedTestHelper, KyuubiSQLException, Utils}
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf._
+import org.apache.kyuubi.service.ServiceUtils
 
-class SparkProcessBuilderSuite extends KyuubiFunSuite {
+class SparkProcessBuilderSuite extends KerberizedTestHelper {
   private val conf = KyuubiConf()
     .set(EMBEDDED_ZK_PORT, 5555)
     .set(EMBEDDED_ZK_TEMP_DIR, "spark_process_test")
@@ -71,4 +72,31 @@ class SparkProcessBuilderSuite extends KyuubiFunSuite {
       error1.getMessage.contains("Caused by: org.apache.hadoop.hive.ql.metadata.HiveException:"))
   }
 
+  test("proxy user or keytab") {
+    val b1 = new SparkProcessBuilder("kentyao", conf)
+    assert(b1.toString.contains("--proxy-user kentyao"))
+
+    val conf1 = conf ++ Map("spark.kerberos.principal" -> testPrincipal)
+    val b2 = new SparkProcessBuilder("kentyao", conf1)
+    assert(b2.toString.contains("--proxy-user kentyao"))
+
+    val conf2 = conf ++ Map("spark.kerberos.keytab" -> testKeytab)
+    val b3 = new SparkProcessBuilder("kentyao", conf2)
+    assert(b3.toString.contains("--proxy-user kentyao"))
+
+    tryWithSecurityEnabled {
+      val conf3 = conf ++ Map("spark.kerberos.principal" -> testPrincipal,
+        "spark.kerberos.keytab" -> "testKeytab")
+      val b4 = new SparkProcessBuilder(Utils.currentUser, conf3)
+      assert(b4.toString.contains(s"--proxy-user ${Utils.currentUser}"))
+
+      val conf4 = conf ++ Map("spark.kerberos.principal" -> testPrincipal,
+        "spark.kerberos.keytab" -> testKeytab)
+      val b5 = new SparkProcessBuilder("kentyao", conf4)
+      assert(b5.toString.contains("--proxy-user kentyao"))
+
+      val b6 = new SparkProcessBuilder(ServiceUtils.getShortName(testPrincipal), conf4)
+      assert(!b6.toString.contains("--proxy-user kentyao"))
+    }
+  }
 }