From 87b370ffe79a8b3e749dbecd4cbb1c04cfb26b43 Mon Sep 17 00:00:00 2001
From: Cheng Pan <chengpan@apache.org>
Date: Wed, 17 Jan 2024 21:57:04 +0800
Subject: [PATCH] Always take the first declared SASL/PLAIN auth type

---
 .../org/apache/kyuubi/config/KyuubiConf.scala    |  6 +++---
 .../KyuubiAuthenticationFactory.scala            | 16 ++++++++++------
 .../KyuubiAuthenticationFactorySuite.scala       | 10 +++++-----
 .../kyuubi/server/http/ThriftHttpServlet.scala   |  2 +-
 .../org/apache/kyuubi/RestClientTestHelper.scala |  2 +-
 ...yuubiOperationKerberosAndPlainAuthSuite.scala |  2 +-
 ...tionThriftHttpKerberosAndPlainAuthSuite.scala |  2 +-
 .../server/api/v1/AdminResourceSuite.scala       |  2 +-
 .../server/api/v1/BatchesResourceSuite.scala     |  2 +-
 .../server/rest/client/AdminCtlSuite.scala       |  2 +-
 .../server/rest/client/AdminRestApiSuite.scala   |  2 +-
 11 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala
index cdf6db3e57f..9f075275472 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala
@@ -757,7 +757,7 @@ object KyuubiConf {
       .stringConf
       .createWithDefault("X-Real-IP")
 
-  val AUTHENTICATION_METHOD: ConfigEntry[Set[String]] = buildConf("kyuubi.authentication")
+  val AUTHENTICATION_METHOD: ConfigEntry[Seq[String]] = buildConf("kyuubi.authentication")
     .doc("A comma-separated list of client authentication types." +
       "<ul>" +
       " <li>NOSASL: raw transport.</li>" +
@@ -793,9 +793,9 @@ object KyuubiConf {
     .serverOnly
     .stringConf
     .transformToUpperCase
-    .toSet()
+    .toSequence()
     .checkValues(AuthTypes)
-    .createWithDefault(Set(AuthTypes.NONE.toString))
+    .createWithDefault(Seq(AuthTypes.NONE.toString))
 
   val AUTHENTICATION_CUSTOM_CLASS: OptionalConfigEntry[String] =
     buildConf("kyuubi.authentication.custom.class")
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/KyuubiAuthenticationFactory.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/KyuubiAuthenticationFactory.scala
index 736f8e1e15e..070318bb625 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/KyuubiAuthenticationFactory.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/KyuubiAuthenticationFactory.scala
@@ -37,11 +37,15 @@ import org.apache.kyuubi.shaded.thrift.transport.{TSaslServerTransport, TTranspo
 
 class KyuubiAuthenticationFactory(conf: KyuubiConf, isServer: Boolean = true) extends Logging {
 
-  val authTypes: Set[AuthType] = conf.get(AUTHENTICATION_METHOD).map(AuthTypes.withName)
-  val noSaslEnabled: Boolean = authTypes == Set(NOSASL)
+  val authTypes: Seq[AuthType] = conf.get(AUTHENTICATION_METHOD).map(AuthTypes.withName)
+  val saslDisabled: Boolean = authTypes == Seq(NOSASL)
   val kerberosEnabled: Boolean = authTypes.contains(KERBEROS)
-  private val plainAuthTypeOpt = authTypes.filterNot(_.equals(KERBEROS))
-    .filterNot(_.equals(NOSASL)).headOption
+
+  // take the first declared SASL/PLAIN auth type
+  private val effectivePlainAuthType = authTypes.find {
+    case NOSASL | KERBEROS => false
+    case _ => true
+  }
 
   private val hadoopAuthServer: Option[HadoopThriftAuthBridgeServer] = {
     if (kerberosEnabled) {
@@ -70,7 +74,7 @@ class KyuubiAuthenticationFactory(conf: KyuubiConf, isServer: Boolean = true) ex
   }
 
   def getTTransportFactory: TTransportFactory = {
-    if (noSaslEnabled) {
+    if (saslDisabled) {
       new TTransportFactory()
     } else {
       var transportFactory: TSaslServerTransport.Factory = null
@@ -87,7 +91,7 @@ class KyuubiAuthenticationFactory(conf: KyuubiConf, isServer: Boolean = true) ex
         case _ =>
       }
 
-      plainAuthTypeOpt match {
+      effectivePlainAuthType match {
         case Some(plainAuthType) =>
           transportFactory = PlainSASLHelper.getTransportFactory(
             plainAuthType.toString,
diff --git a/kyuubi-common/src/test/scala/org/apache/kyuubi/service/authentication/KyuubiAuthenticationFactorySuite.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/service/authentication/KyuubiAuthenticationFactorySuite.scala
index 607c397d81f..e9aad2b948f 100644
--- a/kyuubi-common/src/test/scala/org/apache/kyuubi/service/authentication/KyuubiAuthenticationFactorySuite.scala
+++ b/kyuubi-common/src/test/scala/org/apache/kyuubi/service/authentication/KyuubiAuthenticationFactorySuite.scala
@@ -56,21 +56,21 @@ class KyuubiAuthenticationFactorySuite extends KyuubiFunSuite {
   }
 
   test("AuthType Other") {
-    val conf = KyuubiConf().set(KyuubiConf.AUTHENTICATION_METHOD, Set("INVALID"))
+    val conf = KyuubiConf().set(KyuubiConf.AUTHENTICATION_METHOD, Seq("INVALID"))
     interceptEquals[IllegalArgumentException] { new KyuubiAuthenticationFactory(conf) }(
       "The value of kyuubi.authentication should be one of" +
         " NOSASL, NONE, LDAP, JDBC, KERBEROS, CUSTOM, but was INVALID")
   }
 
   test("AuthType LDAP") {
-    val conf = KyuubiConf().set(KyuubiConf.AUTHENTICATION_METHOD, Set("LDAP"))
+    val conf = KyuubiConf().set(KyuubiConf.AUTHENTICATION_METHOD, Seq("LDAP"))
     val authFactory = new KyuubiAuthenticationFactory(conf)
     authFactory.getTTransportFactory
     assert(Security.getProviders.exists(_.isInstanceOf[SaslPlainProvider]))
   }
 
   test("AuthType KERBEROS w/o keytab/principal") {
-    val conf = KyuubiConf().set(KyuubiConf.AUTHENTICATION_METHOD, Set("KERBEROS"))
+    val conf = KyuubiConf().set(KyuubiConf.AUTHENTICATION_METHOD, Seq("KERBEROS"))
 
     val factory = new KyuubiAuthenticationFactory(conf)
     val e = intercept[LoginException](factory.getTTransportFactory)
@@ -78,11 +78,11 @@ class KyuubiAuthenticationFactorySuite extends KyuubiFunSuite {
   }
 
   test("AuthType is NOSASL if only NOSASL is specified") {
-    val conf = KyuubiConf().set(KyuubiConf.AUTHENTICATION_METHOD, Set("NOSASL"))
+    val conf = KyuubiConf().set(KyuubiConf.AUTHENTICATION_METHOD, Seq("NOSASL"))
     var factory = new KyuubiAuthenticationFactory(conf)
     !factory.getTTransportFactory.isInstanceOf[TSaslServerTransport.Factory]
 
-    conf.set(KyuubiConf.AUTHENTICATION_METHOD, Set("NOSASL", "NONE"))
+    conf.set(KyuubiConf.AUTHENTICATION_METHOD, Seq("NOSASL", "NONE"))
     factory = new KyuubiAuthenticationFactory(conf)
     factory.getTTransportFactory.isInstanceOf[TSaslServerTransport.Factory]
   }
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/http/ThriftHttpServlet.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/http/ThriftHttpServlet.scala
index eb8fb2caa69..980f35d70b5 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/http/ThriftHttpServlet.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/http/ThriftHttpServlet.scala
@@ -136,7 +136,7 @@ class ThriftHttpServlet(
       } else SessionManager.setForwardedAddresses(List.empty[String])
 
       // Generate new cookie and add it to the response
-      if (requireNewCookie && !authFactory.noSaslEnabled) {
+      if (requireNewCookie && !authFactory.saslDisabled) {
         val cookieToken = HttpAuthUtils.createCookieToken(clientUserName)
         val hs2Cookie = createCookie(signer.signCookie(cookieToken))
         if (isHttpOnlyCookie) response.setHeader("SET-COOKIE", getHttpOnlyCookieHeader(hs2Cookie))
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/RestClientTestHelper.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/RestClientTestHelper.scala
index 1c78b9fa612..8344cdef01d 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/RestClientTestHelper.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/RestClientTestHelper.scala
@@ -48,7 +48,7 @@ trait RestClientTestHelper extends RestFrontendTestHelper with KerberizedTestHel
     UserGroupInformation.setConfiguration(config)
     assert(UserGroupInformation.isSecurityEnabled)
 
-    val conf = KyuubiConf().set(KyuubiConf.AUTHENTICATION_METHOD, Set("KERBEROS", "LDAP", "CUSTOM"))
+    val conf = KyuubiConf().set(KyuubiConf.AUTHENTICATION_METHOD, Seq("KERBEROS", "LDAP", "CUSTOM"))
       .set(KyuubiConf.SERVER_KEYTAB.key, testKeytab)
       .set(KyuubiConf.SERVER_PRINCIPAL, testPrincipal)
       .set(KyuubiConf.SERVER_SPNEGO_KEYTAB, testKeytab)
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationKerberosAndPlainAuthSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationKerberosAndPlainAuthSuite.scala
index 1791b492e25..31cde639734 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationKerberosAndPlainAuthSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationKerberosAndPlainAuthSuite.scala
@@ -64,7 +64,7 @@ class KyuubiOperationKerberosAndPlainAuthSuite extends WithKyuubiServer with Ker
     assert(UserGroupInformation.isSecurityEnabled)
 
     KyuubiConf()
-      .set(KyuubiConf.AUTHENTICATION_METHOD, Set("KERBEROS", "LDAP", "CUSTOM"))
+      .set(KyuubiConf.AUTHENTICATION_METHOD, Seq("KERBEROS", "LDAP", "CUSTOM"))
       .set(KyuubiConf.SERVER_KEYTAB, testKeytab)
       .set(KyuubiConf.SERVER_PRINCIPAL, testPrincipal)
       .set(KyuubiConf.AUTHENTICATION_LDAP_URL, ldapUrl)
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/thrift/http/KyuubiOperationThriftHttpKerberosAndPlainAuthSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/thrift/http/KyuubiOperationThriftHttpKerberosAndPlainAuthSuite.scala
index cee43bf5cf1..941e121a6cd 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/thrift/http/KyuubiOperationThriftHttpKerberosAndPlainAuthSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/thrift/http/KyuubiOperationThriftHttpKerberosAndPlainAuthSuite.scala
@@ -49,7 +49,7 @@ class KyuubiOperationThriftHttpKerberosAndPlainAuthSuite
     UserGroupInformation.setConfiguration(config)
     assert(UserGroupInformation.isSecurityEnabled)
 
-    KyuubiConf().set(KyuubiConf.AUTHENTICATION_METHOD, Set("KERBEROS", "LDAP", "CUSTOM"))
+    KyuubiConf().set(KyuubiConf.AUTHENTICATION_METHOD, Seq("KERBEROS", "LDAP", "CUSTOM"))
       .set(KyuubiConf.SERVER_KEYTAB, testKeytab)
       .set(KyuubiConf.SERVER_PRINCIPAL, testPrincipal)
       .set(KyuubiConf.AUTHENTICATION_LDAP_URL, ldapUrl)
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/AdminResourceSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/AdminResourceSuite.scala
index 0951d82727c..c69a97cefd6 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/AdminResourceSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/AdminResourceSuite.scala
@@ -51,7 +51,7 @@ class AdminResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper {
   private val engineMgr = new KyuubiApplicationManager()
 
   override protected lazy val conf: KyuubiConf = KyuubiConf()
-    .set(AUTHENTICATION_METHOD, Set("CUSTOM"))
+    .set(AUTHENTICATION_METHOD, Seq("CUSTOM"))
     .set(AUTHENTICATION_CUSTOM_CLASS, classOf[AnonymousAuthenticationProviderImpl].getName)
     .set(SERVER_ADMINISTRATORS, Set("admin001"))
     .set(ENGINE_IDLE_TIMEOUT, Duration.ofMinutes(3).toMillis)
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala
index 6ae2bd04063..6d33093997c 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala
@@ -85,7 +85,7 @@ abstract class BatchesResourceSuiteBase extends KyuubiFunSuite
   override protected lazy val conf: KyuubiConf = {
     val testResourceDir = Paths.get(sparkBatchTestResource.get).getParent
     val kyuubiConf = KyuubiConf()
-      .set(AUTHENTICATION_METHOD, Set("CUSTOM"))
+      .set(AUTHENTICATION_METHOD, Seq("CUSTOM"))
       .set(AUTHENTICATION_CUSTOM_CLASS, classOf[AnonymousAuthenticationProviderImpl].getName)
       .set(SERVER_ADMINISTRATORS, Set("admin"))
       .set(BATCH_IMPL_VERSION, batchVersion)
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/AdminCtlSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/AdminCtlSuite.scala
index 32bb6fbb152..986b171c142 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/AdminCtlSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/AdminCtlSuite.scala
@@ -56,7 +56,7 @@ class AdminCtlSuite extends RestClientTestHelper with TestPrematureExit {
     val id = UUID.randomUUID().toString
     conf.set(HighAvailabilityConf.HA_NAMESPACE, "kyuubi_test")
     conf.set(KyuubiConf.ENGINE_IDLE_TIMEOUT, 180000L)
-    conf.set(KyuubiConf.AUTHENTICATION_METHOD, Set("LDAP", "CUSTOM"))
+    conf.set(KyuubiConf.AUTHENTICATION_METHOD, Seq("LDAP", "CUSTOM"))
     conf.set(KyuubiConf.GROUP_PROVIDER, "hadoop")
 
     val user = ldapUser
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/AdminRestApiSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/AdminRestApiSuite.scala
index c8f1d68e67e..4a05dc079c0 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/AdminRestApiSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/AdminRestApiSuite.scala
@@ -50,7 +50,7 @@ class AdminRestApiSuite extends RestClientTestHelper {
     val id = UUID.randomUUID().toString
     conf.set(HighAvailabilityConf.HA_NAMESPACE, "kyuubi_test")
     conf.set(KyuubiConf.ENGINE_IDLE_TIMEOUT, 180000L)
-    conf.set(KyuubiConf.AUTHENTICATION_METHOD, Set("LDAP", "CUSTOM"))
+    conf.set(KyuubiConf.AUTHENTICATION_METHOD, Seq("LDAP", "CUSTOM"))
     conf.set(KyuubiConf.GROUP_PROVIDER, "hadoop")
     val user = ldapUser
     val engine = new EngineRef(conf.clone, user, PluginLoader.loadGroupProvider(conf), id, null)