[MINOR][NETWORK][3.5] Improve Transport Context Stream Manager Delegation

common/network-common/src/main/java/org/apache/spark/network/TransportContext.java

@@ -204,8 +204,11 @@ public TransportChannelHandler initializePipeline(
         .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);
       }

common/network-common/src/main/java/org/apache/spark/network/crypto/AuthRpcHandler.java

@@ -66,7 +66,7 @@ class AuthRpcHandler extends AbstractAuthRpcHandler {
       Channel channel,
       RpcHandler delegate,
       SecretKeyHolder secretKeyHolder) {
-    super(delegate);
+    super(delegate, conf.requireAuthForStreamRequests());
     this.conf = conf;
     this.channel = channel;
     this.secretKeyHolder = secretKeyHolder;

common/network-common/src/main/java/org/apache/spark/network/sasl/SaslRpcHandler.java

@@ -64,7 +64,7 @@ public SaslRpcHandler(
       Channel channel,
       RpcHandler delegate,
       SecretKeyHolder secretKeyHolder) {
-    super(delegate);
+    super(delegate, conf.requireAuthForStreamRequests());
     this.conf = conf;
     this.channel = channel;
     this.secretKeyHolder = secretKeyHolder;

common/network-common/src/main/java/org/apache/spark/network/server/AbstractAuthRpcHandler.java

@@ -19,9 +19,13 @@
 
 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;
+import org.apache.spark.network.util.TransportConf;
 
 /**
  * RPC Handler which performs authentication, and when it's successful, delegates further
@@ -32,10 +36,23 @@ public abstract class AbstractAuthRpcHandler extends RpcHandler {
   /** RpcHandler we will delegate to for authenticated connections. */
   private final RpcHandler delegate;
 
+  /**
+   * Whether stream and chunk fetch requests fail closed until the channel has authenticated.
+   * See {@link TransportConf#requireAuthForStreamRequests()}.
+   */
+  private final boolean requireAuthForStreamRequests;
+
   private boolean isAuthenticated;
 
+  // Kept for external subclasses on this maintenance branch: defaults to the historical
+  // behavior of serving stream requests before authentication completes.
   protected AbstractAuthRpcHandler(RpcHandler delegate) {
+    this(delegate, false);
+  }
+
+  protected AbstractAuthRpcHandler(RpcHandler delegate, boolean requireAuthForStreamRequests) {
     this.delegate = delegate;
+    this.requireAuthForStreamRequests = requireAuthForStreamRequests;
   }
 
   /**
@@ -83,7 +100,11 @@ public final StreamCallbackWithID receiveStream(
 
   @Override
   public StreamManager getStreamManager() {
-    return delegate.getStreamManager();
+    StreamManager streamManager = delegate.getStreamManager();
+    if (requireAuthForStreamRequests) {
+      return new AuthCheckingStreamManager(streamManager);
+    }
+    return streamManager;
   }
 
   @Override
@@ -117,4 +138,84 @@ public MergedBlockMetaReqHandler 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. It is only installed when
+   * {@link TransportConf#requireAuthForStreamRequests()} is enabled. 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);
+    }
+  }
 }

common/network-common/src/main/java/org/apache/spark/network/util/TransportConf.java

@@ -150,6 +150,15 @@ public int authRTTimeoutMs() {
       conf.get(SPARK_NETWORK_SASL_TIMEOUT_KEY, "30s"))) * 1000;
   }
 
+  /**
+   * Whether stream and chunk fetch requests on a channel behind an authentication bootstrap
+   * are only served once the channel has completed authentication. Opt-in on branch-3.5 so a
+   * patch release does not change runtime behavior for existing deployments.
+   */
+  public boolean requireAuthForStreamRequests() {
+    return conf.getBoolean("spark.network.auth.requireAuthForStreamRequests", false);
+  }
+
   /**
    * Max number of times we will try IO exceptions (such as connection timeouts) per request.
    * If set to 0, we will not do any retries.

common/network-common/src/test/java/org/apache/spark/network/crypto/AuthIntegrationSuite.java

@@ -190,6 +190,20 @@ public void validateBlockMetaReqHandlerMatchesReceive() throws Exception {
     verify(delegate, never()).getMergedBlockMetaReqHandler();
   }
 
+  @Test
+  public void testStreamManagerFailsClosedWhenConfigured() {
+    // Pins the wiring: with the conf set, the handler must return the fail-closed wrapper.
+    RpcHandler delegate = mock(RpcHandler.class);
+    when(delegate.getStreamManager()).thenReturn(mock(StreamManager.class));
+    TransportConf conf = new TransportConf("rpc", new MapConfigProvider(
+      ImmutableMap.of("spark.network.auth.requireAuthForStreamRequests", "true")));
+    AuthRpcHandler handler = new AuthRpcHandler(
+      conf, mock(Channel.class), delegate, mock(SecretKeyHolder.class));
+
+    assertThrows(SecurityException.class,
+      () -> handler.getStreamManager().openStream("/jars/app.jar"));
+  }
+
   private static class DummyRpcHandler extends RpcHandler {
     @Override
     public void receive(

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

@@ -360,7 +360,8 @@ public void testRpcHandlerDelegate() throws Exception {
     // Tests all delegates exception for receive(), which is more complicated and already handled
     // by all other tests.
     RpcHandler handler = mock(RpcHandler.class);
-    RpcHandler saslHandler = new SaslRpcHandler(null, null, handler, null);
+    RpcHandler saslHandler = new SaslRpcHandler(
+      new TransportConf("shuffle", MapConfigProvider.EMPTY), null, handler, null);
 
     saslHandler.getStreamManager();
     verify(handler).getStreamManager();
@@ -372,6 +373,19 @@ public void testRpcHandlerDelegate() throws Exception {
     verify(handler).exceptionCaught(isNull(), isNull());
   }
 
+  @Test
+  public void testStreamManagerFailsClosedWhenConfigured() {
+    // Pins the wiring: with the conf set, the handler must return the fail-closed wrapper.
+    RpcHandler handler = mock(RpcHandler.class);
+    when(handler.getStreamManager()).thenReturn(mock(StreamManager.class));
+    TransportConf conf = new TransportConf("shuffle", new MapConfigProvider(
+      ImmutableMap.of("spark.network.auth.requireAuthForStreamRequests", "true")));
+    RpcHandler saslHandler = new SaslRpcHandler(conf, null, handler, null);
+
+    assertThrows(SecurityException.class,
+      () -> saslHandler.getStreamManager().openStream("/jars/app.jar"));
+  }
+
   @Test
   public void testDelegates() throws Exception {
     Method[] rpcHandlerMethods = RpcHandler.class.getDeclaredMethods();

common/network-common/src/test/java/org/apache/spark/network/server/AbstractAuthRpcHandlerSuite.java

@@ -0,0 +1,111 @@
+/*
+ * 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.Assert.*;
+import static org.mockito.Mockito.*;
+
+import java.nio.ByteBuffer;
+
+import io.netty.channel.Channel;
+import org.junit.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, boolean requireAuthForStreamRequests) {
+      super(delegate, requireAuthForStreamRequests);
+    }
+
+    @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, true);
+
+    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, true);
+    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);
+  }
+
+  @Test
+  public void testStreamManagerServedBeforeAuthWhenDisabled() {
+    // With the flag off the historical behavior is kept: the delegate's stream manager is
+    // returned unwrapped, so requests are served before authentication completes.
+    RpcHandler delegate = mock(RpcHandler.class);
+    StreamManager delegateManager = mock(StreamManager.class);
+    when(delegate.getStreamManager()).thenReturn(delegateManager);
+    AbstractAuthRpcHandler handler = new TestAuthRpcHandler(delegate, false);
+
+    assertSame(delegateManager, handler.getStreamManager());
+    assertFalse(handler.isAuthenticated());
+  }
+}

common/network-common/src/test/java/org/apache/spark/network/util/TransportConfSuite.java

@@ -0,0 +1,38 @@
+/*
+ * 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.util;
+
+import static org.junit.Assert.*;
+
+import com.google.common.collect.ImmutableMap;
+import org.junit.Test;
+
+public class TransportConfSuite {
+
+  @Test
+  public void testRequireAuthForStreamRequests() {
+    // The default must stay false on branch-3.5: the fail-closed check is opt-in so the
+    // patch release does not change runtime behavior for existing deployments.
+    assertFalse(new TransportConf("shuffle", MapConfigProvider.EMPTY)
+      .requireAuthForStreamRequests());
+    TransportConf enabled = new TransportConf("shuffle",
+      new MapConfigProvider(
+        ImmutableMap.of("spark.network.auth.requireAuthForStreamRequests", "true")));
+    assertTrue(enabled.requireAuthForStreamRequests());
+  }
+}