(celeborn) branch main updated: [CELEBORN-2382] Authorize `PbReviseLostShuffles` to prevent cross-application metadata modification

[ni...@apache.org]

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

SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git


The following commit(s) were added to refs/heads/main by this push:
     new d77543ef4b [CELEBORN-2382] Authorize `PbReviseLostShuffles` to prevent cross-application metadata modification
d77543ef4b is described below

commit d77543ef4bbd3a62aa8f247b6434e7c438d3850a
Author: Kalvin2077 <wk...@outlook.com>
AuthorDate: Wed Jul 15 17:56:32 2026 +0800

    [CELEBORN-2382] Authorize `PbReviseLostShuffles` to prevent cross-application metadata modification
    
    ### What changes were proposed in this pull request?
    
    Call `checkAuth(context, pb.getAppId)` before serving `PbReviseLostShuffles` in the Master, consistent with the other application-scoped handlers.
    
    ### Why are the changes needed?
    
    The Master handled `PbReviseLostShuffles` without an authorization check. With authentication enabled, an application could specify another application's ID and add shuffle IDs to its `registeredAppAndShuffles` metadata.
    
    ### Does this PR resolve a correctness bug?
    
    - [ ] Yes
    
    ### Does this PR introduce _any_ user-facing change?
    
    - [ ] Yes
    
    ### How was this patch tested?
    
    - New case in `MasterApplicationMetaAuthSuite` verifying same-application requests succeed and cross-application requests are rejected with authentication enabled
    - New case in `MasterSuite` verifying same-application requests succeed, cross-application requests are rejected, and authentication-disabled behavior is preserved.
    
    Closes #3761 from Kalvin2077/fix/pbrevise-checkauth.
    
    Authored-by: Kalvin2077 <wk...@outlook.com>
    Signed-off-by: Nicholas Jiang <pr...@163.com>
---
 .../celeborn/service/deploy/master/Master.scala    |  1 +
 .../master/MasterApplicationMetaAuthSuite.scala    | 35 +++++++++-
 .../service/deploy/master/MasterSuite.scala        | 79 ++++++++++++++++++----
 3 files changed, 99 insertions(+), 16 deletions(-)

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 8960934dd3..4203041b74 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
@@ -599,6 +599,7 @@ private[celeborn] class Master(
         handleWorkerDecommission(context, workers, requestId))
 
     case pb: PbReviseLostShuffles =>
+      checkAuth(context, pb.getAppId)
       executeWithLeaderChecker(
         context,
         handleReviseLostShuffle(context, pb.getAppId, pb.getLostShufflesList, pb.getRequestId))
diff --git a/master/src/test/scala/org/apache/celeborn/service/deploy/master/MasterApplicationMetaAuthSuite.scala b/master/src/test/scala/org/apache/celeborn/service/deploy/master/MasterApplicationMetaAuthSuite.scala
index b672851cb0..50662f0383 100644
--- a/master/src/test/scala/org/apache/celeborn/service/deploy/master/MasterApplicationMetaAuthSuite.scala
+++ b/master/src/test/scala/org/apache/celeborn/service/deploy/master/MasterApplicationMetaAuthSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.celeborn.service.deploy.master
 
 import java.io.{PrintWriter, StringWriter}
+import java.util
 
 import scala.collection.mutable.ArrayBuffer
 
@@ -27,12 +28,13 @@ import org.scalatest.funsuite.AnyFunSuite
 import org.apache.celeborn.common.CelebornConf
 import org.apache.celeborn.common.metrics.source.Role
 import org.apache.celeborn.common.network.sasl.registration.RegistrationInfo
-import org.apache.celeborn.common.protocol.{PbApplicationMeta, PbApplicationMetaRequest, RpcNameConstants, TransportModuleConstants}
+import org.apache.celeborn.common.protocol.{PbApplicationMeta, PbApplicationMetaRequest, PbReviseLostShufflesResponse, RpcNameConstants, TransportModuleConstants}
+import org.apache.celeborn.common.protocol.message.ControlMessages.ReviseLostShuffles
 import org.apache.celeborn.common.rpc.{ClientSaslContextBuilder, RpcAddress, RpcEndpointRef, RpcEnv, RpcSecurityContextBuilder}
 
 /**
- * End-to-end authorization check for PbApplicationMetaRequest with auth enabled, driving
- * the real SASL registration path rather than a mocked client id (as in [[MasterSuite]]).
+ * End-to-end authorization checks for application-scoped Master RPCs with auth enabled,
+ * driving the real SASL registration path rather than a mocked client id (as in [[MasterSuite]]).
  *
  * It covers the whole chain the security guarantee rests on: registration sets the
  * connection's client id, and checkAuth enforces it. A regression that stopped setting
@@ -143,4 +145,31 @@ class MasterApplicationMetaAuthSuite extends AnyFunSuite
       workerInternalRef().askSync[PbApplicationMeta](metaRequest(victimApp)).getSecret
         == victimSecret)
   }
+
+  test("PbReviseLostShuffles is authorized against the registered application") {
+    val victimApp = "revise-victim-app"
+    val attackerApp = "revise-attacker-app"
+
+    val victimRef = registeredAppRef(victimApp, "revise-victim-secret")
+    val ownResponse = victimRef.askSync[PbReviseLostShufflesResponse](
+      ReviseLostShuffles(victimApp, util.Arrays.asList[Integer](1), "victim-request"))
+    assert(ownResponse.getSuccess)
+
+    val victimShuffles = master.statusSystem.registeredAppAndShuffles.get(victimApp)
+    assert(victimShuffles.size() == 1)
+    assert(victimShuffles.contains(1))
+
+    val attackerRef = registeredAppRef(attackerApp, "revise-attacker-secret")
+    val attackerResponse = attackerRef.askSync[PbReviseLostShufflesResponse](
+      ReviseLostShuffles(attackerApp, util.Arrays.asList[Integer](2), "attacker-request"))
+    assert(attackerResponse.getSuccess)
+
+    val e = intercept[Exception] {
+      attackerRef.askSync[PbReviseLostShufflesResponse](
+        ReviseLostShuffles(victimApp, util.Arrays.asList[Integer](3), "cross-app-request"))
+    }
+    assert(stackTraceOf(e).contains(s"not authorized for application $victimApp"))
+    assert(victimShuffles.size() == 1)
+    assert(!victimShuffles.contains(3))
+  }
 }
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 d915320c95..48ad9fe8d9 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
@@ -29,7 +29,7 @@ import org.apache.celeborn.common.CelebornConf
 import org.apache.celeborn.common.identity.UserIdentifier
 import org.apache.celeborn.common.network.client.{RpcResponseCallback, TransportClient}
 import org.apache.celeborn.common.protocol.{PbApplicationMetaRequest, PbCheckForWorkerTimeout, PbRegisterWorker}
-import org.apache.celeborn.common.protocol.message.ControlMessages.{RequestSlots, RequestSlotsResponse}
+import org.apache.celeborn.common.protocol.message.ControlMessages.{RequestSlots, RequestSlotsResponse, ReviseLostShuffles}
 import org.apache.celeborn.common.protocol.message.StatusCode
 import org.apache.celeborn.common.rpc.{RpcAddress, RpcCallContext}
 import org.apache.celeborn.common.rpc.netty.{NettyRpcEnv, RemoteNettyRpcCallContext}
@@ -40,6 +40,18 @@ class MasterSuite extends AnyFunSuite
   with BeforeAndAfterEach
   with MasterClusterFeature {
 
+  // Builds a remote call context whose connection is authenticated as `clientId`;
+  // null models a connection when authentication is disabled or an internal worker connection.
+  private def contextForClient(clientId: String): RemoteNettyRpcCallContext = {
+    val client = mock(classOf[TransportClient])
+    when(client.getClientId).thenReturn(clientId)
+    new RemoteNettyRpcCallContext(
+      mock(classOf[NettyRpcEnv]),
+      mock(classOf[RpcResponseCallback]),
+      RpcAddress("localhost", 1234),
+      client)
+  }
+
   def getTmpDir(): String = {
     val tmpDir = Files.createTempDirectory(null).toFile
     tmpDir.deleteOnExit()
@@ -212,18 +224,6 @@ class MasterSuite extends AnyFunSuite
     val masterArgs = new MasterArguments(args, conf)
     val master = new Master(conf, masterArgs)
 
-    // Builds a remote call context whose connection is authenticated as `clientId`;
-    // null models a worker on the internal channel, which sets no client id.
-    def contextForClient(clientId: String): RemoteNettyRpcCallContext = {
-      val client = mock(classOf[TransportClient])
-      when(client.getClientId).thenReturn(clientId)
-      new RemoteNettyRpcCallContext(
-        mock(classOf[NettyRpcEnv]),
-        mock(classOf[RpcResponseCallback]),
-        RpcAddress("localhost", 1234),
-        client)
-    }
-
     val request = PbApplicationMetaRequest.newBuilder().setAppId("victim-app").build()
     val unhandled = (_: Any) => fail("PbApplicationMetaRequest was not handled")
 
@@ -241,4 +241,57 @@ class MasterSuite extends AnyFunSuite
       master.rpcEnv.shutdown()
     }
   }
+
+  test("PbReviseLostShuffles authorizes application metadata changes") {
+    val conf = new CelebornConf()
+    val randomMasterPort = selectRandomPort()
+    val randomHttpPort = selectRandomPort()
+    conf.set(CelebornConf.HA_ENABLED.key, "false")
+    conf.set(CelebornConf.MASTER_HTTP_HOST.key, "127.0.0.1")
+    conf.set(CelebornConf.MASTER_HTTP_PORT.key, randomHttpPort.toString)
+
+    val args = Array("-h", "localhost", "-p", randomMasterPort.toString)
+    val masterArgs = new MasterArguments(args, conf)
+    val master = new Master(conf, masterArgs)
+
+    val victimApp = "revise-victim-app"
+    val ownRequest = ReviseLostShuffles(
+      victimApp,
+      util.Arrays.asList[Integer](1),
+      "own-request")
+    val crossAppRequest = ReviseLostShuffles(
+      victimApp,
+      util.Arrays.asList[Integer](2),
+      "cross-app-request")
+    val authDisabledApp = "auth-disabled-app"
+    val authDisabledRequest = ReviseLostShuffles(
+      authDisabledApp,
+      util.Arrays.asList[Integer](3),
+      "auth-disabled-request")
+    val unhandled = (_: Any) => fail("PbReviseLostShuffles was not handled")
+
+    try {
+      master.receiveAndReply(contextForClient(victimApp)).applyOrElse(ownRequest, unhandled)
+      val victimShuffles = master.statusSystem.registeredAppAndShuffles.get(victimApp)
+      assert(victimShuffles.size() == 1)
+      assert(victimShuffles.contains(1))
+
+      val e = intercept[IllegalStateException] {
+        master.receiveAndReply(contextForClient("attacker-app"))
+          .applyOrElse(crossAppRequest, unhandled)
+      }
+      assert(e.getMessage.contains(s"not authorized for application $victimApp"))
+      assert(victimShuffles.size() == 1)
+      assert(!victimShuffles.contains(2))
+
+      // Authentication-disabled connections have no client id, so existing behavior is preserved.
+      master.receiveAndReply(contextForClient(null)).applyOrElse(authDisabledRequest, unhandled)
+      val authDisabledShuffles =
+        master.statusSystem.registeredAppAndShuffles.get(authDisabledApp)
+      assert(authDisabledShuffles.size() == 1)
+      assert(authDisabledShuffles.contains(3))
+    } finally {
+      master.rpcEnv.shutdown()
+    }
+  }
 }