[MINOR] Handle client identity after the handshake is complete

Co-authored with Kimi & Claude

Co-authored-by: Holden Karau <holden@pigscanfly.ca>

Author: sfc-gh-hkarau

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 boolean doAuthChallenge(
       }
 
       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 boolean doAuthChallenge(
     // 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);

common/network-common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java

@@ -315,6 +315,45 @@ public void testDataEncryptionIsActuallyEnabled() throws Exception {
     }
   }
 
+  @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