diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/InternalSecurityAccessor.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/InternalSecurityAccessor.scala index 62680e6a6..afc1dde1f 100644 --- a/kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/InternalSecurityAccessor.scala +++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/service/authentication/InternalSecurityAccessor.scala @@ -20,6 +20,8 @@ package org.apache.kyuubi.service.authentication import javax.crypto.Cipher import javax.crypto.spec.{IvParameterSpec, SecretKeySpec} +import org.apache.hadoop.classification.VisibleForTesting + import org.apache.kyuubi.{KyuubiSQLException, Logging} import org.apache.kyuubi.config.KyuubiConf import org.apache.kyuubi.config.KyuubiConf._ @@ -121,4 +123,9 @@ object InternalSecurityAccessor extends Logging { def get(): InternalSecurityAccessor = { _engineSecurityAccessor } + + @VisibleForTesting + def reset(): Unit = { + _engineSecurityAccessor = null + } } diff --git a/kyuubi-common/src/test/scala/org/apache/kyuubi/KyuubiFunSuite.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/KyuubiFunSuite.scala index 2789d7f89..8d0a14c16 100644 --- a/kyuubi-common/src/test/scala/org/apache/kyuubi/KyuubiFunSuite.scala +++ b/kyuubi-common/src/test/scala/org/apache/kyuubi/KyuubiFunSuite.scala @@ -30,6 +30,7 @@ import org.scalatest.funsuite.AnyFunSuite import org.slf4j.bridge.SLF4JBridgeHandler import org.apache.kyuubi.config.internal.Tests.IS_TESTING +import org.apache.kyuubi.service.authentication.InternalSecurityAccessor trait KyuubiFunSuite extends AnyFunSuite with BeforeAndAfterAll @@ -46,6 +47,7 @@ trait KyuubiFunSuite extends AnyFunSuite override def beforeAll(): Unit = { System.setProperty(IS_TESTING.key, "true") doThreadPreAudit() + InternalSecurityAccessor.reset() super.beforeAll() } diff --git a/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/JDBCTestHelper.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/JDBCTestHelper.scala index 97330837d..e7802f2fe 100644 --- a/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/JDBCTestHelper.scala +++ b/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/JDBCTestHelper.scala @@ -53,6 +53,7 @@ trait JDBCTestHelper extends KyuubiFunSuite { def withMultipleConnectionJdbcStatement( tableNames: String*)(fs: (Statement => Unit)*): Unit = { + info(s"Create JDBC connection using: $jdbcUrlWithConf") val connections = fs.map { _ => DriverManager.getConnection(jdbcUrlWithConf, user, password) } val statements = connections.map(_.createStatement()) diff --git a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClient.scala b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClient.scala index 33aa5149d..0127f2df1 100644 --- a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClient.scala +++ b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClient.scala @@ -305,6 +305,10 @@ class ZookeeperDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient { basePath, initData.getBytes(StandardCharsets.UTF_8)) secretNode.start() + val znodeTimeout = conf.get(HA_ZK_NODE_TIMEOUT) + if (!secretNode.waitForInitialCreate(znodeTimeout, TimeUnit.MILLISECONDS)) { + throw new KyuubiException(s"Max znode creation wait time $znodeTimeout s exhausted") + } } override def getAndIncrement(path: String, delta: Int = 1): Int = { diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationWithEngineSecurity.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationWithEngineSecuritySuite.scala similarity index 80% rename from kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationWithEngineSecurity.scala rename to kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationWithEngineSecuritySuite.scala index 63369f4b2..da6367bc4 100644 --- a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationWithEngineSecurity.scala +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationWithEngineSecuritySuite.scala @@ -17,25 +17,26 @@ package org.apache.kyuubi.operation +import java.nio.charset.StandardCharsets + import org.apache.kyuubi.WithKyuubiServer import org.apache.kyuubi.config.KyuubiConf import org.apache.kyuubi.ha.HighAvailabilityConf import org.apache.kyuubi.ha.client.DiscoveryClientProvider -import org.apache.kyuubi.service.authentication.{InternalSecurityAccessor, ZooKeeperEngineSecuritySecretProviderImpl} +import org.apache.kyuubi.service.authentication.InternalSecurityAccessor -class KyuubiOperationWithEngineSecurity extends WithKyuubiServer with HiveJDBCTestHelper { +class KyuubiOperationWithEngineSecuritySuite extends WithKyuubiServer with HiveJDBCTestHelper { import DiscoveryClientProvider._ override protected def jdbcUrl: String = getJdbcUrl private val engineSecretNode = "/SECRET" + private val engineSecret = "_ENGINE_SECRET_" override protected val conf: KyuubiConf = { KyuubiConf() .set(KyuubiConf.ENGINE_SECURITY_ENABLED, false) - .set( - KyuubiConf.ENGINE_SECURITY_SECRET_PROVIDER, - classOf[ZooKeeperEngineSecuritySecretProviderImpl].getCanonicalName) + .set(KyuubiConf.ENGINE_SECURITY_SECRET_PROVIDER, "zookeeper") .set(HighAvailabilityConf.HA_ZK_ENGINE_SECURE_SECRET_NODE, engineSecretNode) } @@ -43,7 +44,9 @@ class KyuubiOperationWithEngineSecurity extends WithKyuubiServer with HiveJDBCTe super.beforeAll() withDiscoveryClient(conf) { discoveryClient => discoveryClient.create(engineSecretNode, "PERSISTENT", false) - discoveryClient.startSecretNode("PERSISTENT", engineSecretNode, "_ENGINE_SECRET_") + discoveryClient.startSecretNode("PERSISTENT", engineSecretNode, engineSecret) + val expected = engineSecret.getBytes(StandardCharsets.UTF_8) + assert(discoveryClient.getData(engineSecretNode) === expected) } conf.set(KyuubiConf.ENGINE_SECURITY_ENABLED, true) diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiRestAuthenticationSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiRestAuthenticationSuite.scala index 61de82251..089b756f5 100644 --- a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiRestAuthenticationSuite.scala +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiRestAuthenticationSuite.scala @@ -37,12 +37,20 @@ import org.apache.kyuubi.session.KyuubiSession class KyuubiRestAuthenticationSuite extends RestClientTestHelper { override protected val otherConfigs: Map[String, String] = { - // allow to impersonate other users with spnego authentication Map( + KyuubiConf.ENGINE_SECURITY_ENABLED.key -> "true", + KyuubiConf.ENGINE_SECURITY_SECRET_PROVIDER.key -> "simple", + KyuubiConf.SIMPLE_SECURITY_SECRET_PROVIDER_PROVIDER_SECRET.key -> "_KYUUBI_REST_", + // allow to impersonate other users with spnego authentication s"hadoop.proxyuser.$clientPrincipalUser.groups" -> "*", s"hadoop.proxyuser.$clientPrincipalUser.hosts" -> "*") } + override def beforeAll(): Unit = { + super.beforeAll() + InternalSecurityAccessor.initialize(conf, true) + } + test("test with LDAP authorization") { val encodeAuthorization = new String( Base64.getEncoder.encode( 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 b6bc1af52..dd62ee48d 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 @@ -45,7 +45,7 @@ import org.apache.kyuubi.operation.OperationState.OperationState import org.apache.kyuubi.server.KyuubiRestFrontendService import org.apache.kyuubi.server.http.authentication.AuthenticationHandler.AUTHORIZATION_HEADER import org.apache.kyuubi.server.metadata.api.{Metadata, MetadataFilter} -import org.apache.kyuubi.service.authentication.KyuubiAuthenticationFactory +import org.apache.kyuubi.service.authentication.{InternalSecurityAccessor, KyuubiAuthenticationFactory} import org.apache.kyuubi.session.{KyuubiBatchSession, KyuubiSessionManager, SessionHandle, SessionType} class BatchesResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper with BatchTestHelper { @@ -57,6 +57,11 @@ class BatchesResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper wi KyuubiConf.SESSION_LOCAL_DIR_ALLOW_LIST, Seq(Paths.get(sparkBatchTestResource.get).getParent.toString)) + override def beforeAll(): Unit = { + super.beforeAll() + InternalSecurityAccessor.initialize(conf, true) + } + override def afterEach(): Unit = { val sessionManager = fe.be.sessionManager.asInstanceOf[KyuubiSessionManager] sessionManager.allSessions().foreach { session =>