(spark) branch branch-4.1 updated: [MINOR] Handle client identity after the handshake is complete

[ho...@apache.org]

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

asf-gitbox-commits pushed a commit to branch branch-4.1
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.1 by this push:
     new 7482cdb6c9c7 [MINOR] Handle client identity after the handshake is complete
7482cdb6c9c7 is described below

commit 7482cdb6c9c7d46b797f20e51818cd1d49914d62
Author: Holden Karau <ho...@snowflake.com>
AuthorDate: Sun Aug 30 05:21:46 2026 +0000

    [MINOR] Handle client identity after the handshake is complete
    
    Co-authored with Kimi & Claude
    
    Co-authored-by: Holden Karau <ho...@pigscanfly.ca>
---
 .../apache/spark/network/sasl/SaslRpcHandler.java  |  8 +++--
 .../apache/spark/network/sasl/SparkSaslSuite.java  | 39 ++++++++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/common/network-common/src/main/java/org/apache/spark/network/sasl/SaslRpcHandler.java b/common/network-common/src/main/java/org/apache/spark/network/sasl/SaslRpcHandler.java
index b5fffe583ec6..336800dcdbe3 100644
--- a/common/network-common/src/main/java/org/apache/spark/network/sasl/SaslRpcHandler.java
+++ b/common/network-common/src/main/java/org/apache/spark/network/sasl/SaslRpcHandler.java
@@ -56,6 +56,9 @@ public class SaslRpcHandler extends AbstractAuthRpcHandler {
 
   private SparkSaslServer saslServer;
 
+  /** Claimed AppID from initial connection */
+  private String claimedAppId;
+
   public SaslRpcHandler(
       TransportConf conf,
       Channel channel,
@@ -83,8 +86,7 @@ public class SaslRpcHandler extends AbstractAuthRpcHandler {
       }
 
       if (saslServer == null) {
-        // First message in the handshake, setup the necessary state.
-        client.setClientId(saslMessage.appId);
+        claimedAppId = saslMessage.appId;
         saslServer = new SparkSaslServer(saslMessage.appId, secretKeyHolder,
           conf.saslServerAlwaysEncrypt());
       }
@@ -105,6 +107,8 @@ public class SaslRpcHandler extends AbstractAuthRpcHandler {
     // method returns. This assumes that the code ensures, through other means, that no outbound
     // messages are being written to the channel while negotiation is still going on.
     if (saslServer.isComplete()) {
+      // Set app Id once the handshake is completed.
+      client.setClientId(claimedAppId);
       if (!SparkSaslServer.QOP_AUTH_CONF.equals(saslServer.getNegotiatedProperty(Sasl.QOP))) {
         logger.debug("SASL authentication successful for channel {}", client);
         complete(true);
diff --git a/common/network-common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java b/common/network-common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java
index 4feaf5cef3f2..c87c7e517cb2 100644
--- a/common/network-common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java
+++ b/common/network-common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java
@@ -315,6 +315,45 @@ public class SparkSaslSuite {
     }
   }
 
+  @Test
+  public void testClientIdBoundAfter() throws Exception {
+    TransportConf conf = new TransportConf("shuffle", MapConfigProvider.EMPTY);
+    SecretKeyHolder keyHolder = mock(SecretKeyHolder.class);
+    when(keyHolder.getSaslUser(anyString())).thenReturn("user");
+    when(keyHolder.getSecretKey(anyString())).thenReturn("secret");
+    TransportClient client = mock(TransportClient.class);
+    SaslRpcHandler handler = new SaslRpcHandler(conf, null, null, keyHolder);
+    SparkSaslClient saslClient = new SparkSaslClient("app-1", keyHolder, false);
+    try {
+      byte[] clientToken = saslClient.firstToken();
+      boolean authenticated = false;
+      for (int i = 0; i < 10 && !authenticated; i++) {
+        SaslMessage msg = new SaslMessage("app-1", clientToken);
+        ByteBuf buf = Unpooled.buffer(msg.encodedLength() + clientToken.length);
+        msg.encode(buf);
+        buf.writeBytes(msg.body().nioByteBuffer());
+
+        AtomicReference<ByteBuffer> response = new AtomicReference<>();
+        RpcResponseCallback callback = mock(RpcResponseCallback.class);
+        doAnswer(invocation -> {
+          response.set((ByteBuffer) invocation.getArguments()[0]);
+          return null;
+        }).when(callback).onSuccess(any(ByteBuffer.class));
+
+        authenticated = handler.doAuthChallenge(client, buf.nioBuffer(), callback);
+        if (!authenticated) {
+          // No client ID set since we did not finish the challenge.
+          verify(client, never()).setClientId(anyString());
+          clientToken = saslClient.response(JavaUtils.bufferToArray(response.get()));
+        }
+      }
+      assertTrue(authenticated);
+      verify(client).setClientId("app-1");
+    } finally {
+      saslClient.dispose();
+    }
+  }
+
   @Test
   public void testRpcHandlerDelegate() throws Exception {
     // Tests all delegates exception for receive(), which is more complicated and already handled


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org