support rest api user defined authentication

docs/configuration/settings.md

@@ -247,6 +247,7 @@ You can configure the Kyuubi properties in `$KYUUBI_HOME/conf/kyuubi-defaults.co
 | kyuubi.frontend.protocols                              | THRIFT_BINARY,REST | A comma-separated list for all frontend protocols <ul> <li>THRIFT_BINARY - HiveServer2 compatible thrift binary protocol.</li> <li>THRIFT_HTTP - HiveServer2 compatible thrift http protocol.</li> <li>REST - Kyuubi defined REST API(experimental).</li>  <li>MYSQL - MySQL compatible text protocol(experimental).</li>  <li>TRINO - Trino compatible http protocol(experimental).</li> </ul>                                                                                                                             | seq      | 1.4.0 |
 | kyuubi.frontend.proxy.http.client.ip.header            | X-Real-IP          | The HTTP header to record the real client IP address. If your server is behind a load balancer or other proxy, the server will see this load balancer or proxy IP address as the client IP address, to get around this common issue, most load balancers or proxies offer the ability to record the real remote IP address in an HTTP header that will be added to the request for other devices to use. Note that, because the header value can be specified to any IP address, so it will not be used for authentication. | string   | 1.6.0 |
 | kyuubi.frontend.rest.authentication                    | NONE               | A comma-separated list of rest protocol client authentication types. It fallback to `kyuubi.authentication` if not configure.                                                                                                                                                                                                                                                                                                                                                                                               | seq      | 1.9.0 |
+| kyuubi.frontend.rest.authentication.custom.class       | &lt;undefined&gt;  | User-defined authentication implementation of org.apache.kyuubi.service.authentication.PasswdAuthenticationProvider for rest protocol.                                                                                                                                                                                                                                                                                                                                                                                      | string   | 1.9.0 |
 | kyuubi.frontend.rest.bind.host                         | &lt;undefined&gt;  | Hostname or IP of the machine on which to run the REST frontend service.                                                                                                                                                                                                                                                                                                                                                                                                                                                    | string   | 1.4.0 |
 | kyuubi.frontend.rest.bind.port                         | 10099              | Port of the machine on which to run the REST frontend service.                                                                                                                                                                                                                                                                                                                                                                                                                                                              | int      | 1.4.0 |
 | kyuubi.frontend.rest.max.worker.threads                | 999                | Maximum number of threads in the frontend worker thread pool for the rest frontend service                                                                                                                                                                                                                                                                                                                                                                                                                                  | int      | 1.6.2 |

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

@@ -814,6 +814,14 @@ object KyuubiConf {
       .stringConf
       .createOptional
 
+  val FRONTEND_REST_AUTHENTICATION_CUSTOM_CLASS: ConfigEntry[Option[String]] =
+    buildConf("kyuubi.frontend.rest.authentication.custom.class")
+      .doc("User-defined authentication implementation of " +
+        "org.apache.kyuubi.service.authentication.PasswdAuthenticationProvider for rest protocol.")
+      .version("1.9.0")
+      .serverOnly
+      .fallbackConf(AUTHENTICATION_CUSTOM_CLASS)
+
   val AUTHENTICATION_LDAP_URL: OptionalConfigEntry[String] =
     buildConf("kyuubi.authentication.ldap.url")
       .doc("SPACE character separated LDAP connection URL(s).")

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

@@ -28,6 +28,7 @@ import org.apache.hadoop.conf.Configuration
 import org.apache.kyuubi.{KyuubiSQLException, Logging, Utils}
 import org.apache.kyuubi.Utils.stringifyException
 import org.apache.kyuubi.config.KyuubiConf.{FRONTEND_ADVERTISED_HOST, FRONTEND_CONNECTION_URL_USE_HOSTNAME, PROXY_USER, SESSION_CLOSE_ON_DISCONNECT}
+import org.apache.kyuubi.config.KyuubiConf.FrontendProtocols.THRIFT_BINARY
 import org.apache.kyuubi.config.KyuubiReservedKeys._
 import org.apache.kyuubi.operation.{FetchOrientation, OperationHandle}
 import org.apache.kyuubi.service.authentication.{AuthUtils, KyuubiAuthenticationFactory}
@@ -57,7 +58,7 @@ abstract class TFrontendService(name: String)
   protected lazy val serverSocket = new ServerSocket(portNum, -1, serverAddr)
   protected lazy val actualPort: Int = serverSocket.getLocalPort
   protected lazy val authFactory: KyuubiAuthenticationFactory =
-    new KyuubiAuthenticationFactory(conf, isServer())
+    new KyuubiAuthenticationFactory(conf, THRIFT_BINARY, isServer())
 
   protected def hadoopConf: Configuration = _hadoopConf
 

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

@@ -20,6 +20,7 @@ package org.apache.kyuubi.service.authentication
 import javax.security.sasl.AuthenticationException
 
 import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf.FrontendProtocols.{FrontendProtocol, REST}
 import org.apache.kyuubi.service.authentication.AuthMethods.AuthMethod
 import org.apache.kyuubi.util.ClassUtils
 
@@ -31,22 +32,27 @@ object AuthenticationProviderFactory {
   def getAuthenticationProvider(
       method: AuthMethod,
       conf: KyuubiConf,
+      protocol: FrontendProtocol,
       isServer: Boolean = true): PasswdAuthenticationProvider = {
     if (isServer) {
-      getAuthenticationProviderForServer(method, conf)
+      getAuthenticationProviderForServer(method, conf, protocol)
     } else {
       getAuthenticationProviderForEngine(conf)
     }
   }
 
   private def getAuthenticationProviderForServer(
       method: AuthMethod,
-      conf: KyuubiConf): PasswdAuthenticationProvider = method match {
+      conf: KyuubiConf,
+      protocol: FrontendProtocol): PasswdAuthenticationProvider = method match {
     case AuthMethods.NONE => new AnonymousAuthenticationProviderImpl
     case AuthMethods.LDAP => new LdapAuthenticationProviderImpl(conf)
     case AuthMethods.JDBC => new JdbcAuthenticationProviderImpl(conf)
     case AuthMethods.CUSTOM =>
-      val className = conf.get(KyuubiConf.AUTHENTICATION_CUSTOM_CLASS)
+      val className = protocol match {
+        case REST => conf.get(KyuubiConf.FRONTEND_REST_AUTHENTICATION_CUSTOM_CLASS)
+        case _ => conf.get(KyuubiConf.AUTHENTICATION_CUSTOM_CLASS)
+      }
       if (className.isEmpty) {
         throw new AuthenticationException(
           "authentication.custom.class must be set when auth method was CUSTOM.")

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

@@ -24,12 +24,16 @@ import javax.security.sasl.Sasl
 import org.apache.kyuubi.Logging
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf._
+import org.apache.kyuubi.config.KyuubiConf.FrontendProtocols.FrontendProtocol
 import org.apache.kyuubi.service.authentication.AuthTypes._
 import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TCLIService.Iface
 import org.apache.kyuubi.shaded.thrift.TProcessorFactory
 import org.apache.kyuubi.shaded.thrift.transport.{TSaslServerTransport, TTransportException, TTransportFactory}
 
-class KyuubiAuthenticationFactory(conf: KyuubiConf, isServer: Boolean = true) extends Logging {
+class KyuubiAuthenticationFactory(
+    conf: KyuubiConf,
+    protocol: FrontendProtocol,
+    isServer: Boolean = true) extends Logging {
 
   val authTypes: Seq[AuthType] = conf.get(AUTHENTICATION_METHOD).map(AuthTypes.withName)
   val saslDisabled: Boolean = AuthUtils.saslDisabled(authTypes)
@@ -85,6 +89,7 @@ class KyuubiAuthenticationFactory(conf: KyuubiConf, isServer: Boolean = true) ex
           transportFactory = PlainSASLHelper.getTransportFactory(
             plainAuthType.toString,
             conf,
+            protocol,
             Option(transportFactory),
             isServer).asInstanceOf[TSaslServerTransport.Factory]
 

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

@@ -23,6 +23,7 @@ import javax.security.auth.callback.{Callback, CallbackHandler, NameCallback, Pa
 import javax.security.sasl.AuthorizeCallback
 
 import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf.FrontendProtocols.FrontendProtocol
 import org.apache.kyuubi.service.authentication.AuthMethods.AuthMethod
 import org.apache.kyuubi.service.authentication.PlainSASLServer.SaslPlainProvider
 import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TCLIService.Iface
@@ -42,11 +43,16 @@ object PlainSASLHelper {
   private class PlainServerCallbackHandler private (
       authMethod: AuthMethod,
       conf: KyuubiConf,
+      protocol: FrontendProtocol,
       isServer: Boolean)
     extends CallbackHandler {
 
-    def this(authMethodStr: String, conf: KyuubiConf, isServer: Boolean) =
-      this(AuthMethods.withName(authMethodStr), conf, isServer)
+    def this(
+        authMethodStr: String,
+        conf: KyuubiConf,
+        protocol: FrontendProtocol,
+        isServer: Boolean) =
+      this(AuthMethods.withName(authMethodStr), conf, protocol, isServer)
 
     @throws[UnsupportedCallbackException]
     override def handle(callbacks: Array[Callback]): Unit = {
@@ -64,7 +70,11 @@ object PlainSASLHelper {
         }
       }
       val provider =
-        AuthenticationProviderFactory.getAuthenticationProvider(authMethod, conf, isServer)
+        AuthenticationProviderFactory.getAuthenticationProvider(
+          authMethod,
+          conf,
+          protocol,
+          isServer)
       provider.authenticate(username, password)
       if (ac != null) ac.setAuthorized(true)
     }
@@ -77,11 +87,12 @@ object PlainSASLHelper {
   def getTransportFactory(
       authTypeStr: String,
       conf: KyuubiConf,
+      protocol: FrontendProtocol,
       transportFactory: Option[TSaslServerTransport.Factory] = None,
       isServer: Boolean = true): TTransportFactory = {
     val saslFactory = transportFactory.getOrElse(new TSaslServerTransport.Factory())
     try {
-      val handler = new PlainServerCallbackHandler(authTypeStr, conf, isServer)
+      val handler = new PlainServerCallbackHandler(authTypeStr, conf, protocol, isServer)
       val props = new java.util.HashMap[String, String]
       saslFactory.addServerDefinition("PLAIN", authTypeStr, null, props, handler)
     } catch {

kyuubi-common/src/test/scala/org/apache/kyuubi/service/authentication/AuthenticationProviderFactorySuite.scala

@@ -21,20 +21,21 @@ import javax.security.sasl.AuthenticationException
 
 import org.apache.kyuubi.{KyuubiFunSuite, Utils}
 import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf.FrontendProtocols.THRIFT_BINARY
 
 class AuthenticationProviderFactorySuite extends KyuubiFunSuite {
 
   import AuthenticationProviderFactory._
 
   test("get auth provider") {
     val conf = KyuubiConf()
-    val p1 = getAuthenticationProvider(AuthMethods.withName("NONE"), conf)
+    val p1 = getAuthenticationProvider(AuthMethods.withName("NONE"), conf, THRIFT_BINARY)
     p1.authenticate(Utils.currentUser, "")
-    val p2 = getAuthenticationProvider(AuthMethods.withName("LDAP"), conf)
+    val p2 = getAuthenticationProvider(AuthMethods.withName("LDAP"), conf, THRIFT_BINARY)
     val e1 = intercept[AuthenticationException](p2.authenticate("test", "test"))
     assert(e1.getMessage.contains("Error validating LDAP user:"))
     val e2 = intercept[AuthenticationException](
-      AuthenticationProviderFactory.getAuthenticationProvider(null, conf))
+      AuthenticationProviderFactory.getAuthenticationProvider(null, conf, THRIFT_BINARY))
     assert(e2.getMessage === "Not a valid authentication method")
   }
 

kyuubi-common/src/test/scala/org/apache/kyuubi/service/authentication/CustomAuthenticationProviderImplSuite.scala

@@ -21,21 +21,43 @@ import javax.security.sasl.AuthenticationException
 
 import org.apache.kyuubi.KyuubiFunSuite
 import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf.FrontendProtocols.{REST, THRIFT_BINARY}
 import org.apache.kyuubi.service.authentication.AuthenticationProviderFactory.getAuthenticationProvider
 
 class CustomAuthenticationProviderImplSuite extends KyuubiFunSuite {
   test("Test user defined authentication") {
     val conf = KyuubiConf()
 
     val e1 = intercept[AuthenticationException](
-      getAuthenticationProvider(AuthMethods.withName("CUSTOM"), conf))
+      getAuthenticationProvider(AuthMethods.withName("CUSTOM"), conf, THRIFT_BINARY))
     assert(e1.getMessage.contains(
       "authentication.custom.class must be set when auth method was CUSTOM."))
 
     conf.set(
       KyuubiConf.AUTHENTICATION_CUSTOM_CLASS,
       classOf[UserDefineAuthenticationProviderImpl].getCanonicalName)
-    val p1 = getAuthenticationProvider(AuthMethods.withName("CUSTOM"), conf)
+    val p1 = getAuthenticationProvider(AuthMethods.withName("CUSTOM"), conf, THRIFT_BINARY)
+    val e2 = intercept[AuthenticationException](p1.authenticate("test", "test"))
+    assert(e2.getMessage.contains("Username or password is not valid!"))
+
+    p1.authenticate("user", "password")
+  }
+}
+
+class RestCustomAuthenticationProviderImplSuite extends KyuubiFunSuite {
+  test("Test user defined authentication") {
+    val conf = KyuubiConf()
+
+    conf.set(
+      KyuubiConf.FRONTEND_REST_AUTHENTICATION_CUSTOM_CLASS,
+      Some(classOf[UserDefineAuthenticationProviderImpl].getCanonicalName))
+
+    val e1 = intercept[AuthenticationException](
+      getAuthenticationProvider(AuthMethods.withName("CUSTOM"), conf, THRIFT_BINARY))
+    assert(e1.getMessage.contains(
+      "authentication.custom.class must be set when auth method was CUSTOM."))
+
+    val p1 = getAuthenticationProvider(AuthMethods.withName("CUSTOM"), conf, REST)
     val e2 = intercept[AuthenticationException](p1.authenticate("test", "test"))
     assert(e2.getMessage.contains("Username or password is not valid!"))
 

kyuubi-common/src/test/scala/org/apache/kyuubi/service/authentication/KyuubiAuthenticationFactorySuite.scala

@@ -22,6 +22,7 @@ import javax.security.auth.login.LoginException
 
 import org.apache.kyuubi.{KyuubiFunSuite, KyuubiSQLException}
 import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf.FrontendProtocols.THRIFT_BINARY
 import org.apache.kyuubi.service.authentication.PlainSASLServer.SaslPlainProvider
 import org.apache.kyuubi.shaded.thrift.transport.TSaslServerTransport
 import org.apache.kyuubi.util.AssertionUtils._
@@ -46,7 +47,7 @@ class KyuubiAuthenticationFactorySuite extends KyuubiFunSuite {
 
   test("AuthType NONE") {
     val kyuubiConf = KyuubiConf()
-    val auth = new KyuubiAuthenticationFactory(kyuubiConf)
+    val auth = new KyuubiAuthenticationFactory(kyuubiConf, THRIFT_BINARY)
     auth.getTTransportFactory
     assert(Security.getProviders.exists(_.isInstanceOf[SaslPlainProvider]))
 
@@ -56,33 +57,34 @@ class KyuubiAuthenticationFactorySuite extends KyuubiFunSuite {
 
   test("AuthType Other") {
     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")
+    interceptEquals[IllegalArgumentException] {
+      new KyuubiAuthenticationFactory(conf, THRIFT_BINARY)
+    }("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, Seq("LDAP"))
-    val authFactory = new KyuubiAuthenticationFactory(conf)
+    val authFactory = new KyuubiAuthenticationFactory(conf, THRIFT_BINARY)
     authFactory.getTTransportFactory
     assert(Security.getProviders.exists(_.isInstanceOf[SaslPlainProvider]))
   }
 
   test("AuthType KERBEROS w/o keytab/principal") {
     val conf = KyuubiConf().set(KyuubiConf.AUTHENTICATION_METHOD, Seq("KERBEROS"))
 
-    val factory = new KyuubiAuthenticationFactory(conf)
+    val factory = new KyuubiAuthenticationFactory(conf, THRIFT_BINARY)
     val e = intercept[LoginException](factory.getTTransportFactory)
     assert(e.getMessage startsWith "Kerberos principal should have 3 parts")
   }
 
   test("AuthType is NOSASL if only NOSASL is specified") {
     val conf = KyuubiConf().set(KyuubiConf.AUTHENTICATION_METHOD, Seq("NOSASL"))
-    var factory = new KyuubiAuthenticationFactory(conf)
+    var factory = new KyuubiAuthenticationFactory(conf, THRIFT_BINARY)
     !factory.getTTransportFactory.isInstanceOf[TSaslServerTransport.Factory]
 
     conf.set(KyuubiConf.AUTHENTICATION_METHOD, Seq("NOSASL", "NONE"))
-    factory = new KyuubiAuthenticationFactory(conf)
+    factory = new KyuubiAuthenticationFactory(conf, THRIFT_BINARY)
     factory.getTTransportFactory.isInstanceOf[TSaslServerTransport.Factory]
   }
 }

kyuubi-common/src/test/scala/org/apache/kyuubi/service/authentication/PlainSASLHelperSuite.scala

@@ -21,6 +21,7 @@ import java.security.Security
 
 import org.apache.kyuubi.{KYUUBI_VERSION, KyuubiFunSuite}
 import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf.FrontendProtocols.THRIFT_BINARY
 import org.apache.kyuubi.service.{NoopTBinaryFrontendServer, TBinaryFrontendService}
 import org.apache.kyuubi.service.authentication.PlainSASLServer.SaslPlainProvider
 import org.apache.kyuubi.shaded.thrift.transport.{TSaslServerTransport, TSocket}
@@ -39,20 +40,20 @@ class PlainSASLHelperSuite extends KyuubiFunSuite {
     val tProcessor = tProcessorFactory.getProcessor(tSocket)
     assert(tProcessor.isInstanceOf[TSetIpAddressProcessor[_]])
     val e = intercept[IllegalArgumentException] {
-      PlainSASLHelper.getTransportFactory("KERBEROS", conf)
+      PlainSASLHelper.getTransportFactory("KERBEROS", conf, THRIFT_BINARY)
     }
     assert(e.getMessage === "Illegal authentication type KERBEROS for plain transport")
     val e2 = intercept[IllegalArgumentException] {
-      PlainSASLHelper.getTransportFactory("NOSASL", conf)
+      PlainSASLHelper.getTransportFactory("NOSASL", conf, THRIFT_BINARY)
     }
     assert(e2.getMessage === "Illegal authentication type NOSASL for plain transport")
 
     val e3 = intercept[IllegalArgumentException] {
-      PlainSASLHelper.getTransportFactory("ELSE", conf)
+      PlainSASLHelper.getTransportFactory("ELSE", conf, THRIFT_BINARY)
     }
     assert(e3.getMessage === "Illegal authentication type ELSE for plain transport")
 
-    val tTransportFactory = PlainSASLHelper.getTransportFactory("NONE", conf)
+    val tTransportFactory = PlainSASLHelper.getTransportFactory("NONE", conf, THRIFT_BINARY)
     assert(tTransportFactory.isInstanceOf[TSaslServerTransport.Factory])
     Security.getProviders.exists(_.isInstanceOf[SaslPlainProvider])
   }