This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch branch-4.0
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.0 by this push:
new 5e5d40f5698c [MINOR][4.0][CORE] Improve Transport Context Stream
Manager Delegation
5e5d40f5698c is described below
commit 5e5d40f5698ccd903bdabf4741b5304044428c2a
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Tue Oct 27 22:16:17 2026 -0700
[MINOR][4.0][CORE] Improve Transport Context Stream Manager Delegation
### What changes were proposed in this pull request?
Backport of 3b47eee2570ac130ff35c09ec205304df691f237, which was missed on
branch-4.0.
### Why are the changes needed?
Missed backport.
### Does this PR introduce _any_ user-facing change?
Same as the original commit.
### How was this patch tested?
Existing tests.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #59850 from viirya/backport-3b47eee-branch-4.0.
Lead-authored-by: Liang-Chi Hsieh <[email protected]>
Co-authored-by: Holden Karau <[email protected]>
Signed-off-by: Liang-Chi Hsieh <[email protected]>
---
.../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 6914a3f03eef..1798814f3a2a 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
@@ -117,4 +120,82 @@ public abstract class AbstractAuthRpcHandler extends
RpcHandler {
}
};
}
+
+ /**
+ * 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: [email protected]
For additional commands, e-mail: [email protected]
