(spark) branch branch-4.x updated: [MINOR] Improve Transport Context Stream Manager Delegation

[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.x
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.x by this push:
     new 0356335660cf [MINOR] Improve Transport Context Stream Manager Delegation
0356335660cf is described below

commit 0356335660cfb43c63e32f426772552cdb5678a0
Author: Holden Karau <ho...@snowflake.com>
AuthorDate: Sat Aug 29 00:44:53 2026 +0000

    [MINOR] Improve Transport Context Stream Manager Delegation
    
    Co-authored-by: Holden Karau <ho...@pigscanfly.ca>
    (cherry picked from commit 3b47eee2570ac130ff35c09ec205304df691f237)
---
 .../org/apache/spark/network/TransportContext.java |  5 +-
 .../network/server/AbstractAuthRpcHandler.java     | 83 +++++++++++++++++-
 .../server/AbstractAuthRpcHandlerSuite.java        | 98 ++++++++++++++++++++++
 3 files changed, 184 insertions(+), 2 deletions(-)

diff --git a/common/network-common/src/main/java/org/apache/spark/network/TransportContext.java b/common/network-common/src/main/java/org/apache/spark/network/TransportContext.java
index e8ce6840e3fc..2b99e76b49a1 100644
--- a/common/network-common/src/main/java/org/apache/spark/network/TransportContext.java
+++ b/common/network-common/src/main/java/org/apache/spark/network/TransportContext.java
@@ -237,8 +237,11 @@ public class TransportContext implements Closeable {
         .addLast("handler", channelHandler);
       // Use a separate EventLoopGroup to handle ChunkFetchRequest messages for shuffle rpcs.
       if (chunkFetchWorkers != null) {
+        // Use the per-channel handler (which may be wrapped by an authentication bootstrap),
+        // not the context-level one, so chunk fetches also fail closed until the channel has
+        // authenticated.
         ChunkFetchRequestHandler chunkFetchHandler = new ChunkFetchRequestHandler(
-          channelHandler.getClient(), rpcHandler.getStreamManager(),
+          channelHandler.getClient(), channelRpcHandler.getStreamManager(),
           conf.maxChunksBeingTransferred(), true /* syncModeEnabled */);
         pipeline.addLast(chunkFetchWorkers, "chunkFetchHandler", chunkFetchHandler);
       }
diff --git a/common/network-common/src/main/java/org/apache/spark/network/server/AbstractAuthRpcHandler.java b/common/network-common/src/main/java/org/apache/spark/network/server/AbstractAuthRpcHandler.java
index 95fde677624f..bb806a9c5371 100644
--- a/common/network-common/src/main/java/org/apache/spark/network/server/AbstractAuthRpcHandler.java
+++ b/common/network-common/src/main/java/org/apache/spark/network/server/AbstractAuthRpcHandler.java
@@ -19,6 +19,9 @@ package org.apache.spark.network.server;
 
 import java.nio.ByteBuffer;
 
+import io.netty.channel.Channel;
+
+import org.apache.spark.network.buffer.ManagedBuffer;
 import org.apache.spark.network.client.RpcResponseCallback;
 import org.apache.spark.network.client.StreamCallbackWithID;
 import org.apache.spark.network.client.TransportClient;
@@ -83,7 +86,7 @@ public abstract class AbstractAuthRpcHandler extends RpcHandler {
 
   @Override
   public StreamManager getStreamManager() {
-    return delegate.getStreamManager();
+    return new AuthCheckingStreamManager(delegate.getStreamManager());
   }
 
   @Override
@@ -109,4 +112,82 @@ public abstract class AbstractAuthRpcHandler extends RpcHandler {
   public MergedBlockMetaReqHandler getMergedBlockMetaReqHandler() {
     return delegate.getMergedBlockMetaReqHandler();
   }
+
+  /**
+   * Wraps the delegate's StreamManager so that no chunk or stream is served on a channel that
+   * has not completed authentication. Historically only receive() and receiveStream() were
+   * gated on authentication, which left StreamRequest and ChunkFetchRequest served pre-auth by
+   * StreamManagers whose checkAuthorization is a no-op (e.g. the NettyRpcEnv file server that
+   * distributes jars, files and REPL classes), so enabling spark.authenticate did not protect
+   * the file-distribution channel. This wrapper makes every StreamManager behind an
+   * authentication bootstrap fail closed instead. Lifecycle and accounting callbacks are
+   * always delegated so per-channel state is cleaned up regardless of authentication state.
+   */
+  private class AuthCheckingStreamManager extends StreamManager {
+    private final StreamManager delegate;
+
+    AuthCheckingStreamManager(StreamManager delegate) {
+      this.delegate = delegate;
+    }
+
+    private void checkAuthenticated() {
+      if (!isAuthenticated) {
+        throw new SecurityException("Unauthenticated call to stream manager.");
+      }
+    }
+
+    @Override
+    public ManagedBuffer getChunk(long streamId, int chunkIndex) {
+      checkAuthenticated();
+      return delegate.getChunk(streamId, chunkIndex);
+    }
+
+    @Override
+    public ManagedBuffer openStream(String streamId) {
+      checkAuthenticated();
+      return delegate.openStream(streamId);
+    }
+
+    @Override
+    public void checkAuthorization(TransportClient client, long streamId) {
+      checkAuthenticated();
+      delegate.checkAuthorization(client, streamId);
+    }
+
+    @Override
+    public void checkAuthorization(TransportClient client, String streamId) {
+      checkAuthenticated();
+      delegate.checkAuthorization(client, streamId);
+    }
+
+    @Override
+    public void connectionTerminated(Channel channel) {
+      delegate.connectionTerminated(channel);
+    }
+
+    @Override
+    public long chunksBeingTransferred() {
+      return delegate.chunksBeingTransferred();
+    }
+
+    @Override
+    public void chunkBeingSent(long streamId) {
+      delegate.chunkBeingSent(streamId);
+    }
+
+    @Override
+    public void streamBeingSent(String streamId) {
+      delegate.streamBeingSent(streamId);
+    }
+
+    @Override
+    public void chunkSent(long streamId) {
+      delegate.chunkSent(streamId);
+    }
+
+    @Override
+    public void streamSent(String streamId) {
+      delegate.streamSent(streamId);
+    }
+  }
 }
diff --git a/common/network-common/src/test/java/org/apache/spark/network/server/AbstractAuthRpcHandlerSuite.java b/common/network-common/src/test/java/org/apache/spark/network/server/AbstractAuthRpcHandlerSuite.java
new file mode 100644
index 000000000000..385c584b3e03
--- /dev/null
+++ b/common/network-common/src/test/java/org/apache/spark/network/server/AbstractAuthRpcHandlerSuite.java
@@ -0,0 +1,98 @@
+/*
+ * 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.spark.network.server;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
+import java.nio.ByteBuffer;
+
+import io.netty.channel.Channel;
+import org.junit.jupiter.api.Test;
+
+import org.apache.spark.network.client.RpcResponseCallback;
+import org.apache.spark.network.client.TransportClient;
+
+/**
+ * Tests that RPC handlers behind an authentication bootstrap fail closed: no chunk or stream
+ * request may be served on a channel that has not completed authentication, regardless of how
+ * lax the delegate StreamManager's own authorization checks are (the default
+ * StreamManager.checkAuthorization is a no-op).
+ */
+public class AbstractAuthRpcHandlerSuite {
+
+  private static class TestAuthRpcHandler extends AbstractAuthRpcHandler {
+    TestAuthRpcHandler(RpcHandler delegate) {
+      super(delegate);
+    }
+
+    @Override
+    protected boolean doAuthChallenge(
+        TransportClient client,
+        ByteBuffer message,
+        RpcResponseCallback callback) {
+      return true;
+    }
+  }
+
+  @Test
+  public void testStreamManagerFailsClosedBeforeAuth() {
+    RpcHandler delegate = mock(RpcHandler.class);
+    StreamManager delegateManager = mock(StreamManager.class);
+    when(delegate.getStreamManager()).thenReturn(delegateManager);
+    AbstractAuthRpcHandler handler = new TestAuthRpcHandler(delegate);
+
+    StreamManager sm = handler.getStreamManager();
+    TransportClient client = mock(TransportClient.class);
+
+    assertThrows(SecurityException.class, () -> sm.openStream("/jars/app.jar"));
+    assertThrows(SecurityException.class, () -> sm.getChunk(0L, 0));
+    assertThrows(SecurityException.class, () -> sm.checkAuthorization(client, 0L));
+    assertThrows(SecurityException.class, () -> sm.checkAuthorization(client, "/jars/app.jar"));
+    verify(delegateManager, never()).openStream(anyString());
+    verify(delegateManager, never()).getChunk(anyLong(), anyInt());
+
+    // Lifecycle callbacks still reach the delegate so per-channel state can be cleaned up
+    // even for channels that never authenticated.
+    Channel channel = mock(Channel.class);
+    sm.connectionTerminated(channel);
+    verify(delegateManager).connectionTerminated(channel);
+  }
+
+  @Test
+  public void testStreamManagerDelegatesAfterAuth() {
+    RpcHandler delegate = mock(RpcHandler.class);
+    StreamManager delegateManager = mock(StreamManager.class);
+    when(delegate.getStreamManager()).thenReturn(delegateManager);
+    AbstractAuthRpcHandler handler = new TestAuthRpcHandler(delegate);
+    StreamManager sm = handler.getStreamManager();
+
+    // Complete the (test) auth handshake; the wrapper obtained pre-auth must observe it.
+    handler.receive(mock(TransportClient.class), ByteBuffer.allocate(0),
+      mock(RpcResponseCallback.class));
+    assertTrue(handler.isAuthenticated());
+
+    sm.openStream("/jars/app.jar");
+    verify(delegateManager).openStream("/jars/app.jar");
+    sm.getChunk(1L, 2);
+    verify(delegateManager).getChunk(1L, 2);
+    TransportClient client = mock(TransportClient.class);
+    sm.checkAuthorization(client, 1L);
+    verify(delegateManager).checkAuthorization(client, 1L);
+  }
+}


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