Add security to shuffle service (INCOMPLETE)

This commit adds logic that allows the shuffle server to use
SASL to authenticate shuffle requests. This is currently not
working yet for two reasons:

(1) The ExternalShuffleClient doesn't actually have the
authentication logic yet. This will be implemented shortly
in a separate PR.

(2) All supported Yarn versions use a different version of
guava that does not have the base encoding util functions
that the SASL server uses. This will also be fixed in a
separate PR.

Author: Andrew Or

core/src/main/scala/org/apache/spark/storage/BlockManager.scala

@@ -95,7 +95,11 @@ private[spark] class BlockManager(
 
   private[spark]
   val externalShuffleServiceEnabled = conf.getBoolean("spark.shuffle.service.enabled", false)
-  private val externalShuffleServicePort = Utils.getExternalShuffleServicePort(conf)
+
+  // Port used by the external shuffle service. In Yarn mode, this may be already be
+  // set through the Hadoop configuration as the server is launched in the Yarn NM.
+  private val externalShuffleServicePort =
+    Utils.getSparkOrYarnConfig(conf, "spark.shuffle.service.port", "7337").toInt
 
   // Check that we're not using external shuffle service with consolidated shuffle files.
   if (externalShuffleServiceEnabled

core/src/main/scala/org/apache/spark/util/Utils.scala

@@ -1783,19 +1783,19 @@ private[spark] object Utils extends Logging {
     }.getOrElse("Unknown")
 
   /**
-   * Return the port used in the external shuffle service as specified through
-   * `spark.shuffle.service.port`. In Yarn, this is set in the Hadoop configuration.
+   * Return the value of a config either through the SparkConf or the Hadoop configuration
+   * if this is Yarn mode. In the latter case, this defaults to the value set through SparkConf
+   * if the key is not set in the Hadoop configuration.
    */
-  def getExternalShuffleServicePort(conf: SparkConf): Int = {
-    val shuffleServicePortKey = "spark.shuffle.service.port"
-    val sparkPort = conf.getInt(shuffleServicePortKey, 7337)
+  def getSparkOrYarnConfig(conf: SparkConf, key: String, default: String): String = {
+    val sparkValue = conf.get(key, default)
     if (SparkHadoopUtil.get.isYarnMode) {
-      val hadoopConf = SparkHadoopUtil.get.newConfiguration(conf)
-      hadoopConf.getInt(shuffleServicePortKey, sparkPort)
+      SparkHadoopUtil.get.newConfiguration(conf).get(key, sparkValue)
     } else {
-      sparkPort
+      sparkValue
     }
   }
+
 }
 
 /**

network/shuffle/src/main/java/org/apache/spark/network/sasl/ShuffleSecretManager.java

@@ -0,0 +1,117 @@
+/*
+ * 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.sasl;
+
+import java.lang.Override;
+import java.nio.ByteBuffer;
+import java.nio.charset.Charset;
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.spark.network.sasl.SecretKeyHolder;
+
+/**
+ * A class that manages shuffle secret used by the external shuffle service.
+ */
+public class ShuffleSecretManager implements SecretKeyHolder {
+  private final Logger logger = LoggerFactory.getLogger(ShuffleSecretManager.class);
+  private final ConcurrentHashMap<String, String> shuffleSecretMap;
+
+  private static final Charset UTF8_CHARSET = Charset.forName("UTF-8");
+
+  // Spark user used for authenticating SASL connections
+  // Note that this should match the value in org.apache.spark.SecurityManager
+  private static final String SPARK_SASL_USER = "sparkSaslUser";
+
+  /**
+   * Convert the given string to a byte buffer that can be converted back to a string
+   * through {@link #bytesToString(ByteBuffer)}. This is used if the external shuffle
+   * service represents shuffle secrets as bytes buffers instead of strings.
+   */
+  public static ByteBuffer stringToBytes(String s) {
+    return ByteBuffer.wrap(s.getBytes(UTF8_CHARSET));
+  }
+
+  /**
+   * Convert the given byte buffer to a string that can be converted back to a byte
+   * buffer through {@link #stringToBytes(String)}. This is used if the external shuffle
+   * service represents shuffle secrets as bytes buffers instead of strings.
+   */
+  public static String bytesToString(ByteBuffer b) {
+    return new String(b.array(), UTF8_CHARSET);
+  }
+
+  public ShuffleSecretManager() {
+    shuffleSecretMap = new ConcurrentHashMap<String, String>();
+  }
+
+  /**
+   * Register the specified application with its secret.
+   * Executors need to first authenticate themselves with the same secret before
+   * the fetching shuffle files written by other executors in this application.
+   */
+  public void registerApp(String appId, String shuffleSecret) {
+    if (!shuffleSecretMap.contains(appId)) {
+      shuffleSecretMap.put(appId, shuffleSecret);
+      logger.info("Registered shuffle secret for application {}", appId);
+    } else {
+      logger.debug("Application {} already registered", appId);
+    }
+  }
+
+  /**
+   * Register the specified application with its secret specified as a byte buffer.
+   */
+  public void registerApp(String appId, ByteBuffer shuffleSecret) {
+    registerApp(appId, bytesToString(shuffleSecret));
+  }
+
+  /**
+   * Unregister the specified application along with its secret.
+   * This is called when an application terminates.
+   */
+  public void unregisterApp(String appId) {
+    if (shuffleSecretMap.contains(appId)) {
+      shuffleSecretMap.remove(appId);
+      logger.info("Unregistered shuffle secret for application {}", appId);
+    } else {
+      logger.warn("Attempted to unregister application {} when it is not registered", appId);
+    }
+  }
+
+  /**
+   * Return the Spark user for authenticating SASL connections.
+   */
+  @Override
+  public String getSaslUser(String appId) {
+    return SPARK_SASL_USER;
+  }
+
+  /**
+   * Return the secret key registered with the specified application.
+   * This key is used to authenticate the executors in the application
+   * before they can fetch shuffle files from the external shuffle service.
+   * If the application is not registered, return null.
+   */
+  @Override
+  public String getSecretKey(String appId) {
+    return shuffleSecretMap.get(appId);
+  }
+}

network/yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleService.java

@@ -32,6 +32,8 @@
 import org.slf4j.LoggerFactory;
 
 import org.apache.spark.network.TransportContext;
+import org.apache.spark.network.sasl.SaslRpcHandler;
+import org.apache.spark.network.sasl.ShuffleSecretManager;
 import org.apache.spark.network.server.RpcHandler;
 import org.apache.spark.network.server.TransportServer;
 import org.apache.spark.network.shuffle.ExternalShuffleBlockHandler;
@@ -44,9 +46,18 @@
 public class YarnShuffleService extends AuxiliaryService {
   private final Logger logger = LoggerFactory.getLogger(YarnShuffleService.class);
 
+  // Port on which the shuffle server listens for fetch requests
   private static final String SPARK_SHUFFLE_SERVICE_PORT_KEY = "spark.shuffle.service.port";
   private static final int DEFAULT_SPARK_SHUFFLE_SERVICE_PORT = 7337;
 
+  // Whether the shuffle server should authenticate fetch requests
+  private static final String SPARK_AUTHENTICATE_KEY = "spark.authenticate";
+  private static final boolean DEFAULT_SPARK_AUTHENTICATE = false;
+
+  // An entity that manages the shuffle secret per application
+  // This is used only if authentication is enabled
+  private ShuffleSecretManager secretManager;
+
   // Actual server that serves the shuffle files
   private TransportServer shuffleServer = null;
 
@@ -55,58 +66,88 @@ public YarnShuffleService() {
     logger.info("Initializing Yarn shuffle service for Spark");
   }
 
+  /**
+   * Return whether authentication is enabled as specified by the configuration.
+   * If so, fetch requests will fail unless the appropriate authentication secret
+   * for the application is provided.
+   */
+  private boolean isAuthenticationEnabled() {
+    return secretManager != null;
+  }
+
   /**
    * Start the shuffle server with the given configuration.
    */
   @Override
   protected void serviceInit(Configuration conf) {
     try {
+      // If authentication is enabled, set up the shuffle server to use a
+      // special RPC handler that filters out unauthenticated fetch requests
+      boolean authEnabled = conf.getBoolean(SPARK_AUTHENTICATE_KEY, DEFAULT_SPARK_AUTHENTICATE);
+      RpcHandler rpcHandler = new ExternalShuffleBlockHandler();
+      if (authEnabled) {
+        secretManager = new ShuffleSecretManager();
+        rpcHandler = new SaslRpcHandler(rpcHandler, secretManager);
+      }
+
       int port = conf.getInt(
         SPARK_SHUFFLE_SERVICE_PORT_KEY, DEFAULT_SPARK_SHUFFLE_SERVICE_PORT);
       TransportConf transportConf = new TransportConf(new HadoopConfigProvider(conf));
-      RpcHandler rpcHandler = new ExternalShuffleBlockHandler();
       TransportContext transportContext = new TransportContext(transportConf, rpcHandler);
       shuffleServer = transportContext.createServer(port);
-      logger.info("Started Yarn shuffle service for Spark on port " + port);
+      String authEnabledString = authEnabled ? "enabled" : "not enabled";
+      logger.info("Started Yarn shuffle service for Spark on port {}. " +
+        "Authentication is {}.", port, authEnabledString);
     } catch (Exception e) {
       logger.error("Exception in starting Yarn shuffle service for Spark", e);
     }
   }
 
   @Override
   public void initializeApplication(ApplicationInitializationContext context) {
-    ApplicationId appId = context.getApplicationId();
-    logger.debug("Initializing application " + appId + "!");
+    String appId = context.getApplicationId().toString();
+    ByteBuffer shuffleSecret = context.getApplicationDataForService();
+    logger.debug("Initializing application {}", appId);
+    if (isAuthenticationEnabled()) {
+      secretManager.registerApp(appId, shuffleSecret);
+    }
   }
 
   @Override
   public void stopApplication(ApplicationTerminationContext context) {
-    ApplicationId appId = context.getApplicationId();
-    logger.debug("Stopping application " + appId + "!");
-  }
-
-  @Override
-  public ByteBuffer getMetaData() {
-    logger.debug("Getting meta data");
-    return ByteBuffer.allocate(0);
+    String appId = context.getApplicationId().toString();
+    logger.debug("Stopping application {}", appId);
+    if (isAuthenticationEnabled()) {
+      secretManager.unregisterApp(appId);
+    }
   }
 
   @Override
   public void initializeContainer(ContainerInitializationContext context) {
     ContainerId containerId = context.getContainerId();
-    logger.debug("Initializing container " + containerId + "!");
+    logger.debug("Initializing container {}", containerId);
   }
 
   @Override
   public void stopContainer(ContainerTerminationContext context) {
     ContainerId containerId = context.getContainerId();
-    logger.debug("Stopping container " + containerId + "!");
+    logger.debug("Stopping container {}", containerId);
   }
 
+  /**
+   * Close the shuffle server to clean up any associated state.
+   */
   @Override
   protected void serviceStop() {
     if (shuffleServer != null) {
       shuffleServer.close();
     }
   }
+
+  // Not currently used
+  @Override
+  public ByteBuffer getMetaData() {
+    return ByteBuffer.allocate(0);
+  }
+
 }

yarn/alpha/src/main/scala/org/apache/spark/deploy/yarn/ExecutorRunnable.scala

@@ -36,6 +36,7 @@ import org.apache.hadoop.yarn.ipc.YarnRPC
 import org.apache.hadoop.yarn.util.{Apps, ConverterUtils, Records, ProtoUtils}
 
 import org.apache.spark.{SecurityManager, SparkConf, Logging}
+import org.apache.spark.network.sasl.ShuffleSecretManager
 
 @deprecated("use yarn/stable", "1.2.0")
 class ExecutorRunnable(
@@ -93,7 +94,15 @@ class ExecutorRunnable(
     // If external shuffle service is enabled, register with the
     // Yarn shuffle service already started on the node manager
     if (sparkConf.getBoolean("spark.shuffle.service.enabled", false)) {
-      ctx.setServiceData(Map[String, ByteBuffer]("spark_shuffle" -> ByteBuffer.allocate(0)))
+      val secretString = securityMgr.getSecretKey()
+      val secretBytes =
+        if (secretString != null) {
+          ShuffleSecretManager.stringToBytes(secretString)
+        } else {
+          // Authentication is not enabled, so just provide dummy metadata
+          ByteBuffer.allocate(0)
+        }
+      ctx.setServiceData(Map[String, ByteBuffer]("spark_shuffle" -> secretBytes))
     }
 
     // Send the start request to the ContainerManager

yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/ExecutorRunnable.scala

@@ -36,6 +36,7 @@ import org.apache.hadoop.yarn.ipc.YarnRPC
 import org.apache.hadoop.yarn.util.{Apps, ConverterUtils, Records}
 
 import org.apache.spark.{SecurityManager, SparkConf, Logging}
+import org.apache.spark.network.sasl.ShuffleSecretManager
 
 
 class ExecutorRunnable(
@@ -92,7 +93,15 @@ class ExecutorRunnable(
     // If external shuffle service is enabled, register with the
     // Yarn shuffle service already started on the node manager
     if (sparkConf.getBoolean("spark.shuffle.service.enabled", false)) {
-      ctx.setServiceData(Map[String, ByteBuffer]("spark_shuffle" -> ByteBuffer.allocate(0)))
+      val secretString = securityMgr.getSecretKey()
+      val secretBytes =
+        if (secretString != null) {
+          ShuffleSecretManager.stringToBytes(secretString)
+        } else {
+          // Authentication is not enabled, so just provide dummy metadata
+          ByteBuffer.allocate(0)
+        }
+      ctx.setServiceData(Map[String, ByteBuffer]("spark_shuffle" -> secretBytes))
     }
 
     // Send the start request to the ContainerManager