This is an automated email from the ASF dual-hosted git repository.
rexxiong pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new b7d6704cc [CELEBORN-1251] Connect the server and client bootstraps to
RpcEnv
b7d6704cc is described below
commit b7d6704cc8e9ab1034dc5aab2e2539c8a80d993e
Author: Chandni Singh <[email protected]>
AuthorDate: Thu Jan 25 17:28:48 2024 +0800
[CELEBORN-1251] Connect the server and client bootstraps to RpcEnv
### What changes were proposed in this pull request?
This connects client/server bootstraps to the RpcEnv in Celeborn. This is a
prerequisite for leveraging RPC security in subsequent PRs where we will add
Sasl authentication to the communication between the client and Celeborn
Master/Workers.
It is part of the epic: https://issues.apache.org/jira/browse/CELEBORN-1011.
### Why are the changes needed?
The changes are needed for adding authentication to Celeborn. See
[CELEBORN-1011](https://issues.apache.org/jira/browse/CELEBORN-1011).
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
Added some UTs
Closes #2257 from otterc/CELEBORN-1251.
Authored-by: Chandni Singh <[email protected]>
Signed-off-by: Shuang <[email protected]>
---
.../apache/celeborn/client/ShuffleClientImpl.java | 8 +-
.../apache/celeborn/client/LifecycleManager.scala | 3 +-
.../common/network/sasl/SecretRegistryImpl.java | 11 +-
.../org/apache/celeborn/common/rpc/RpcEnv.scala | 15 +-
.../celeborn/common/rpc/RpcSecurityContext.scala | 161 +++++++++++++++++++++
.../celeborn/common/rpc/netty/NettyRpcEnv.scala | 50 ++++++-
.../common/network/sasl/CelebornSaslSuiteJ.java | 7 +-
.../common/network/sasl/RegistrationSuiteJ.java | 2 +-
.../celeborn/common/network/sasl/SaslTestBase.java | 6 +-
.../common/rpc/RpcSecurityContextSuite.scala | 136 +++++++++++++++++
.../common/rpc/netty/NettyRpcEnvSuite.scala | 7 +-
11 files changed, 373 insertions(+), 33 deletions(-)
diff --git
a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index b995bc13a..b7dc720b4 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -180,7 +180,13 @@ public class ShuffleClientImpl extends ShuffleClient {
}
// init rpc env
- rpcEnv = RpcEnv.create(RpcNameConstants.SHUFFLE_CLIENT_SYS,
Utils.localHostName(conf), 0, conf);
+ rpcEnv =
+ RpcEnv.create(
+ RpcNameConstants.SHUFFLE_CLIENT_SYS,
+ Utils.localHostName(conf),
+ 0,
+ conf,
+ scala.None$.empty());
String module = TransportModuleConstants.DATA_MODULE;
TransportConf dataTransportConf =
diff --git
a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
index a716b62a6..ad2a036d0 100644
--- a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
@@ -153,7 +153,8 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
RpcNameConstants.LIFECYCLE_MANAGER_SYS,
lifecycleHost,
conf.shuffleManagerPort,
- conf)
+ conf,
+ None)
rpcEnv.setupEndpoint(RpcNameConstants.LIFECYCLE_MANAGER_EP, this)
logInfo(s"Starting LifecycleManager on ${rpcEnv.address}")
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/sasl/SecretRegistryImpl.java
b/common/src/main/java/org/apache/celeborn/common/network/sasl/SecretRegistryImpl.java
index 811ccbf98..18342b7b5 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/sasl/SecretRegistryImpl.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/sasl/SecretRegistryImpl.java
@@ -19,18 +19,9 @@ package org.apache.celeborn.common.network.sasl;
import java.util.concurrent.ConcurrentHashMap;
-/**
- * A simple implementation of {@link SecretRegistry} that stores secrets in
memory. It is designed
- * as a singleton.
- */
+/** A simple implementation of {@link SecretRegistry} that stores secrets in
memory. */
public class SecretRegistryImpl implements SecretRegistry {
- private static final SecretRegistryImpl INSTANCE = new SecretRegistryImpl();
-
- public static SecretRegistryImpl getInstance() {
- return INSTANCE;
- }
-
private final ConcurrentHashMap<String, String> secrets = new
ConcurrentHashMap<>();
@Override
diff --git a/common/src/main/scala/org/apache/celeborn/common/rpc/RpcEnv.scala
b/common/src/main/scala/org/apache/celeborn/common/rpc/RpcEnv.scala
index f0d0f0452..7a911bd58 100644
--- a/common/src/main/scala/org/apache/celeborn/common/rpc/RpcEnv.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/rpc/RpcEnv.scala
@@ -23,6 +23,7 @@ import scala.concurrent.Future
import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.rpc.netty.NettyRpcEnvFactory
+import org.apache.celeborn.common.security.RpcSecurityContext
/**
* A RpcEnv implementation must have a [[RpcEnvFactory]] implementation with
an empty constructor
@@ -34,8 +35,9 @@ object RpcEnv {
name: String,
host: String,
port: Int,
- conf: CelebornConf): RpcEnv = {
- create(name, host, host, port, conf, 0)
+ conf: CelebornConf,
+ securityContext: Option[RpcSecurityContext]): RpcEnv = {
+ create(name, host, host, port, conf, 0, securityContext)
}
def create(
@@ -44,8 +46,10 @@ object RpcEnv {
advertiseAddress: String,
port: Int,
conf: CelebornConf,
- numUsableCores: Int): RpcEnv = {
- val config = RpcEnvConfig(conf, name, bindAddress, advertiseAddress, port,
numUsableCores)
+ numUsableCores: Int,
+ securityContext: Option[RpcSecurityContext] = None): RpcEnv = {
+ val config =
+ RpcEnvConfig(conf, name, bindAddress, advertiseAddress, port,
numUsableCores, securityContext)
new NettyRpcEnvFactory().create(config)
}
}
@@ -176,4 +180,5 @@ private[celeborn] case class RpcEnvConfig(
bindAddress: String,
advertiseAddress: String,
port: Int,
- numUsableCores: Int)
+ numUsableCores: Int,
+ securityContext: Option[RpcSecurityContext])
diff --git
a/common/src/main/scala/org/apache/celeborn/common/rpc/RpcSecurityContext.scala
b/common/src/main/scala/org/apache/celeborn/common/rpc/RpcSecurityContext.scala
new file mode 100644
index 000000000..34176e260
--- /dev/null
+++
b/common/src/main/scala/org/apache/celeborn/common/rpc/RpcSecurityContext.scala
@@ -0,0 +1,161 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.celeborn.common.security
+
+import org.apache.celeborn.common.network.sasl.{SaslCredentials,
SecretRegistry}
+import org.apache.celeborn.common.network.sasl.registration.RegistrationInfo
+
+/**
+ * Represents the security context, combining both client and server contexts.
+ *
+ * @param clientSaslContext Optional client sasl context.
+ * @param serverSaslContext Optional server sasl context.
+ */
+private[celeborn] case class RpcSecurityContext(
+ clientSaslContext: Option[ClientSaslContext] = None,
+ serverSaslContext: Option[ServerSaslContext] = None)
+
+/**
+ * Represents the SASL context.
+ */
+private[celeborn] trait SaslContext {}
+
+/**
+ * Represents the client SASL context.
+ * @param appId The application id.
+ * @param saslCredentials sasl credentials.
+ * @param addRegistrationBootstrap Whether to add registration bootstrap.
+ */
+private[celeborn] case class ClientSaslContext(
+ appId: String,
+ saslCredentials: SaslCredentials,
+ addRegistrationBootstrap: Boolean = false,
+ registrationInfo: RegistrationInfo = null) extends SaslContext
+
+/**
+ * Represents the server SASL context.
+ * @param secretRegistry The secret registry.
+ * @param addRegistrationBootstrap Whether to add registration bootstrap.
+ */
+private[celeborn] case class ServerSaslContext(
+ secretRegistry: SecretRegistry,
+ addRegistrationBootstrap: Boolean = false) extends SaslContext
+
+/**
+ * Builder for [[ClientSaslContext]].
+ */
+private[celeborn] class ClientSaslContextBuilder {
+ private var saslUser: String = _
+ private var saslPassword: String = _
+ private var appId: String = _
+ private var addRegistrationBootstrap: Boolean = false
+ private var registrationInfo: RegistrationInfo = _
+
+ def withSaslUser(user: String): ClientSaslContextBuilder = {
+ this.saslUser = user
+ this
+ }
+
+ def withSaslPassword(password: String): ClientSaslContextBuilder = {
+ this.saslPassword = password
+ this
+ }
+
+ def withAppId(appId: String): ClientSaslContextBuilder = {
+ this.appId = appId
+ this
+ }
+
+ def withAddRegistrationBootstrap(addRegistrationBootstrap: Boolean):
ClientSaslContextBuilder = {
+ this.addRegistrationBootstrap = addRegistrationBootstrap
+ this
+ }
+
+ def withRegistrationInfo(registrationInfo: RegistrationInfo):
ClientSaslContextBuilder = {
+ this.registrationInfo = registrationInfo
+ this
+ }
+
+ def build(): ClientSaslContext = {
+ if (saslUser == null || saslPassword == null) {
+ throw new IllegalArgumentException("Sasl user/password is not set.")
+ }
+ if (appId == null) {
+ throw new IllegalArgumentException("App id is not set.")
+ }
+ if (addRegistrationBootstrap && registrationInfo == null) {
+ throw new IllegalArgumentException("Registration info is not set.")
+ }
+ ClientSaslContext(
+ appId,
+ new SaslCredentials(saslUser, saslPassword),
+ addRegistrationBootstrap,
+ registrationInfo)
+ }
+}
+
+/**
+ * Builder for [[ServerSaslContext]].
+ */
+private[celeborn] class ServerSaslContextBuilder {
+ private var secretRegistry: SecretRegistry = _
+ private var addRegistrationBootstrap: Boolean = false
+
+ def withSecretRegistry(secretRegistry: SecretRegistry):
ServerSaslContextBuilder = {
+ this.secretRegistry = secretRegistry
+ this
+ }
+
+ def withAddRegistrationBootstrap(addRegistrationBootstrap: Boolean):
ServerSaslContextBuilder = {
+ this.addRegistrationBootstrap = addRegistrationBootstrap
+ this
+ }
+
+ def build(): ServerSaslContext = {
+ if (secretRegistry == null) {
+ throw new IllegalArgumentException("Secret registry is not set.")
+ }
+ ServerSaslContext(
+ secretRegistry,
+ addRegistrationBootstrap)
+ }
+}
+
+/**
+ * Builder for [[RpcSecurityContext]].
+ */
+private[celeborn] class RpcSecurityContextBuilder {
+ private var clientSaslContext: Option[ClientSaslContext] = None
+ private var serverSaslContext: Option[ServerSaslContext] = None
+
+ def withClientSaslContext(context: ClientSaslContext):
RpcSecurityContextBuilder = {
+ this.clientSaslContext = Some(context)
+ this
+ }
+
+ def withServerSaslContext(context: ServerSaslContext):
RpcSecurityContextBuilder = {
+ this.serverSaslContext = Some(context)
+ this
+ }
+
+ def build(): RpcSecurityContext = {
+ if (clientSaslContext.nonEmpty && serverSaslContext.nonEmpty) {
+ throw new IllegalArgumentException("Both client and server sasl context
cannot be set.")
+ }
+ RpcSecurityContext(clientSaslContext, serverSaslContext)
+ }
+}
diff --git
a/common/src/main/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnv.scala
b/common/src/main/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnv.scala
index 293e113cb..712ca8245 100644
---
a/common/src/main/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnv.scala
+++
b/common/src/main/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnv.scala
@@ -24,19 +24,19 @@ import java.util.concurrent._
import java.util.concurrent.atomic.AtomicBoolean
import javax.annotation.Nullable
+import scala.collection.JavaConverters.seqAsJavaListConverter
import scala.concurrent.{Future, Promise}
import scala.reflect.ClassTag
import scala.util.{DynamicVariable, Failure, Success}
import scala.util.control.NonFatal
-import com.google.common.base.Throwables
-
import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.internal.Logging
import org.apache.celeborn.common.network.TransportContext
-import org.apache.celeborn.common.network.buffer.NioManagedBuffer
import org.apache.celeborn.common.network.client._
-import org.apache.celeborn.common.network.protocol.{RequestMessage =>
NRequestMessage, RpcFailure => NRpcFailure, RpcRequest}
+import org.apache.celeborn.common.network.protocol.{RequestMessage =>
NRequestMessage, RpcRequest}
+import org.apache.celeborn.common.network.sasl.{SaslClientBootstrap,
SaslServerBootstrap}
+import
org.apache.celeborn.common.network.sasl.registration.{RegistrationClientBootstrap,
RegistrationServerBootstrap}
import org.apache.celeborn.common.network.server._
import org.apache.celeborn.common.protocol.{RpcNameConstants,
TransportModuleConstants}
import org.apache.celeborn.common.rpc._
@@ -48,6 +48,7 @@ class NettyRpcEnv(
javaSerializerInstance: JavaSerializerInstance) extends RpcEnv(config)
with Logging {
val celebornConf = config.conf
+ val securityContext = config.securityContext
private[celeborn] val transportConf = Utils.fromCelebornConf(
celebornConf.clone,
@@ -61,7 +62,27 @@ class NettyRpcEnv(
private val transportContext =
new TransportContext(transportConf, new NettyRpcHandler(dispatcher, this))
- val clientFactory = transportContext.createClientFactory()
+ private def createClientBootstraps():
java.util.List[TransportClientBootstrap] = {
+ val bootstrapOpt = securityContext.flatMap(_.clientSaslContext.map {
clientSaslContext =>
+ if (clientSaslContext.addRegistrationBootstrap) {
+ logInfo("Add registration client bootstrap")
+ new RegistrationClientBootstrap(
+ transportConf,
+ clientSaslContext.appId,
+ clientSaslContext.saslCredentials,
+ clientSaslContext.registrationInfo)
+ } else {
+ logInfo("Add sasl client bootstrap")
+ new SaslClientBootstrap(
+ transportConf,
+ clientSaslContext.appId,
+ clientSaslContext.saslCredentials)
+ }
+ })
+ bootstrapOpt.toList.asJava
+ }
+
+ val clientFactory =
transportContext.createClientFactory(createClientBootstraps())
private val timeoutScheduler =
ThreadUtils.newDaemonSingleThreadScheduledExecutor("celeborn-netty-rpc-env-timeout-checker")
@@ -92,8 +113,25 @@ class NettyRpcEnv(
}
}
+ private def createServerBootstraps():
java.util.List[TransportServerBootstrap] = {
+ val bootstrapOpt = securityContext.flatMap(_.serverSaslContext.map {
serverSaslContext =>
+ if (serverSaslContext.addRegistrationBootstrap) {
+ logInfo("Add registration server bootstrap")
+ new RegistrationServerBootstrap(
+ transportConf,
+ serverSaslContext.secretRegistry)
+ } else {
+ logInfo("Add sasl server bootstrap")
+ new SaslServerBootstrap(
+ transportConf,
+ serverSaslContext.secretRegistry)
+ }
+ })
+ bootstrapOpt.toList.asJava
+ }
+
def startServer(bindAddress: String, port: Int): Unit = {
- server = transportContext.createServer(bindAddress, port)
+ server = transportContext.createServer(bindAddress, port,
createServerBootstraps())
dispatcher.registerRpcEndpoint(
RpcEndpointVerifier.NAME,
new RpcEndpointVerifier(this, dispatcher))
diff --git
a/common/src/test/java/org/apache/celeborn/common/network/sasl/CelebornSaslSuiteJ.java
b/common/src/test/java/org/apache/celeborn/common/network/sasl/CelebornSaslSuiteJ.java
index b7d328541..b97da2e86 100644
---
a/common/src/test/java/org/apache/celeborn/common/network/sasl/CelebornSaslSuiteJ.java
+++
b/common/src/test/java/org/apache/celeborn/common/network/sasl/CelebornSaslSuiteJ.java
@@ -43,7 +43,7 @@ public class CelebornSaslSuiteJ extends SaslTestBase {
new CelebornSaslServer(
DIGEST_MD5,
DEFAULT_SASL_SERVER_PROPS,
- new
CelebornSaslServer.DigestCallbackHandler(SecretRegistryImpl.getInstance()));
+ new CelebornSaslServer.DigestCallbackHandler(secretRegistry));
assertFalse(client.isComplete());
assertFalse(server.isComplete());
@@ -73,7 +73,7 @@ public class CelebornSaslSuiteJ extends SaslTestBase {
new CelebornSaslServer(
DIGEST_MD5,
DEFAULT_SASL_SERVER_PROPS,
- new
CelebornSaslServer.DigestCallbackHandler(SecretRegistryImpl.getInstance()));
+ new CelebornSaslServer.DigestCallbackHandler(secretRegistry));
assertFalse(client.isComplete());
assertFalse(server.isComplete());
@@ -95,8 +95,7 @@ public class CelebornSaslSuiteJ extends SaslTestBase {
@Test
public void testSaslAuth() throws Throwable {
TransportConf conf = new TransportConf("shuffle", new CelebornConf());
- SaslServerBootstrap serverBootstrap =
- new SaslServerBootstrap(conf, SecretRegistryImpl.getInstance());
+ SaslServerBootstrap serverBootstrap = new SaslServerBootstrap(conf,
secretRegistry);
SaslClientBootstrap clientBootstrap =
new SaslClientBootstrap(conf, TEST_USER, new
SaslCredentials(TEST_USER, TEST_SECRET));
authHelper(conf, serverBootstrap, clientBootstrap);
diff --git
a/common/src/test/java/org/apache/celeborn/common/network/sasl/RegistrationSuiteJ.java
b/common/src/test/java/org/apache/celeborn/common/network/sasl/RegistrationSuiteJ.java
index 8240324bd..f41ded84d 100644
---
a/common/src/test/java/org/apache/celeborn/common/network/sasl/RegistrationSuiteJ.java
+++
b/common/src/test/java/org/apache/celeborn/common/network/sasl/RegistrationSuiteJ.java
@@ -56,7 +56,7 @@ public class RegistrationSuiteJ extends SaslTestBase {
// The SecretRegistryImpl already has the entry for TEST_USER so
re-registering the app should
// fail.
RegistrationServerBootstrap serverBootstrap =
- new RegistrationServerBootstrap(conf,
SecretRegistryImpl.getInstance());
+ new RegistrationServerBootstrap(conf, secretRegistry);
RegistrationClientBootstrap clientBootstrap =
new RegistrationClientBootstrap(
conf, TEST_USER, new SaslCredentials(TEST_USER, TEST_SECRET), new
RegistrationInfo());
diff --git
a/common/src/test/java/org/apache/celeborn/common/network/sasl/SaslTestBase.java
b/common/src/test/java/org/apache/celeborn/common/network/sasl/SaslTestBase.java
index 561af5cd6..c6ca51831 100644
---
a/common/src/test/java/org/apache/celeborn/common/network/sasl/SaslTestBase.java
+++
b/common/src/test/java/org/apache/celeborn/common/network/sasl/SaslTestBase.java
@@ -43,14 +43,16 @@ import org.apache.celeborn.common.util.JavaUtils;
public class SaslTestBase {
+ protected static final SecretRegistry secretRegistry = new
SecretRegistryImpl();
+
@BeforeClass
public static void setup() {
- SecretRegistryImpl.getInstance().register(TEST_USER, TEST_SECRET);
+ secretRegistry.register(TEST_USER, TEST_SECRET);
}
@AfterClass
public static void teardown() {
- SecretRegistryImpl.getInstance().unregister(TEST_USER);
+ secretRegistry.unregister(TEST_USER);
}
static final String TEST_USER = "appId";
diff --git
a/common/src/test/scala/org/apache/celeborn/common/rpc/RpcSecurityContextSuite.scala
b/common/src/test/scala/org/apache/celeborn/common/rpc/RpcSecurityContextSuite.scala
new file mode 100644
index 000000000..5ea2fa0d4
--- /dev/null
+++
b/common/src/test/scala/org/apache/celeborn/common/rpc/RpcSecurityContextSuite.scala
@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.common.rpc
+
+import org.scalatest.matchers.must.Matchers.be
+import org.scalatest.matchers.should.Matchers.{an, convertToAnyShouldWrapper}
+
+import org.apache.celeborn.CelebornFunSuite
+import org.apache.celeborn.common.network.sasl.{SaslCredentials,
SecretRegistryImpl}
+import org.apache.celeborn.common.network.sasl.registration.RegistrationInfo
+import org.apache.celeborn.common.security.{ClientSaslContext,
ClientSaslContextBuilder, RpcSecurityContextBuilder, ServerSaslContext,
ServerSaslContextBuilder}
+
+class RpcSecurityContextSuite extends CelebornFunSuite {
+
+ test("RpcSecurityContext should be created with either client and server
sasl contexts") {
+ val clientContext = ClientSaslContext(
+ "clientAppId",
+ new SaslCredentials("user", "password"),
+ addRegistrationBootstrap = true)
+
+ val rpcSecurityContext = new RpcSecurityContextBuilder()
+ .withClientSaslContext(clientContext)
+ .build()
+
+ rpcSecurityContext.clientSaslContext shouldBe Some(clientContext)
+ rpcSecurityContext.serverSaslContext shouldBe None
+ }
+
+ test("RpcSecurityContext should be created with only client sasl context") {
+ val clientContext = ClientSaslContext(
+ "clientAppId",
+ new SaslCredentials("user", "password"),
+ addRegistrationBootstrap = true)
+
+ val rpcSecurityContext = new RpcSecurityContextBuilder()
+ .withClientSaslContext(clientContext)
+ .build()
+
+ rpcSecurityContext.clientSaslContext shouldBe Some(clientContext)
+ rpcSecurityContext.serverSaslContext shouldBe None
+ }
+
+ test("RpcSecurityContext should be created with only server sasl context") {
+ val serverContext = ServerSaslContext(new SecretRegistryImpl())
+
+ val rpcSecurityContext = new RpcSecurityContextBuilder()
+ .withServerSaslContext(serverContext)
+ .build()
+
+ rpcSecurityContext.clientSaslContext shouldBe None
+ rpcSecurityContext.serverSaslContext shouldBe Some(serverContext)
+ }
+
+ test("ClientSaslContext build with valid parameters") {
+ val clientContext = new ClientSaslContextBuilder()
+ .withSaslUser("user")
+ .withSaslPassword("password")
+ .withAppId("clientAppId")
+ .withAddRegistrationBootstrap(true)
+ .withRegistrationInfo(new RegistrationInfo())
+ .build()
+
+ clientContext.appId shouldBe "clientAppId"
+ clientContext.saslCredentials.getUserId shouldBe "user"
+ clientContext.saslCredentials.getPassword shouldBe "password"
+ clientContext.addRegistrationBootstrap shouldBe true
+ clientContext.registrationInfo shouldNot be(null)
+ }
+
+ test("ClientSaslContext build should throw IllegalArgumentException when
sasl user/password is not set") {
+ an[IllegalArgumentException] should be thrownBy {
+ new ClientSaslContextBuilder()
+ .withAppId("clientAppId")
+ .withAddRegistrationBootstrap(true)
+ .withRegistrationInfo(new RegistrationInfo())
+ .build()
+ }
+ }
+
+ test("ClientSaslContext build should throw IllegalArgumentException when app
id is not set") {
+ an[IllegalArgumentException] should be thrownBy {
+ new ClientSaslContextBuilder()
+ .withSaslUser("user")
+ .withSaslPassword("password")
+ .withAddRegistrationBootstrap(true)
+ .withRegistrationInfo(new RegistrationInfo())
+ .build()
+ }
+ }
+
+ test("ClientSaslContext build should throw IllegalArgumentException when
addRegistrationBootstrap is true but registration info is not set") {
+ an[IllegalArgumentException] should be thrownBy {
+ new ClientSaslContextBuilder()
+ .withSaslUser("user")
+ .withSaslPassword("password")
+ .withAppId("clientAppId")
+ .withAddRegistrationBootstrap(true)
+ .build()
+ }
+ }
+
+ test("ServerSaslContext build should build ServerSaslContext with valid
parameters") {
+ val serverContext = new ServerSaslContextBuilder()
+ .withSecretRegistry(new SecretRegistryImpl())
+ .withAddRegistrationBootstrap(true)
+ .build()
+
+ serverContext.secretRegistry shouldNot be(null)
+ serverContext.addRegistrationBootstrap shouldBe true
+ }
+
+ test(
+ "ServerSaslContext build should throw IllegalArgumentException when secret
registry is not set") {
+ an[IllegalArgumentException] should be thrownBy {
+ new ServerSaslContextBuilder()
+ .withAddRegistrationBootstrap(true)
+ .build()
+ }
+ }
+
+}
diff --git
a/common/src/test/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnvSuite.scala
b/common/src/test/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnvSuite.scala
index 8afbf5980..4e0fa5300 100644
---
a/common/src/test/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnvSuite.scala
+++
b/common/src/test/scala/org/apache/celeborn/common/rpc/netty/NettyRpcEnvSuite.scala
@@ -39,7 +39,7 @@ class NettyRpcEnvSuite extends RpcEnvSuite with TimeLimits {
name: String,
port: Int,
clientMode: Boolean = false): RpcEnv = {
- val config = RpcEnvConfig(conf, "test", "localhost", "localhost", port, 0)
+ val config = RpcEnvConfig(conf, "test", "localhost", "localhost", port, 0,
None)
new NettyRpcEnvFactory().create(config)
}
@@ -54,7 +54,7 @@ class NettyRpcEnvSuite extends RpcEnvSuite with TimeLimits {
test("advertise address different from bind address") {
val celebornConf = createCelebornConf()
- val config = RpcEnvConfig(celebornConf, "test", "localhost",
"example.com", 0, 0)
+ val config = RpcEnvConfig(celebornConf, "test", "localhost",
"example.com", 0, 0, None)
val env = new NettyRpcEnvFactory().create(config)
try {
assert(env.address.hostPort.startsWith("example.com:"))
@@ -101,7 +101,8 @@ class NettyRpcEnvSuite extends RpcEnvSuite with TimeLimits {
"localhost",
"localhost",
0,
- numUsableCores)
+ numUsableCores,
+ None)
val anotherEnv = new NettyRpcEnvFactory().create(config)
anotherEnv.setupEndpoint(
"StackOverflowError",
