[CELEBORN-2468] Support pluggable transport authentication and authorization

client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/client/FlinkShuffleClientImpl.java

@@ -84,7 +84,7 @@ public class FlinkShuffleClientImpl extends ShuffleClientImpl {
       JavaUtils.newConcurrentHashMap();
   private long driverTimestamp;
 
-  private final TransportContext context;
+  private TransportContext context;
 
   /** The buffer size bytes in flink, default value is 32KB. */
   private final int bufferSizeBytes;
@@ -163,6 +163,9 @@ public void shutdown() {
     if (flinkTransportClientFactory != null) {
       flinkTransportClientFactory.close();
     }
+    if (context != null) {
+      context.close();
+    }
     if (readClientHandler != null) {
       readClientHandler.close();
     }
@@ -191,13 +194,12 @@ public FlinkShuffleClientImpl(
               BATCH_HEADER_SIZE, BufferUtils.HEADER_LENGTH_PREFIX));
     }
     this.bufferSizeBytes = bufferSizeBytes;
-    String module = TransportModuleConstants.DATA_MODULE;
-    TransportConf dataTransportConf =
-        Utils.fromCelebornConf(conf, module, conf.networkIoThreads(module));
-    this.context =
-        new TransportContext(
-            dataTransportConf, readClientHandler, conf.clientCloseIdleConnections());
-    this.setupLifecycleManagerRef(driverHost, port);
+    try {
+      this.setupLifecycleManagerRef(driverHost, port);
+    } catch (RuntimeException | Error e) {
+      super.shutdown();
+      throw e;
+    }
     this.driverTimestamp = driverTimestamp;
     this.openStreamThreadPool =
         ThreadUtils.newDaemonCachedThreadPool(
@@ -206,8 +208,21 @@ public FlinkShuffleClientImpl(
 
   private void initializeTransportClientFactory() {
     if (null == flinkTransportClientFactory) {
-      flinkTransportClientFactory =
-          new FlinkTransportClientFactory(context, createBootstraps(), bufferSizeBytes);
+      String module = TransportModuleConstants.DATA_MODULE;
+      TransportConf dataTransportConf =
+          Utils.fromCelebornConf(conf, module, conf.networkIoThreads(module));
+      TransportContext newContext =
+          new TransportContext(
+              dataTransportConf, readClientHandler, conf.clientCloseIdleConnections());
+      try {
+        FlinkTransportClientFactory factory =
+            new FlinkTransportClientFactory(newContext, createBootstraps(), bufferSizeBytes);
+        context = newContext;
+        flinkTransportClientFactory = factory;
+      } catch (RuntimeException | Error e) {
+        newContext.close();
+        throw e;
+      }
     }
   }
 

client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/FlinkBootstrapSuiteJ.java

@@ -0,0 +1,141 @@
+/*
+ * 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.celeborn.plugin.flink;
+
+import static org.junit.Assert.*;
+import static org.mockito.Mockito.mock;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.network.client.TransportClient;
+import org.apache.celeborn.common.network.client.TransportClientBootstrap;
+import org.apache.celeborn.common.network.util.TransportConf;
+import org.apache.celeborn.common.rpc.RpcEndpointRef;
+import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl;
+
+public class FlinkBootstrapSuiteJ {
+  @Before
+  public void resetCounts() {
+    CountingBootstrap.created.set(0);
+    CountingBootstrap.closed.set(0);
+  }
+
+  @Test
+  public void shutdownClosesBothDataContextsExactlyOnce() {
+    FlinkShuffleClientImpl client = new TestClient(configuration(false));
+    try {
+      client.getDataClientFactory();
+      assertEquals(2, CountingBootstrap.created.get());
+    } finally {
+      client.shutdown();
+      client.shutdown();
+    }
+    assertEquals(2, CountingBootstrap.closed.get());
+  }
+
+  @Test
+  public void readContextWaitsForLifecycleManagerAndRecoversFromInitializationFailure() {
+    AtomicInteger calls = new AtomicInteger();
+    FlinkShuffleClientImpl client =
+        new TestClient(configuration(true)) {
+          @Override
+          protected List<TransportClientBootstrap> createBootstraps() {
+            // The first call initializes push; the second initializes Flink read.
+            if (calls.incrementAndGet() == 2) {
+              throw new IllegalStateException("Application registration is unavailable");
+            }
+            return Collections.emptyList();
+          }
+        };
+    try {
+      assertEquals(0, CountingBootstrap.created.get());
+      RpcEndpointRef endpoint = mock(RpcEndpointRef.class);
+      assertThrows(IllegalStateException.class, () -> client.setupLifecycleManagerRef(endpoint));
+      assertEquals(CountingBootstrap.closed.get() + 1, CountingBootstrap.created.get());
+      client.setupLifecycleManagerRef(endpoint);
+      assertNotNull(client.getDataClientFactory());
+      assertEquals(CountingBootstrap.closed.get() + 2, CountingBootstrap.created.get());
+    } finally {
+      client.shutdown();
+    }
+    assertEquals(CountingBootstrap.created.get(), CountingBootstrap.closed.get());
+  }
+
+  @Test
+  public void lifecycleManagerSetupFailureClosesTheConstructedPushContext() {
+    assertThrows(
+        IllegalStateException.class,
+        () ->
+            new FlinkShuffleClientImpl(
+                "application", "localhost", 1232, 1L, configuration(false), null, 64) {
+              @Override
+              public void setupLifecycleManagerRef(String host, int port) {
+                throw new IllegalStateException("LifecycleManager is unavailable");
+              }
+            });
+    assertEquals(1, CountingBootstrap.created.get());
+    assertEquals(1, CountingBootstrap.closed.get());
+  }
+
+  private static CelebornConf configuration(boolean authenticated) {
+    return new CelebornConf(false)
+        .set("celeborn.auth.enabled", Boolean.toString(authenticated))
+        .set("celeborn.internal.port.enabled", "true")
+        .set("celeborn.data.client.bootstrap.classes", CountingBootstrap.class.getName())
+        .set("celeborn.data.io.clientThreads", "1")
+        .set("celeborn.rpc_app_client.io.clientThreads", "1")
+        .set("celeborn.rpc_app_client.dispatcher.threads", "1");
+  }
+
+  private static class TestClient extends FlinkShuffleClientImpl {
+    TestClient(CelebornConf conf) {
+      super("application", "localhost", 1232, 1L, conf, null, 64);
+    }
+
+    @Override
+    public void setupLifecycleManagerRef(String host, int port) {}
+
+    @Override
+    protected List<TransportClientBootstrap> createBootstraps() {
+      return Collections.emptyList();
+    }
+  }
+
+  public static class CountingBootstrap implements TransportClientBootstrap {
+    static final AtomicInteger created = new AtomicInteger();
+    static final AtomicInteger closed = new AtomicInteger();
+
+    public CountingBootstrap(TransportConf conf) {
+      created.incrementAndGet();
+    }
+
+    @Override
+    public void doBootstrap(TransportClient client) {}
+
+    @Override
+    public void close() {
+      closed.incrementAndGet();
+    }
+  }
+}

client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java

@@ -25,7 +25,10 @@
 import java.util.Map;
 import java.util.Random;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
 
+import org.apache.commons.lang3.reflect.FieldUtils;
 import org.apache.flink.api.common.BatchShuffleMode;
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.configuration.Configuration;
@@ -53,6 +56,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import org.apache.celeborn.client.ApplicationHeartbeater;
 import org.apache.celeborn.common.CelebornConf;
 import org.apache.celeborn.common.protocol.FallbackPolicy;
 import org.apache.celeborn.common.util.Utils$;
@@ -72,10 +76,6 @@ public void setUp() {
     int startPort = Utils$.MODULE$.selectRandomInt(1024, 65535);
     configuration.setString("celeborn.master.port", String.valueOf(startPort));
     configuration.setString("celeborn.master.endpoints", "localhost:" + startPort);
-    // Set a large app heartbeat interval to avoid the app heartbeater resetting
-    // LifecycleManager#shuffleCount/shuffleFallbackCounts via sumThenReset in the
-    // middle of a test, which makes the count assertions below flaky.
-    configuration.setString("celeborn.client.application.heartbeatInterval", "1h");
     remoteShuffleMaster = createShuffleMaster(configuration);
   }
 
@@ -111,10 +111,12 @@ public void testRegisterJobWithForceFallbackPolicy() {
 
   @Test
   public void testRegisterPartitionWithProducer()
-      throws UnknownHostException, ExecutionException, InterruptedException {
+      throws UnknownHostException, ExecutionException, InterruptedException,
+          IllegalAccessException {
     JobID jobID = JobID.generate();
     JobShuffleContext jobShuffleContext = createJobShuffleContext(jobID);
     remoteShuffleMaster.registerJob(jobShuffleContext);
+    stopApplicationHeartbeater();
 
     IntermediateDataSetID intermediateDataSetID = new IntermediateDataSetID();
     PartitionDescriptor partitionDescriptor = createPartitionDescriptor(intermediateDataSetID, 0);
@@ -165,13 +167,15 @@ public void testRegisterPartitionWithProducer()
 
   @Test
   public void testRegisterPartitionWithProducerForForceFallbackPolicy()
-      throws UnknownHostException, ExecutionException, InterruptedException {
+      throws UnknownHostException, ExecutionException, InterruptedException,
+          IllegalAccessException {
     configuration.setString(
         CelebornConf.FLINK_SHUFFLE_FALLBACK_POLICY().key(), FallbackPolicy.ALWAYS.name());
     remoteShuffleMaster = createShuffleMaster(configuration, new NettyShuffleServiceFactory());
     JobID jobID = JobID.generate();
     JobShuffleContext jobShuffleContext = createJobShuffleContext(jobID);
     remoteShuffleMaster.registerJob(jobShuffleContext);
+    stopApplicationHeartbeater();
 
     IntermediateDataSetID intermediateDataSetID = new IntermediateDataSetID();
     PartitionDescriptor partitionDescriptor = createPartitionDescriptor(intermediateDataSetID, 0);
@@ -305,6 +309,20 @@ public void testInvalidShuffleConfig() {
         () -> createShuffleMaster(configuration));
   }
 
+  private void stopApplicationHeartbeater() throws IllegalAccessException, InterruptedException {
+    // The first heartbeat is immediate regardless of its interval and drains these counters.
+    // Stop it before checking cumulative counts, and wait for any in-progress callback to exit.
+    ApplicationHeartbeater heartbeater =
+        (ApplicationHeartbeater)
+            FieldUtils.readField(remoteShuffleMaster.lifecycleManager(), "heartbeater", true);
+    ExecutorService executor =
+        (ExecutorService) FieldUtils.readField(heartbeater, "appHeartbeatHandlerThread", true);
+    heartbeater.stop();
+    Assert.assertTrue(
+        "Application heartbeater did not terminate",
+        executor.awaitTermination(10, TimeUnit.SECONDS));
+  }
+
   @After
   public void tearDown() {
     if (remoteShuffleMaster != null) {

client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java

@@ -230,9 +230,14 @@ public ShuffleClientImpl(String appUniqueId, CelebornConf conf, UserIdentifier u
             Role.CLIENT(),
             scala.None$.empty());
 
-    String module = TransportModuleConstants.DATA_MODULE;
-    dataTransportConf = Utils.fromCelebornConf(conf, module, conf.networkIoThreads(module));
-    initDataClientFactoryIfNeeded();
+    try {
+      String module = TransportModuleConstants.DATA_MODULE;
+      dataTransportConf = Utils.fromCelebornConf(conf, module, conf.networkIoThreads(module));
+      initDataClientFactoryIfNeeded();
+    } catch (RuntimeException | Error e) {
+      rpcEnv.shutdown();
+      throw e;
+    }
     int pushDataRetryThreads = conf.clientPushRetryThreads();
     pushDataRetryPool =
         ThreadUtils.newDaemonCachedThreadPool("celeborn-retry-sender", pushDataRetryThreads, 60);
@@ -271,19 +276,20 @@ protected List<TransportClientBootstrap> createBootstraps() {
   }
 
   private void initDataClientFactoryIfNeeded() {
-    if (dataClientFactory != null) {
+    if (dataClientFactory != null || (authEnabled && lifecycleManagerRef == null)) {
       return;
     }
-    this.transportContext =
+    TransportContext context =
         new TransportContext(
             dataTransportConf, new BaseMessageHandler(), conf.clientCloseIdleConnections());
-    if (!authEnabled) {
+    try {
       logger.info("Initializing data client factory for {}.", appUniqueId);
-      dataClientFactory = transportContext.createClientFactory();
-    } else if (lifecycleManagerRef != null) {
-      logger.info("Initializing data client factory for secured {}.", appUniqueId);
-      List<TransportClientBootstrap> bootstraps = createBootstraps();
-      dataClientFactory = transportContext.createClientFactory(bootstraps);
+      TransportClientFactory factory = context.createClientFactory(createBootstraps());
+      transportContext = context;
+      dataClientFactory = factory;
+    } catch (RuntimeException | Error e) {
+      context.close();
+      throw e;
     }
   }
 

client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala

@@ -48,6 +48,8 @@ import org.apache.celeborn.common.meta.{ApplicationMeta, ShufflePartitionLocatio
 import org.apache.celeborn.common.metrics.source.Role
 import org.apache.celeborn.common.network.protocol.{SerdeVersion, TransportMessagesHelper}
 import org.apache.celeborn.common.network.sasl.registration.RegistrationInfo
+import org.apache.celeborn.common.network.security.AuthorizationRequest
+import org.apache.celeborn.common.network.security.SecurityOperation._
 import org.apache.celeborn.common.protocol._
 import org.apache.celeborn.common.protocol.RpcNameConstants.WORKER_EP
 import org.apache.celeborn.common.protocol.message.ControlMessages._
@@ -357,6 +359,34 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
     shufflePartitionType.getOrDefault(shuffleId, partitionType)
   }
 
+  // Resources and secrets belong to this LM's appUniqueId, regardless of a request's app claim.
+  // Cleanup and stage completion require a local origin.
+  override def authorize(context: RpcRequestContext, message: Any): Any = message match {
+    case RemoveExpiredShuffle | _: StageEnd =>
+      context.requireLocal()
+      message
+    case _ =>
+      val operation = message match {
+        case _: RegisterShuffle => REGISTER_SHUFFLE
+        case _: PbRegisterMapPartitionTask => REGISTER_MAP_PARTITION_TASK
+        case _: Revive => REVIVE
+        case _: PbPartitionSplit => PARTITION_SPLIT
+        case _: MapperEnd => MAPPER_END
+        case _: ReadReducerPartitionEnd => READ_REDUCER_PARTITION_END
+        case _: GetReducerFileGroup => GET_REDUCER_FILE_GROUP
+        case _: PbGetStageEnd => GET_STAGE_END
+        case _: PbGetShuffleId => GET_SHUFFLE_ID
+        case _: PbReportShuffleFetchFailure => REPORT_SHUFFLE_FETCH_FAILURE
+        case _: PbReportBarrierStageAttemptFailure => REPORT_BARRIER_STAGE_ATTEMPT_FAILURE
+        case _: PbApplicationMetaRequest => GET_APPLICATION_META
+        case _ =>
+          throw new SecurityException(
+            s"Unsupported LifecycleManager request: ${message.getClass.getName}")
+      }
+      context.authorize(AuthorizationRequest.forApplication(operation, appUniqueId))
+      message
+  }
+
   override def receive: PartialFunction[Any, Unit] = {
     case RemoveExpiredShuffle =>
       removeExpiredShuffle()