This is an automated email from the ASF dual-hosted git repository.

zhouky 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 ab4c0bc85 [CELEBORN-1257] Adds a secured port in Celeborn Master for 
secure communication with LifecycleManager
ab4c0bc85 is described below

commit ab4c0bc85bcdd7add792f8bcdc0f6110b088b9e3
Author: Chandni Singh <[email protected]>
AuthorDate: Tue Feb 6 14:53:28 2024 +0800

    [CELEBORN-1257] Adds a secured port in Celeborn Master for secure 
communication with LifecycleManager
    
    ### What changes were proposed in this pull request?
    This adds a secured port to Celeborn Master which is used for secure 
communication with LifecycleManager.
    This is part of adding authentication support in Celeborn (see 
CELEBORN-1011).
    
    This change targets just adding the secured port to Master. The following 
items from the proposal are still pending:
    1. Persisting the app secrets in Ratis.
    2. Forwarding secrets to Workers and having ability for the workers to pull 
registration info from the Master.
    3. Secured and internal port in Workers.
    4. Secured communication between workers and clients.
    
    In addition, since we are supporting both secured and unsecured 
communication for backward compatibility and seamless rolling upgrades, there 
is an additional change needed. An app which registers with the Master can try 
to talk to the workers on unsecured ports which is a security breach. So, the 
workers need to know whether an app registered with Master or not and for that 
Master has to propagate list of un-secured apps to Celeborn workers as well. We 
can discuss this more with http [...]
    
    ### Why are the changes needed?
    It is needed for adding authentication support to Celeborn (CELEBORN-1011)
    
    ### Does this PR introduce _any_ user-facing change?
    Yes
    
    ### How was this patch tested?
    Added a simple UT.
    
    Closes #2281 from otterc/CELEBORN-1257.
    
    Authored-by: Chandni Singh <[email protected]>
    Signed-off-by: zky.zhoukeyong <[email protected]>
---
 .../apache/celeborn/client/LifecycleManager.scala  |  74 +++++++++++++-
 .../celeborn/common/client/MasterClient.java       |  33 ++++++-
 .../celeborn/common/protocol/RpcNameConstants.java |   4 +
 .../org/apache/celeborn/common/CelebornConf.scala  | 103 ++++++++++++++++---
 .../celeborn/common/client/MasterClientSuiteJ.java |  14 +--
 docs/configuration/client.md                       |   1 +
 docs/configuration/ha.md                           |   1 +
 docs/configuration/master.md                       |   1 +
 docs/configuration/worker.md                       |   1 +
 .../celeborn/service/deploy/master/Master.scala    | 110 +++++++++++++++++----
 .../service/deploy/master/MasterArguments.scala    |  12 +++
 .../service/deploy/master/SecuredRpcEndpoint.scala |  91 +++++++++++++++++
 .../service/deploy/master/MasterSuite.scala        |  43 +++++++-
 .../celeborn/service/deploy/worker/Worker.scala    |   3 +-
 14 files changed, 444 insertions(+), 47 deletions(-)

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 48b227343..ad23d1dbc 100644
--- a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
@@ -17,7 +17,9 @@
 
 package org.apache.celeborn.client
 
+import java.lang.{Byte => JByte}
 import java.nio.ByteBuffer
+import java.security.SecureRandom
 import java.util
 import java.util.{function, List => JList}
 import java.util.concurrent.{Callable, ConcurrentHashMap, LinkedBlockingQueue, 
ScheduledFuture, TimeUnit}
@@ -41,12 +43,14 @@ import org.apache.celeborn.common.client.MasterClient
 import org.apache.celeborn.common.identity.{IdentityProvider, UserIdentifier}
 import org.apache.celeborn.common.internal.Logging
 import org.apache.celeborn.common.meta.{ShufflePartitionLocationInfo, 
WorkerInfo}
+import org.apache.celeborn.common.network.sasl.registration.RegistrationInfo
 import org.apache.celeborn.common.protocol._
 import org.apache.celeborn.common.protocol.RpcNameConstants.WORKER_EP
 import org.apache.celeborn.common.protocol.message.ControlMessages._
 import org.apache.celeborn.common.protocol.message.StatusCode
 import org.apache.celeborn.common.rpc._
 import org.apache.celeborn.common.rpc.netty.{LocalNettyRpcCallContext, 
RemoteNettyRpcCallContext}
+import org.apache.celeborn.common.security.{ClientSaslContextBuilder, 
RpcSecurityContext, RpcSecurityContextBuilder}
 import org.apache.celeborn.common.util.{JavaUtils, PbSerDeUtils, ThreadUtils, 
Utils}
 // Can Remove this if celeborn don't support scala211 in future
 import org.apache.celeborn.common.util.FunctionConverter._
@@ -108,6 +112,7 @@ class LifecycleManager(val appUniqueId: String, val conf: 
CelebornConf) extends
     .build().asInstanceOf[Cache[Int, ByteBuffer]]
 
   private val mockDestroyFailure = conf.testMockDestroySlotsFailure
+  private val authEnabled = conf.authEnabled
 
   @VisibleForTesting
   def workerSnapshots(shuffleId: Int): util.Map[WorkerInfo, 
ShufflePartitionLocationInfo] =
@@ -159,7 +164,32 @@ class LifecycleManager(val appUniqueId: String, val conf: 
CelebornConf) extends
 
   logInfo(s"Starting LifecycleManager on ${rpcEnv.address}")
 
-  private val masterClient = new MasterClient(rpcEnv, conf)
+  private var masterRpcEnvInUse = rpcEnv
+  private var workerRpcEnvInUse = rpcEnv
+  if (authEnabled) {
+    logInfo(s"Authentication is enabled; setting up master and worker RPC 
environments")
+    val appSecret = createSecret()
+    val registrationInfo = new RegistrationInfo()
+    masterRpcEnvInUse =
+      RpcEnv.create(
+        RpcNameConstants.LIFECYCLE_MANAGER_MASTER_SYS,
+        lifecycleHost,
+        0,
+        conf,
+        createRpcSecurityContext(
+          appSecret,
+          addClientRegistrationBootstrap = true,
+          Some(registrationInfo)))
+    workerRpcEnvInUse =
+      RpcEnv.create(
+        RpcNameConstants.LIFECYCLE_MANAGER_WORKER_SYS,
+        lifecycleHost,
+        0,
+        conf,
+        createRpcSecurityContext(appSecret))
+  }
+
+  private val masterClient = new MasterClient(masterRpcEnvInUse, conf, false)
   val commitManager = new CommitManager(appUniqueId, conf, this)
   val workerStatusTracker = new WorkerStatusTracker(conf, this)
   private val heartbeater =
@@ -214,6 +244,36 @@ class LifecycleManager(val appUniqueId: String, val conf: 
CelebornConf) extends
       rpcEnv.shutdown()
       rpcEnv.awaitTermination()
     }
+    if (authEnabled) {
+      if (masterRpcEnvInUse != null) {
+        masterRpcEnvInUse.shutdown()
+        masterRpcEnvInUse.awaitTermination()
+      }
+      if (workerRpcEnvInUse != null) {
+        workerRpcEnvInUse.shutdown()
+        workerRpcEnvInUse.awaitTermination()
+      }
+    }
+  }
+
+  /**
+   * Creates security context for external RPC endpoint.
+   */
+  def createRpcSecurityContext(
+      appSecret: String,
+      addClientRegistrationBootstrap: Boolean = false,
+      registrationInfo: Option[RegistrationInfo] = None): 
Option[RpcSecurityContext] = {
+    val clientSaslContextBuilder = new ClientSaslContextBuilder()
+      .withAddRegistrationBootstrap(addClientRegistrationBootstrap)
+      .withAppId(appUniqueId)
+      .withSaslUser(appUniqueId)
+      .withSaslPassword(appSecret)
+    if (registrationInfo.isDefined) {
+      clientSaslContextBuilder.withRegistrationInfo(registrationInfo.get)
+    }
+    val rpcSecurityContext = new RpcSecurityContextBuilder()
+      .withClientSaslContext(clientSaslContextBuilder.build()).build()
+    Some(rpcSecurityContext)
   }
 
   def getUserIdentifier: UserIdentifier = {
@@ -356,7 +416,7 @@ class LifecycleManager(val appUniqueId: String, val conf: 
CelebornConf) extends
       connectFailedWorkers: ShuffleFailedWorkers): Unit = {
     val futures = new util.LinkedList[(Future[RpcEndpointRef], WorkerInfo)]()
     slots.asScala foreach { case (workerInfo, _) =>
-      val future = rpcEnv.asyncSetupEndpointRefByAddr(RpcEndpointAddress(
+      val future = 
workerRpcEnvInUse.asyncSetupEndpointRefByAddr(RpcEndpointAddress(
         RpcAddress.apply(workerInfo.host, workerInfo.rpcPort),
         WORKER_EP))
       futures.add((future, workerInfo))
@@ -1065,7 +1125,7 @@ class LifecycleManager(val appUniqueId: String, val conf: 
CelebornConf) extends
               s" ${destroyWorkerInfo.readableAddress()}, init according to 
partition info")
             try {
               if (!workerStatusTracker.workerExcluded(destroyWorkerInfo)) {
-                destroyWorkerInfo.endpoint = rpcEnv.setupEndpointRef(
+                destroyWorkerInfo.endpoint = 
workerRpcEnvInUse.setupEndpointRef(
                   RpcAddress.apply(destroyWorkerInfo.host, 
destroyWorkerInfo.rpcPort),
                   WORKER_EP)
               } else {
@@ -1573,4 +1633,12 @@ class LifecycleManager(val appUniqueId: String, val 
conf: CelebornConf) extends
     heartbeater.stop()
     super.stop()
   }
+
+  private def createSecret(): String = {
+    val bits = 256
+    val rnd = new SecureRandom()
+    val secretBytes = new Array[Byte](bits / JByte.SIZE)
+    rnd.nextBytes(secretBytes)
+    JavaUtils.bytesToString(ByteBuffer.wrap(secretBytes))
+  }
 }
diff --git 
a/common/src/main/java/org/apache/celeborn/common/client/MasterClient.java 
b/common/src/main/java/org/apache/celeborn/common/client/MasterClient.java
index 83a727477..90db88a16 100644
--- a/common/src/main/java/org/apache/celeborn/common/client/MasterClient.java
+++ b/common/src/main/java/org/apache/celeborn/common/client/MasterClient.java
@@ -58,10 +58,15 @@ public class MasterClient {
 
   private final AtomicReference<RpcEndpointRef> rpcEndpointRef;
   private final ExecutorService oneWayMessageSender;
+  private final CelebornConf conf;
+  private final boolean isWorker;
+  private String masterEndpointName;
 
-  public MasterClient(RpcEnv rpcEnv, CelebornConf conf) {
+  public MasterClient(RpcEnv rpcEnv, CelebornConf conf, boolean isWorker) {
     this.rpcEnv = rpcEnv;
-    this.masterEndpoints = Arrays.asList(conf.masterEndpoints());
+    this.conf = conf;
+    this.isWorker = isWorker;
+    this.masterEndpoints = resolveMasterEndpoints();
     Collections.shuffle(this.masterEndpoints);
     this.maxRetries = Math.max(masterEndpoints.size(), 
conf.masterClientMaxRetries());
     this.rpcTimeout = conf.masterClientRpcAskTimeout();
@@ -250,7 +255,7 @@ public class MasterClient {
     RpcEndpointRef endpointRef = null;
     try {
       endpointRef =
-          rpcEnv.setupEndpointRef(RpcAddress.fromHostAndPort(endpoint), 
RpcNameConstants.MASTER_EP);
+          rpcEnv.setupEndpointRef(RpcAddress.fromHostAndPort(endpoint), 
masterEndpointName);
     } catch (Exception e) {
       // Catch all exceptions. Because we don't care whether this exception is 
IOException or
       // TimeoutException or other exceptions, so we just try to connect to 
host:port, if fail,
@@ -259,4 +264,26 @@ public class MasterClient {
     }
     return endpointRef;
   }
+
+  private List<String> resolveMasterEndpoints() {
+    if (isWorker) {
+      // For worker, we should use the internal endpoints if internal port is 
enabled.
+      if (conf.internalPortEnabled()) {
+        masterEndpointName = RpcNameConstants.MASTER_INTERNAL_EP;
+        return Arrays.asList(conf.masterInternalEndpoints());
+      } else {
+        masterEndpointName = RpcNameConstants.MASTER_EP;
+        return Arrays.asList(conf.masterEndpoints());
+      }
+    } else {
+      // This is for client, so we should use the secured endpoints if auth is 
enabled.
+      if (conf.authEnabled()) {
+        masterEndpointName = RpcNameConstants.MASTER_SECURED_EP;
+        return Arrays.asList(conf.masterSecuredEndpoints());
+      } else {
+        masterEndpointName = RpcNameConstants.MASTER_EP;
+        return Arrays.asList(conf.masterEndpoints());
+      }
+    }
+  }
 }
diff --git 
a/common/src/main/java/org/apache/celeborn/common/protocol/RpcNameConstants.java
 
b/common/src/main/java/org/apache/celeborn/common/protocol/RpcNameConstants.java
index 6c3898fdc..fed7b827b 100644
--- 
a/common/src/main/java/org/apache/celeborn/common/protocol/RpcNameConstants.java
+++ 
b/common/src/main/java/org/apache/celeborn/common/protocol/RpcNameConstants.java
@@ -21,9 +21,11 @@ public class RpcNameConstants {
   // For Master
   public static String MASTER_SYS = "Master";
   public static String MASTER_INTERNAL_SYS = "MasterInternal";
+  public static String MASTER_SECURED_SYS = "MasterSecured";
   // Master Endpoint Name
   public static String MASTER_EP = "MasterEndpoint";
   public static String MASTER_INTERNAL_EP = "MasterInternalEndpoint";
+  public static String MASTER_SECURED_EP = "MasterSecuredEndpoint";
 
   // For Worker
   public static String WORKER_SYS = "Worker";
@@ -32,6 +34,8 @@ public class RpcNameConstants {
 
   // For LifecycleManager
   public static String LIFECYCLE_MANAGER_SYS = "LifecycleManager";
+  public static String LIFECYCLE_MANAGER_MASTER_SYS = 
"LifecycleManagerMasterSys";
+  public static String LIFECYCLE_MANAGER_WORKER_SYS = 
"LifecycleManagerWorkerSys";
   public static String LIFECYCLE_MANAGER_EP = "LifecycleManagerEndpoint";
 
   // For Shuffle Client
diff --git 
a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala 
b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index 822fe5459..c4aad0306 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -26,6 +26,7 @@ import scala.collection.mutable
 import scala.concurrent.duration._
 import scala.util.Try
 
+import org.apache.celeborn.common.CelebornConf.MASTER_INTERNAL_ENDPOINTS
 import org.apache.celeborn.common.identity.{DefaultIdentityProvider, 
IdentityProvider}
 import org.apache.celeborn.common.internal.Logging
 import org.apache.celeborn.common.internal.config._
@@ -1133,13 +1134,44 @@ class CelebornConf(loadDefaults: Boolean) extends 
Cloneable with Logging with Se
   // //////////////////////////////////////////////////////
   //               Authentication                        //
   // //////////////////////////////////////////////////////
-  def authEnabled: Boolean = get(AUTH_ENABLED)
+  def authEnabled: Boolean = {
+    val authEnabled = get(AUTH_ENABLED)
+    val internalPortEnabled = get(INTERNAL_PORT_ENABLED)
+    if (authEnabled && !internalPortEnabled) {
+      throw new IllegalArgumentException(
+        s"${AUTH_ENABLED.key} is true, but ${INTERNAL_PORT_ENABLED.key} is 
false")
+    }
+    return authEnabled && internalPortEnabled
+  }
+
+  def haMasterNodeSecuredPort(nodeId: String): Int = {
+    val key = HA_MASTER_NODE_SECURED_PORT.key.replace("<id>", nodeId)
+    getInt(key, HA_MASTER_NODE_SECURED_PORT.defaultValue.get)
+  }
+
+  def masterSecuredPort: Int = get(MASTER_SECURED_PORT)
+
+  def masterSecuredEndpoints: Array[String] =
+    get(MASTER_SECURED_ENDPOINTS).toArray.map { endpoint =>
+      Utils.parseHostPort(endpoint.replace("<localhost>", 
Utils.localHostName(this))) match {
+        case (host, 0) => 
s"$host:${HA_MASTER_NODE_SECURED_PORT.defaultValue.get}"
+        case (host, port) => s"$host:$port"
+      }
+    }
 
   // //////////////////////////////////////////////////////
   //                     Internal Port                   //
   // //////////////////////////////////////////////////////
   def internalPortEnabled: Boolean = get(INTERNAL_PORT_ENABLED)
 
+  def masterInternalEndpoints: Array[String] =
+    get(MASTER_INTERNAL_ENDPOINTS).toArray.map { endpoint =>
+      Utils.parseHostPort(endpoint.replace("<localhost>", 
Utils.localHostName(this))) match {
+        case (host, 0) => 
s"$host:${HA_MASTER_NODE_INTERNAL_PORT.defaultValue.get}"
+        case (host, port) => s"$host:$port"
+      }
+    }
+
   // //////////////////////////////////////////////////////
   //                     Rack Resolver                   //
   // //////////////////////////////////////////////////////
@@ -1147,8 +1179,7 @@ class CelebornConf(loadDefaults: Boolean) extends 
Cloneable with Logging with Se
 
   def haMasterNodeInternalPort(nodeId: String): Int = {
     val key = HA_MASTER_NODE_INTERNAL_PORT.key.replace("<id>", nodeId)
-    val legacyKey = 
HA_MASTER_NODE_INTERNAL_PORT.alternatives.head._1.replace("<id>", nodeId)
-    getInt(key, getInt(legacyKey, 
HA_MASTER_NODE_INTERNAL_PORT.defaultValue.get))
+    getInt(key, HA_MASTER_NODE_INTERNAL_PORT.defaultValue.get)
   }
 
   def masterInternalPort: Int = get(MASTER_INTERNAL_PORT)
@@ -4497,14 +4528,6 @@ object CelebornConf extends Logging {
       .timeConf(TimeUnit.MILLISECONDS)
       .createWithDefaultString("30s")
 
-  val AUTH_ENABLED: ConfigEntry[Boolean] =
-    buildConf("celeborn.auth.enabled")
-      .categories("auth")
-      .version("0.5.0")
-      .doc("Whether to enable authentication.")
-      .booleanConf
-      .createWithDefault(false)
-
   val INTERNAL_PORT_ENABLED: ConfigEntry[Boolean] =
     buildConf("celeborn.internal.port.enabled")
       .categories("master", "worker")
@@ -4516,6 +4539,15 @@ object CelebornConf extends Logging {
       .booleanConf
       .createWithDefault(false)
 
+  val AUTH_ENABLED: ConfigEntry[Boolean] =
+    buildConf("celeborn.auth.enabled")
+      .categories("auth")
+      .version("0.5.0")
+      .doc("Whether to enable authentication. Authentication will be enabled 
only when " +
+        s"${INTERNAL_PORT_ENABLED.key} is enabled as well.")
+      .booleanConf
+      .createWithDefault(false)
+
   val MASTER_INTERNAL_PORT: ConfigEntry[Int] =
     buildConf("celeborn.master.internal.port")
       .categories("master")
@@ -4536,6 +4568,20 @@ object CelebornConf extends Logging {
       .checkValue(p => p >= 1024 && p < 65535, "Invalid port")
       .createWithDefault(8097)
 
+  val MASTER_INTERNAL_ENDPOINTS: ConfigEntry[Seq[String]] =
+    buildConf("celeborn.master.internal.endpoints")
+      .categories("worker")
+      .doc("Endpoints of master nodes just for celeborn workers to connect, 
allowed pattern " +
+        "is: `<host1>:<port1>[,<host2>:<port2>]*`, e.g. 
`clb1:8097,clb2:8097,clb3:8097`. " +
+        "If the port is omitted, 8097 will be used.")
+      .version("0.5.0")
+      .stringConf
+      .toSequence
+      .checkValue(
+        endpoints => endpoints.map(_ => 
Try(Utils.parseHostPort(_))).forall(_.isSuccess),
+        "Allowed pattern is: `<host1>:<port1>[,<host2>:<port2>]*`")
+      .createWithDefaultString(s"<localhost>:8097")
+
   val RACKRESOLVER_REFRESH_INTERVAL: ConfigEntry[Long] =
     buildConf("celeborn.master.rackResolver.refresh.interval")
       .categories("master")
@@ -4543,4 +4589,39 @@ object CelebornConf extends Logging {
       .doc("Interval for refreshing the node rack information periodically.")
       .timeConf(TimeUnit.MILLISECONDS)
       .createWithDefaultString("30s")
+
+  val MASTER_SECURED_PORT: ConfigEntry[Int] =
+    buildConf("celeborn.master.secured.port")
+      .categories("master", "auth")
+      .version("0.5.0")
+      .doc(
+        "Secured port on the master where clients connect.")
+      .intConf
+      .checkValue(p => p >= 1024 && p < 65535, "Invalid port")
+      .createWithDefault(19097)
+
+  val HA_MASTER_NODE_SECURED_PORT: ConfigEntry[Int] =
+    buildConf("celeborn.master.ha.node.<id>.secured.port")
+      .categories("ha", "auth")
+      .doc(
+        "Secured port for the clients to bind to a master node <id> in HA 
mode.")
+      .version("0.5.0")
+      .intConf
+      .checkValue(p => p >= 1024 && p < 65535, "Invalid port")
+      .createWithDefault(19097)
+
+  val MASTER_SECURED_ENDPOINTS: ConfigEntry[Seq[String]] =
+    buildConf("celeborn.master.secured.endpoints")
+      .categories("client", "auth")
+      .doc("Endpoints of master nodes for celeborn client to connect for 
secured communication, allowed pattern " +
+        "is: `<host1>:<port1>[,<host2>:<port2>]*`, e.g. 
`clb1:19097,clb2:19097,clb3:19097`. " +
+        "If the port is omitted, 19097 will be used.")
+      .version("0.5.0")
+      .stringConf
+      .toSequence
+      .checkValue(
+        endpoints => endpoints.map(_ => 
Try(Utils.parseHostPort(_))).forall(_.isSuccess),
+        "Allowed pattern is: `<host1>:<port1>[,<host2>:<port2>]*`")
+      .createWithDefaultString(s"<localhost>:19097")
+
 }
diff --git 
a/common/src/test/java/org/apache/celeborn/common/client/MasterClientSuiteJ.java
 
b/common/src/test/java/org/apache/celeborn/common/client/MasterClientSuiteJ.java
index 4a84f7c5f..dacb32c41 100644
--- 
a/common/src/test/java/org/apache/celeborn/common/client/MasterClientSuiteJ.java
+++ 
b/common/src/test/java/org/apache/celeborn/common/client/MasterClientSuiteJ.java
@@ -79,7 +79,7 @@ public class MasterClientSuiteJ {
         });
     prepareForRpcEnvWithoutHA();
 
-    MasterClient client = new MasterClient(rpcEnv, conf);
+    MasterClient client = new MasterClient(rpcEnv, conf, false);
     HeartbeatFromApplication message = 
Mockito.mock(HeartbeatFromApplication.class);
 
     try {
@@ -106,7 +106,7 @@ public class MasterClientSuiteJ {
         });
     prepareForRpcEnvWithoutHA();
 
-    MasterClient client = new MasterClient(rpcEnv, conf);
+    MasterClient client = new MasterClient(rpcEnv, conf, false);
     HeartbeatFromApplication message = 
Mockito.mock(HeartbeatFromApplication.class);
 
     try {
@@ -132,7 +132,7 @@ public class MasterClientSuiteJ {
           return Future$.MODULE$.successful(response);
         });
 
-    MasterClient client = new MasterClient(rpcEnv, conf);
+    MasterClient client = new MasterClient(rpcEnv, conf, false);
     HeartbeatFromApplication message = 
Mockito.mock(HeartbeatFromApplication.class);
 
     try {
@@ -152,7 +152,7 @@ public class MasterClientSuiteJ {
     prepareForEndpointRefWithoutRetry(() -> 
Future$.MODULE$.successful(mockResponse));
     prepareForRpcEnvWithoutHA();
 
-    MasterClient client = new MasterClient(rpcEnv, conf);
+    MasterClient client = new MasterClient(rpcEnv, conf, false);
     HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
 
     HeartbeatFromWorkerResponse response = null;
@@ -174,7 +174,7 @@ public class MasterClientSuiteJ {
     prepareForEndpointRefWithRetry(numTries, () -> 
Future$.MODULE$.successful(mockResponse));
     prepareForRpcEnvWithoutHA();
 
-    MasterClient client = new MasterClient(rpcEnv, conf);
+    MasterClient client = new MasterClient(rpcEnv, conf, false);
     HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
 
     HeartbeatFromWorkerResponse response = null;
@@ -195,7 +195,7 @@ public class MasterClientSuiteJ {
 
     prepareForRpcEnvWithHA(() -> Future$.MODULE$.successful(mockResponse));
 
-    MasterClient client = new MasterClient(rpcEnv, conf);
+    MasterClient client = new MasterClient(rpcEnv, conf, false);
     HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
 
     HeartbeatFromWorkerResponse response = null;
@@ -254,7 +254,7 @@ public class MasterClientSuiteJ {
         .when(rpcEnv)
         .setupEndpointRef(Mockito.any(RpcAddress.class), Mockito.anyString());
 
-    MasterClient client = new MasterClient(rpcEnv, conf);
+    MasterClient client = new MasterClient(rpcEnv, conf, false);
     HeartbeatFromWorker message = Mockito.mock(HeartbeatFromWorker.class);
 
     HeartbeatFromWorkerResponse response = null;
diff --git a/docs/configuration/client.md b/docs/configuration/client.md
index 284edf157..3b4b8836f 100644
--- a/docs/configuration/client.md
+++ b/docs/configuration/client.md
@@ -110,6 +110,7 @@ license: |
 | celeborn.client.spark.shuffle.forceFallback.numPartitionsThreshold | 
2147483647 | Celeborn will only accept shuffle of partition number lower than 
this configuration value. | 0.3.0 | 
celeborn.shuffle.forceFallback.numPartitionsThreshold | 
 | celeborn.client.spark.shuffle.writer | HASH | Celeborn supports the 
following kind of shuffle writers. 1. hash: hash-based shuffle writer works 
fine when shuffle partition count is normal; 2. sort: sort-based shuffle writer 
works fine when memory pressure is high or shuffle partition count is huge. 
This configuration only takes effect when 
celeborn.client.spark.push.dynamicWriteMode.enabled is false. | 0.3.0 | 
celeborn.shuffle.writer | 
 | celeborn.master.endpoints | &lt;localhost&gt;:9097 | Endpoints of master 
nodes for celeborn client to connect, allowed pattern is: 
`<host1>:<port1>[,<host2>:<port2>]*`, e.g. `clb1:9097,clb2:9098,clb3:9099`. If 
the port is omitted, 9097 will be used. | 0.2.0 |  | 
+| celeborn.master.secured.endpoints | &lt;localhost&gt;:19097 | Endpoints of 
master nodes for celeborn client to connect for secured communication, allowed 
pattern is: `<host1>:<port1>[,<host2>:<port2>]*`, e.g. 
`clb1:19097,clb2:19097,clb3:19097`. If the port is omitted, 19097 will be used. 
| 0.5.0 |  | 
 | celeborn.storage.availableTypes | HDD | Enabled storages. Available options: 
MEMORY,HDD,SSD,HDFS. Note: HDD and SSD would be treated as identical. | 0.3.0 | 
celeborn.storage.activeTypes | 
 | celeborn.storage.hdfs.dir | &lt;undefined&gt; | HDFS base directory for 
Celeborn to store shuffle data. | 0.2.0 |  | 
 <!--end-include-->
diff --git a/docs/configuration/ha.md b/docs/configuration/ha.md
index faef0916e..c6ab6aeeb 100644
--- a/docs/configuration/ha.md
+++ b/docs/configuration/ha.md
@@ -24,6 +24,7 @@ license: |
 | celeborn.master.ha.node.&lt;id&gt;.internal.port | 8097 | Internal port for 
the workers and other masters to bind to a master node <id> in HA mode. | 0.5.0 
|  | 
 | celeborn.master.ha.node.&lt;id&gt;.port | 9097 | Port to bind of master node 
<id> in HA mode. | 0.3.0 | celeborn.ha.master.node.&lt;id&gt;.port | 
 | celeborn.master.ha.node.&lt;id&gt;.ratis.port | 9872 | Ratis port to bind of 
master node <id> in HA mode. | 0.3.0 | 
celeborn.ha.master.node.&lt;id&gt;.ratis.port | 
+| celeborn.master.ha.node.&lt;id&gt;.secured.port | 19097 | Secured port for 
the clients to bind to a master node <id> in HA mode. | 0.5.0 |  | 
 | celeborn.master.ha.ratis.raft.rpc.type | netty | RPC type for Ratis, 
available options: netty, grpc. | 0.3.0 | 
celeborn.ha.master.ratis.raft.rpc.type | 
 | celeborn.master.ha.ratis.raft.server.storage.dir | /tmp/ratis |  | 0.3.0 | 
celeborn.ha.master.ratis.raft.server.storage.dir | 
 <!--end-include-->
diff --git a/docs/configuration/master.md b/docs/configuration/master.md
index 091c362f5..a7d7fa292 100644
--- a/docs/configuration/master.md
+++ b/docs/configuration/master.md
@@ -45,6 +45,7 @@ license: |
 | celeborn.master.internal.port | 8097 | Internal port on the master where 
both workers and other master nodes connect. | 0.5.0 |  | 
 | celeborn.master.port | 9097 | Port for master to bind. | 0.2.0 |  | 
 | celeborn.master.rackResolver.refresh.interval | 30s | Interval for 
refreshing the node rack information periodically. | 0.5.0 |  | 
+| celeborn.master.secured.port | 19097 | Secured port on the master where 
clients connect. | 0.5.0 |  | 
 | celeborn.master.slot.assign.extraSlots | 2 | Extra slots number when master 
assign slots. | 0.3.0 | celeborn.slots.assign.extraSlots | 
 | celeborn.master.slot.assign.loadAware.diskGroupGradient | 0.1 | This value 
means how many more workload will be placed into a faster disk group than a 
slower group. | 0.3.0 | celeborn.slots.assign.loadAware.diskGroupGradient | 
 | celeborn.master.slot.assign.loadAware.fetchTimeWeight | 1.0 | Weight of 
average fetch time when calculating ordering in load-aware assignment strategy 
| 0.3.0 | celeborn.slots.assign.loadAware.fetchTimeWeight | 
diff --git a/docs/configuration/worker.md b/docs/configuration/worker.md
index c792840f0..0424e940f 100644
--- a/docs/configuration/worker.md
+++ b/docs/configuration/worker.md
@@ -35,6 +35,7 @@ license: |
 | celeborn.internal.port.enabled | false | Whether to create a internal port 
on Masters/Workers for inter-Masters/Workers communication. This is beneficial 
when SASL authentication is enforced for all interactions between clients and 
Celeborn Services, but the services can exchange messages without being subject 
to SASL authentication. | 0.5.0 |  | 
 | celeborn.master.endpoints | &lt;localhost&gt;:9097 | Endpoints of master 
nodes for celeborn client to connect, allowed pattern is: 
`<host1>:<port1>[,<host2>:<port2>]*`, e.g. `clb1:9097,clb2:9098,clb3:9099`. If 
the port is omitted, 9097 will be used. | 0.2.0 |  | 
 | celeborn.master.estimatedPartitionSize.minSize | 8mb | Ignore partition size 
smaller than this configuration of partition size for estimation. | 0.3.0 | 
celeborn.shuffle.minPartitionSizeToEstimate | 
+| celeborn.master.internal.endpoints | &lt;localhost&gt;:8097 | Endpoints of 
master nodes just for celeborn workers to connect, allowed pattern is: 
`<host1>:<port1>[,<host2>:<port2>]*`, e.g. `clb1:8097,clb2:8097,clb3:8097`. If 
the port is omitted, 8097 will be used. | 0.5.0 |  | 
 | celeborn.shuffle.chunk.size | 8m | Max chunk size of reducer's merged 
shuffle data. For example, if a reducer's shuffle data is 128M and the data 
will need 16 fetch chunk requests to fetch. | 0.2.0 |  | 
 | celeborn.storage.availableTypes | HDD | Enabled storages. Available options: 
MEMORY,HDD,SSD,HDFS. Note: HDD and SSD would be treated as identical. | 0.3.0 | 
celeborn.storage.activeTypes | 
 | celeborn.storage.hdfs.dir | &lt;undefined&gt; | HDFS base directory for 
Celeborn to store shuffle data. | 0.2.0 |  | 
diff --git 
a/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala 
b/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
index b6441d8e0..b4bcfc49b 100644
--- 
a/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
+++ 
b/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
@@ -38,11 +38,13 @@ import org.apache.celeborn.common.internal.Logging
 import org.apache.celeborn.common.meta.{DiskInfo, WorkerInfo, WorkerStatus}
 import org.apache.celeborn.common.metrics.MetricsSystem
 import org.apache.celeborn.common.metrics.source.{JVMCPUSource, JVMSource, 
ResourceConsumptionSource, SystemMiscSource, ThreadPoolSource}
+import org.apache.celeborn.common.network.sasl.SecretRegistryImpl
 import org.apache.celeborn.common.protocol._
 import org.apache.celeborn.common.protocol.message.{ControlMessages, 
StatusCode}
 import org.apache.celeborn.common.protocol.message.ControlMessages._
 import org.apache.celeborn.common.quota.{QuotaManager, ResourceConsumption}
 import org.apache.celeborn.common.rpc._
+import org.apache.celeborn.common.security.{RpcSecurityContextBuilder, 
ServerSaslContextBuilder}
 import org.apache.celeborn.common.util.{CelebornHadoopUtils, CollectionUtils, 
JavaUtils, PbSerDeUtils, ThreadUtils, Utils}
 import org.apache.celeborn.server.common.{HttpService, Service}
 import 
org.apache.celeborn.service.deploy.master.clustermeta.SingleMasterMetaManager
@@ -97,6 +99,29 @@ private[celeborn] class Master(
   }
 
   private val rackResolver = new CelebornRackResolver(conf)
+  private val authEnabled = conf.authEnabled
+  private val secretRegistry = new SecretRegistryImpl()
+  // Visible for testing
+  private[master] var securedRpcEnv: RpcEnv = _
+  if (authEnabled) {
+    val externalSecurityContext = new RpcSecurityContextBuilder()
+      .withServerSaslContext(
+        new ServerSaslContextBuilder()
+          .withAddRegistrationBootstrap(true)
+          .withSecretRegistry(secretRegistry).build()).build()
+
+    securedRpcEnv = RpcEnv.create(
+      RpcNameConstants.MASTER_SECURED_SYS,
+      masterArgs.host,
+      masterArgs.host,
+      masterArgs.securedPort,
+      conf,
+      Math.max(64, Runtime.getRuntime.availableProcessors()),
+      Some(externalSecurityContext))
+    logInfo(
+      s"Secure port enabled ${masterArgs.securedPort} for secured RPC.")
+  }
+
   private val statusSystem =
     if (conf.haEnabled) {
       val sys = new HAMasterMetaManager(internalRpcEnvInUse, conf, 
rackResolver)
@@ -228,6 +253,16 @@ private[celeborn] class Master(
       internalRpcEndpoint)
   }
 
+  // Visible for testing
+  private[master] var securedRpcEndpoint: RpcEndpoint = _
+  private var securedRpcEndpointRef: RpcEndpointRef = _
+  if (authEnabled) {
+    securedRpcEndpoint = new SecuredRpcEndpoint(this, securedRpcEnv, conf)
+    securedRpcEndpointRef = securedRpcEnv.setupEndpoint(
+      RpcNameConstants.MASTER_SECURED_EP,
+      securedRpcEndpoint)
+  }
+
   // start threads to check timeout for workers and applications
   override def onStart(): Unit = {
     if (!threadsStarted.compareAndSet(false, true)) {
@@ -351,16 +386,22 @@ private[celeborn] class Master(
           requestId,
           shouldResponse) =>
       logDebug(s"Received heartbeat from app $appId")
-      executeWithLeaderChecker(
-        context,
-        handleHeartbeatFromApplication(
+      if (checkAuthStatus(appId, context)) {
+        // TODO: [CELEBORN-1261] For the workers to be able to check whether 
an auth-enabled app is talking to it on
+        // unsecured port, Master will need to maintain a list of 
unauthenticated apps and send it to workers.
+        // This wasn't part of the original proposal because that proposal 
didn't target the Celeborn server to support
+        // both secured and unsecured communication.
+        executeWithLeaderChecker(
           context,
-          appId,
-          totalWritten,
-          fileCount,
-          needCheckedWorkerList,
-          requestId,
-          shouldResponse))
+          handleHeartbeatFromApplication(
+            context,
+            appId,
+            totalWritten,
+            fileCount,
+            needCheckedWorkerList,
+            requestId,
+            shouldResponse))
+      }
 
     case pbRegisterWorker: PbRegisterWorker =>
       val requestId = pbRegisterWorker.getRequestId
@@ -394,23 +435,32 @@ private[celeborn] class Master(
       // keep it for compatible reason
       context.reply(ReleaseSlotsResponse(StatusCode.SUCCESS))
 
-    case requestSlots @ RequestSlots(_, _, _, _, _, _, _, _, _, _, _) =>
-      logTrace(s"Received RequestSlots request $requestSlots.")
-      executeWithLeaderChecker(context, handleRequestSlots(context, 
requestSlots))
+    case requestSlots @ RequestSlots(applicationId, _, _, _, _, _, _, _, _, _, 
_) =>
+      if (checkAuthStatus(applicationId, context)) {
+        // TODO: [CELEBORN-1261]
+        logTrace(s"Received RequestSlots request $requestSlots.")
+        executeWithLeaderChecker(context, handleRequestSlots(context, 
requestSlots))
+      }
 
     case pb: PbUnregisterShuffle =>
       val applicationId = pb.getAppId
       val shuffleId = pb.getShuffleId
       val requestId = pb.getRequestId
-      logDebug(s"Received UnregisterShuffle request $requestId, 
$applicationId, $shuffleId")
-      executeWithLeaderChecker(
-        context,
-        handleUnregisterShuffle(context, applicationId, shuffleId, requestId))
+      if (checkAuthStatus(applicationId, context)) {
+        // TODO: [CELEBORN-1261]
+        logDebug(s"Received UnregisterShuffle request $requestId, 
$applicationId, $shuffleId")
+        executeWithLeaderChecker(
+          context,
+          handleUnregisterShuffle(context, applicationId, shuffleId, 
requestId))
+      }
 
     case ApplicationLost(appId, requestId) =>
-      logDebug(
-        s"Received ApplicationLost request $requestId, $appId from 
${context.senderAddress}.")
-      executeWithLeaderChecker(context, handleApplicationLost(context, appId, 
requestId))
+      if (context.senderAddress.equals(self.address) || checkAuthStatus(appId, 
context)) {
+        // TODO: [CELEBORN-1261]
+        logDebug(
+          s"Received ApplicationLost request $requestId, $appId from 
${context.senderAddress}.")
+        executeWithLeaderChecker(context, handleApplicationLost(context, 
appId, requestId))
+      }
 
     case HeartbeatFromWorker(
           host,
@@ -475,6 +525,8 @@ private[celeborn] class Master(
         handleWorkerLost(context, host, rpcPort, pushPort, fetchPort, 
replicatePort, requestId))
 
     case CheckQuota(userIdentifier) =>
+      // TODO: CheckQuota doesn't have application id in it, so we can't check 
auth status here.
+      // Will have to add application id to CheckQuota message to check auth 
status.
       executeWithLeaderChecker(context, handleCheckQuota(userIdentifier, 
context))
 
     case _: PbCheckWorkersAvailable =>
@@ -907,7 +959,7 @@ private[celeborn] class Master(
     }
   }
 
-  private def handleHeartbeatFromApplication(
+  private[master] def handleHeartbeatFromApplication(
       context: RpcCallContext,
       appId: String,
       totalWritten: Long,
@@ -1007,7 +1059,7 @@ private[celeborn] class Master(
     resourceConsumption
   }
 
-  private def handleCheckQuota(
+  private[master] def handleCheckQuota(
       userIdentifier: UserIdentifier,
       context: RpcCallContext): Unit = {
     val userResourceConsumption = handleResourceConsumption(userIdentifier)
@@ -1037,6 +1089,16 @@ private[celeborn] class Master(
     }.asJava
   }
 
+  private def checkAuthStatus(appId: String, context: RpcCallContext): Boolean 
= {
+    if (conf.authEnabled && secretRegistry.isRegistered(appId)) {
+      context.sendFailure(new SecurityException(
+        s"Auth enabled application $appId sending messages on unsecured 
port!"))
+      false
+    } else {
+      true
+    }
+  }
+
   override def getMasterGroupInfo: String = {
     val sb = new StringBuilder
     sb.append("====================== Master Group INFO 
==============================\n")
@@ -1242,6 +1304,9 @@ private[celeborn] class Master(
     if (conf.internalPortEnabled) {
       internalRpcEnvInUse.awaitTermination()
     }
+    if (authEnabled) {
+      securedRpcEnv.awaitTermination()
+    }
   }
 
   override def stop(exitKind: Int): Unit = synchronized {
@@ -1251,6 +1316,9 @@ private[celeborn] class Master(
       if (conf.internalPortEnabled) {
         internalRpcEnvInUse.stop(internalRpcEndpointRef)
       }
+      if (authEnabled) {
+        securedRpcEnv.stop(securedRpcEndpointRef)
+      }
       super.stop(exitKind)
       logInfo("Master stopped.")
       stopped = true
diff --git 
a/master/src/main/scala/org/apache/celeborn/service/deploy/master/MasterArguments.scala
 
b/master/src/main/scala/org/apache/celeborn/service/deploy/master/MasterArguments.scala
index 701116a2e..95dda30b4 100644
--- 
a/master/src/main/scala/org/apache/celeborn/service/deploy/master/MasterArguments.scala
+++ 
b/master/src/main/scala/org/apache/celeborn/service/deploy/master/MasterArguments.scala
@@ -28,6 +28,7 @@ class MasterArguments(args: Array[String], conf: 
CelebornConf) {
   private var _host: Option[String] = None
   private var _port: Option[Int] = None
   private var _internalPort: Option[Int] = None
+  private var _securedPort: Option[Int] = None
   private var _propertiesFile: Option[String] = None
   private var _masterClusterInfo: Option[MasterClusterInfo] = None
 
@@ -47,11 +48,15 @@ class MasterArguments(args: Array[String], conf: 
CelebornConf) {
     _internalPort = _internalPort.orElse {
       if (conf.internalPortEnabled) 
Some(conf.haMasterNodeInternalPort(localNode.nodeId)) else None
     }
+    _securedPort = _securedPort.orElse {
+      if (conf.authEnabled) 
Some(conf.haMasterNodeSecuredPort(localNode.nodeId)) else None
+    }
     _masterClusterInfo = Some(clusterInfo)
   } else {
     _host = _host.orElse(Some(conf.masterHost))
     _port = _port.orElse(Some(conf.masterPort))
     _internalPort = _internalPort.orElse(Some(conf.masterInternalPort))
+    _securedPort = _securedPort.orElse(Some(conf.masterSecuredPort))
   }
 
   def host: String = _host.get
@@ -60,6 +65,8 @@ class MasterArguments(args: Array[String], conf: 
CelebornConf) {
 
   def internalPort: Int = _internalPort.get
 
+  def securedPort: Int = _securedPort.get
+
   def masterClusterInfo: Option[MasterClusterInfo] = _masterClusterInfo
 
   @tailrec
@@ -77,6 +84,10 @@ class MasterArguments(args: Array[String], conf: 
CelebornConf) {
       _internalPort = Some(value)
       parse(tail)
 
+    case ("--secured-port") :: IntParam(value) :: tail =>
+      _securedPort = Some(value)
+      parse(tail)
+
     case "--properties-file" :: value :: tail =>
       _propertiesFile = Some(value)
       parse(tail)
@@ -102,6 +113,7 @@ class MasterArguments(args: Array[String], conf: 
CelebornConf) {
         |  -h HOST, --host HOST   Hostname to listen on
         |  -p PORT, --port PORT   Port to listen on (default: 9097)
         |  --internal-port PORT   Internal port for internal communication 
(default: 8097)
+        |  --secured-port  PORT   Secured port for secured communication 
(default: 19097)
         |  --properties-file FILE Path to a custom Celeborn properties file,
         |                         default is conf/celeborn-defaults.conf.
         |""".stripMargin)
diff --git 
a/master/src/main/scala/org/apache/celeborn/service/deploy/master/SecuredRpcEndpoint.scala
 
b/master/src/main/scala/org/apache/celeborn/service/deploy/master/SecuredRpcEndpoint.scala
new file mode 100644
index 000000000..e8b1eb966
--- /dev/null
+++ 
b/master/src/main/scala/org/apache/celeborn/service/deploy/master/SecuredRpcEndpoint.scala
@@ -0,0 +1,91 @@
+/*
+ * 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.service.deploy.master
+
+import org.apache.celeborn.common.CelebornConf
+import org.apache.celeborn.common.internal.Logging
+import org.apache.celeborn.common.protocol.PbUnregisterShuffle
+import 
org.apache.celeborn.common.protocol.message.ControlMessages.{ApplicationLost, 
CheckQuota, HeartbeatFromApplication, RequestSlots}
+import org.apache.celeborn.common.rpc._
+
+/**
+ * Secured RPC endpoint used by the Master to communicate with the Clients.
+ */
+private[celeborn] class SecuredRpcEndpoint(
+    val master: Master,
+    override val rpcEnv: RpcEnv,
+    val conf: CelebornConf)
+  extends RpcEndpoint with Logging {
+
+  // start threads to check timeout for workers and applications
+  override def onStart(): Unit = {
+    master.onStart()
+  }
+
+  override def onStop(): Unit = {
+    master.onStop()
+  }
+
+  override def onDisconnected(address: RpcAddress): Unit = {
+    logDebug(s"Client $address got disconnected.")
+  }
+
+  override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, 
Unit] = {
+    case HeartbeatFromApplication(
+          appId,
+          totalWritten,
+          fileCount,
+          needCheckedWorkerList,
+          requestId,
+          shouldResponse) =>
+      logDebug(s"Received heartbeat from app $appId")
+      master.executeWithLeaderChecker(
+        context,
+        master.handleHeartbeatFromApplication(
+          context,
+          appId,
+          totalWritten,
+          fileCount,
+          needCheckedWorkerList,
+          requestId,
+          shouldResponse))
+
+    case requestSlots @ RequestSlots(_, _, _, _, _, _, _, _, _, _, _) =>
+      logTrace(s"Received RequestSlots request $requestSlots.")
+      master.executeWithLeaderChecker(context, 
master.handleRequestSlots(context, requestSlots))
+
+    case pb: PbUnregisterShuffle =>
+      val applicationId = pb.getAppId
+      val shuffleId = pb.getShuffleId
+      val requestId = pb.getRequestId
+      logDebug(s"Received UnregisterShuffle request $requestId, 
$applicationId, $shuffleId")
+      master.executeWithLeaderChecker(
+        context,
+        master.handleUnregisterShuffle(context, applicationId, shuffleId, 
requestId))
+
+    case ApplicationLost(appId, requestId) =>
+      logDebug(
+        s"Received ApplicationLost request $requestId, $appId from 
${context.senderAddress}.")
+      master.executeWithLeaderChecker(
+        context,
+        master.handleApplicationLost(context, appId, requestId))
+
+    case CheckQuota(userIdentifier) =>
+      master.executeWithLeaderChecker(context, 
master.handleCheckQuota(userIdentifier, context))
+  }
+}
diff --git 
a/master/src/test/scala/org/apache/celeborn/service/deploy/master/MasterSuite.scala
 
b/master/src/test/scala/org/apache/celeborn/service/deploy/master/MasterSuite.scala
index 2d3e66912..62e733691 100644
--- 
a/master/src/test/scala/org/apache/celeborn/service/deploy/master/MasterSuite.scala
+++ 
b/master/src/test/scala/org/apache/celeborn/service/deploy/master/MasterSuite.scala
@@ -18,13 +18,14 @@
 package org.apache.celeborn.service.deploy.master
 
 import com.google.common.io.Files
-import org.mockito.Mockito.mock
+import org.mockito.Mockito.{mock, times, verify}
 import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}
 import org.scalatest.funsuite.AnyFunSuite
 
 import org.apache.celeborn.common.CelebornConf
 import org.apache.celeborn.common.internal.Logging
 import org.apache.celeborn.common.protocol.{PbCheckForWorkerTimeout, 
PbRegisterWorker}
+import 
org.apache.celeborn.common.protocol.message.ControlMessages.{ApplicationLost, 
HeartbeatFromApplication}
 import org.apache.celeborn.common.util.{CelebornExitKind, Utils}
 
 class MasterSuite extends AnyFunSuite
@@ -96,4 +97,44 @@ class MasterSuite extends AnyFunSuite
     master.rpcEnv.shutdown()
     master.internalRpcEnvInUse.shutdown()
   }
+
+  test("test secured port receives") {
+    val conf = new CelebornConf()
+    conf.set(CelebornConf.HA_ENABLED.key, "false")
+    conf.set(CelebornConf.HA_MASTER_RATIS_STORAGE_DIR.key, getTmpDir())
+    conf.set(CelebornConf.WORKER_STORAGE_DIRS.key, getTmpDir())
+    conf.set(CelebornConf.METRICS_ENABLED.key, "true")
+    conf.set(CelebornConf.INTERNAL_PORT_ENABLED.key, "true")
+    conf.set(CelebornConf.AUTH_ENABLED.key, "true")
+
+    val args =
+      Array("-h", "localhost", "-p", "9097", "--internal-port", "8097", 
"--secured-port", "19097")
+
+    val masterArgs = new MasterArguments(args, conf)
+    val master = new Master(conf, masterArgs)
+    new Thread() {
+      override def run(): Unit = {
+        master.initialize()
+      }
+    }.start()
+    Thread.sleep(5000L)
+    master.securedRpcEndpoint.receiveAndReply(
+      mock(classOf[org.apache.celeborn.common.rpc.RpcCallContext]))
+      .applyOrElse(
+        HeartbeatFromApplication("appId", 0L, 0L, null),
+        (_: Any) => fail("Unexpected message"))
+    master.securedRpcEndpoint.receiveAndReply(
+      mock(classOf[org.apache.celeborn.common.rpc.RpcCallContext]))
+      .applyOrElse(ApplicationLost("appId"), (_: Any) => fail("Unexpected 
message"))
+
+    assertThrows[scala.MatchError] {
+      master.securedRpcEndpoint.receiveAndReply(
+        mock(classOf[org.apache.celeborn.common.rpc.RpcCallContext]))(
+        PbRegisterWorker.newBuilder().build())
+    }
+    master.stop(CelebornExitKind.EXIT_IMMEDIATELY)
+    master.rpcEnv.shutdown()
+    master.internalRpcEnvInUse.shutdown()
+    master.securedRpcEnv.shutdown()
+  }
 }
diff --git 
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala 
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
index 46ff96f4f..30993f3c3 100644
--- 
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
+++ 
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
@@ -237,7 +237,8 @@ private[celeborn] class Worker(
   val shuffleCommitInfos: ConcurrentHashMap[String, ConcurrentHashMap[Long, 
CommitInfo]] =
     JavaUtils.newConcurrentHashMap[String, ConcurrentHashMap[Long, 
CommitInfo]]()
 
-  private val masterClient = new MasterClient(rpcEnv, conf)
+  // TODO: pass the internal rpc env here when internal port is added to the 
worker.
+  private val masterClient = new MasterClient(rpcEnv, conf, true)
 
   // (workerInfo -> last connect timeout timestamp)
   val unavailablePeers: ConcurrentHashMap[WorkerInfo, Long] =


Reply via email to