Validate executor local directories and block ids in the external shuffle service

Add input validation to the external shuffle service for the executor-supplied
values it turns into filesystem paths.

* Block ids passed to removeBlocks must be plain file names (no path separator,
  NUL, "." or ".."). A malformed id is skipped instead of being used to build a
  path, and a batched call continues with the remaining ids.
* The localDirs an executor reports at RegisterExecutor are validated once at the
  RPC boundary in ExternalBlockHandler, before they are handed to either the block
  resolver or the push-based-shuffle merge manager. The check lives in a small
  LocalDirValidator built from the host's configured local-directory roots: each
  entry must canonicalize to a path contained under one of those roots and, when
  app scoping is required, within the registering application's own directory.

The roots are supplied per host through ExternalBlockHandler constructor overloads:
YarnShuffleService passes the NodeManager's local directories
(yarn.nodemanager.local-dirs) with app scoping, and ExternalShuffleService passes
the configured local directories (Utils.getConfiguredLocalDirs) without app
scoping. Block-id validation is covered by ExternalShuffleBlockResolverSuite,
local-directory validation by LocalDirValidatorSuite, and the standalone wiring by
ExternalShuffleServiceLocalDirsSuite; YarnShuffleServiceSuite fixtures are updated
so registered localDirs sit under the NM local dirs and contain the appId.

Co-authored-by: Peter Toth <peter.toth@gmail.com>

Author: huaxingao

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalBlockHandler.java

@@ -78,6 +78,7 @@ public class ExternalBlockHandler extends RpcHandler
   private final OneForOneStreamManager streamManager;
   private final ShuffleMetrics metrics;
   private final MergedShuffleFileManager mergeManager;
+  private final LocalDirValidator localDirValidator;
 
   public ExternalBlockHandler(TransportConf conf, File registeredExecutorFile)
     throws IOException {
@@ -86,6 +87,17 @@ public ExternalBlockHandler(TransportConf conf, File registeredExecutorFile)
       new NoOpMergedShuffleFileManager(conf, null));
   }
 
+  public ExternalBlockHandler(
+      TransportConf conf,
+      File registeredExecutorFile,
+      String[] allowedLocalDirs,
+      boolean requireAppScopedLocalDirs) throws IOException {
+    this(new OneForOneStreamManager(),
+      new ExternalShuffleBlockResolver(conf, registeredExecutorFile),
+      new NoOpMergedShuffleFileManager(conf, null),
+      new LocalDirValidator(allowedLocalDirs, requireAppScopedLocalDirs));
+  }
+
   public ExternalBlockHandler(
       TransportConf conf,
       File registeredExecutorFile,
@@ -94,11 +106,32 @@ public ExternalBlockHandler(
       new ExternalShuffleBlockResolver(conf, registeredExecutorFile), mergeManager);
   }
 
+  public ExternalBlockHandler(
+      TransportConf conf,
+      File registeredExecutorFile,
+      MergedShuffleFileManager mergeManager,
+      String[] allowedLocalDirs,
+      boolean requireAppScopedLocalDirs) throws IOException {
+    this(new OneForOneStreamManager(),
+      new ExternalShuffleBlockResolver(conf, registeredExecutorFile),
+      mergeManager,
+      new LocalDirValidator(allowedLocalDirs, requireAppScopedLocalDirs));
+  }
+
   @VisibleForTesting
   public ExternalShuffleBlockResolver getBlockResolver() {
     return blockManager;
   }
 
+  /**
+   * Validates the local directories an executor reports against the service's configured roots.
+   * Applied at the RPC boundary so both the block resolver and the merged-shuffle manager only
+   * consume validated paths. Throws {@link IllegalArgumentException} if an entry is not allowed.
+   */
+  public void validateLocalDirs(String[] localDirs, String appId) {
+    localDirValidator.validate(localDirs, appId);
+  }
+
   /** Enables mocking out the StreamManager and BlockManager. */
   @VisibleForTesting
   public ExternalBlockHandler(
@@ -113,10 +146,19 @@ public ExternalBlockHandler(
       OneForOneStreamManager streamManager,
       ExternalShuffleBlockResolver blockManager,
       MergedShuffleFileManager mergeManager) {
+    this(streamManager, blockManager, mergeManager, new LocalDirValidator(null, false));
+  }
+
+  private ExternalBlockHandler(
+      OneForOneStreamManager streamManager,
+      ExternalShuffleBlockResolver blockManager,
+      MergedShuffleFileManager mergeManager,
+      LocalDirValidator localDirValidator) {
     this.metrics = new ShuffleMetrics();
     this.streamManager = streamManager;
     this.blockManager = blockManager;
     this.mergeManager = mergeManager;
+    this.localDirValidator = localDirValidator;
   }
 
   @Override
@@ -185,6 +227,9 @@ protected void handleMessage(
         metrics.registerExecutorRequestLatencyMillis.time();
       try {
         checkAuth(client, msg.appId);
+        // Validate executor-supplied localDirs once here, at the RPC boundary, so both the block
+        // resolver and the merged-shuffle manager only ever consume validated paths.
+        validateLocalDirs(msg.executorInfo.localDirs, msg.appId);
         blockManager.registerExecutor(msg.appId, msg.execId, msg.executorInfo);
         mergeManager.registerExecutor(msg.appId, msg.executorInfo);
         callback.onSuccess(ByteBuffer.wrap(new byte[0]));

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleBlockResolver.java

@@ -376,6 +376,13 @@ public int removeBlocks(String appId, String execId, String[] blockIds) {
     }
     int numRemovedBlocks = 0;
     for (String blockId : blockIds) {
+      try {
+        validateBlockId(blockId);
+      } catch (IllegalArgumentException e) {
+        logger.warn("Skipping block with invalid id: {}",
+          MDC.of(LogKeys.BLOCK_ID, blockId));
+        continue;
+      }
       File file = new File(
         ExecutorDiskUtils.getFilePath(executor.localDirs, executor.subDirsPerLocalDir, blockId));
       if (file.delete()) {
@@ -422,6 +429,26 @@ public Cause diagnoseShuffleBlockCorruption(
       algorithm, checksumFile, reduceId, data, checksumByReader);
   }
 
+  /**
+   * Validates that a blockId is a plain file name before it is used to build a file path under
+   * {@code <localDir>/<subDirHex>/}. Legitimate Spark block ids contain only ASCII alphanumerics,
+   * underscores, dots and dashes (see {@code BlockId} in core), so a well-formed blockId never
+   * contains a path separator and is never {@code "."} or {@code ".."}. Such values are rejected.
+   */
+  @VisibleForTesting
+  static void validateBlockId(String blockId) {
+    if (blockId == null || blockId.isEmpty()) {
+      throw new IllegalArgumentException("blockId must not be null or empty");
+    }
+    if (blockId.indexOf('/') >= 0 || blockId.indexOf('\\') >= 0 || blockId.indexOf('\0') >= 0) {
+      throw new IllegalArgumentException(
+        "blockId must not contain a path separator or NUL: " + blockId);
+    }
+    if (blockId.equals("..") || blockId.equals(".")) {
+      throw new IllegalArgumentException("blockId must not be \".\" or \"..\": " + blockId);
+    }
+  }
+
   /** Simply encodes an executor's full ID, which is appId + execId. */
   public static class AppExecId {
     public final String appId;

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/LocalDirValidator.java

@@ -0,0 +1,115 @@
+/*
+ * 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.shuffle;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+
+import org.apache.spark.internal.SparkLogger;
+import org.apache.spark.internal.SparkLoggerFactory;
+import org.apache.spark.internal.LogKeys;
+import org.apache.spark.internal.MDC;
+
+/**
+ * Validates the local directories an executor reports at registration, before the external shuffle
+ * service turns them into filesystem paths. {@link ExternalBlockHandler} owns one of these (built
+ * from the host's configured local-directory roots) and applies it at the RPC boundary, so every
+ * consumer of those directories -- the block resolver and the merged-shuffle manager -- only ever
+ * operates on validated paths.
+ *
+ * Each entry must canonicalize to a path contained under one of the configured roots (YARN passes
+ * {@code yarn.nodemanager.local-dirs}; standalone passes the configured local directories) and,
+ * when app scoping is required, within the registering application's own directory (the
+ * {@code appId} must appear as a path segment). When no roots are configured there is nothing to
+ * contain the entries against and they are accepted. An empty array is always allowed.
+ */
+class LocalDirValidator {
+  private static final SparkLogger logger = SparkLoggerFactory.getLogger(LocalDirValidator.class);
+
+  private final List<Path> allowedLocalDirRoots;
+  private final boolean requireAppScopedLocalDirs;
+
+  LocalDirValidator(String[] allowedLocalDirs, boolean requireAppScopedLocalDirs) {
+    this.allowedLocalDirRoots = canonicalizeRoots(allowedLocalDirs);
+    this.requireAppScopedLocalDirs = requireAppScopedLocalDirs;
+  }
+
+  void validate(String[] localDirs, String appId) {
+    if (allowedLocalDirRoots.isEmpty()) {
+      return;
+    }
+    if (localDirs == null) {
+      throw new IllegalArgumentException("localDirs must not be null");
+    }
+    for (String localDir : localDirs) {
+      if (localDir == null || localDir.isEmpty()) {
+        throw new IllegalArgumentException("localDirs entries must be non-null and non-empty");
+      }
+      validateContained(localDir, appId);
+    }
+  }
+
+  private void validateContained(String localDir, String appId) {
+    Path canonical;
+    try {
+      canonical = new File(localDir).getCanonicalFile().toPath();
+    } catch (IOException e) {
+      throw new IllegalArgumentException(
+        "localDirs entry could not be canonicalized: " + localDir, e);
+    }
+    if (allowedLocalDirRoots.stream().noneMatch(canonical::startsWith)) {
+      throw new IllegalArgumentException(
+        "localDirs entry is not under any of the service's local directories: " + localDir);
+    }
+    if (requireAppScopedLocalDirs &&
+        (appId == null || appId.isEmpty() || !pathContainsSegment(canonical, appId))) {
+      throw new IllegalArgumentException(
+        "localDirs entry is not within application " + appId + "'s directory: " + localDir);
+    }
+  }
+
+  private static boolean pathContainsSegment(Path path, String segment) {
+    for (Path part : path) {
+      if (part.toString().equals(segment)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  private static List<Path> canonicalizeRoots(String[] allowedLocalDirs) {
+    List<Path> roots = new ArrayList<>();
+    if (allowedLocalDirs != null) {
+      for (String root : allowedLocalDirs) {
+        if (root == null || root.isEmpty()) {
+          continue;
+        }
+        try {
+          roots.add(new File(root).getCanonicalFile().toPath());
+        } catch (IOException e) {
+          logger.warn("Ignoring local dir root that could not be canonicalized: {}",
+            MDC.of(LogKeys.PATH, root));
+        }
+      }
+    }
+    return roots;
+  }
+}

common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleBlockResolverSuite.java

@@ -17,9 +17,12 @@
 
 package org.apache.spark.network.shuffle;
 
+import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.util.UUID;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.spark.network.shuffle.protocol.ExecutorShuffleInfo;
@@ -120,4 +123,103 @@ public void jsonSerializationOfExecutorRegistration() throws IOException {
     assertEquals(shuffleInfo, mapper.readValue(legacyShuffleJson, ExecutorShuffleInfo.class));
   }
 
+  /**
+   * removeBlocks only operates on plain block-id file names. A block id containing parent-directory
+   * segments is skipped: it is not removed, and a file located outside the executor's local
+   * directory is left untouched.
+   */
+  @Test
+  public void removeBlocksIgnoresBlockIdWithParentDirSegments() throws IOException {
+    TestShuffleDataContext context = new TestShuffleDataContext(1, 5);
+    context.create();
+    File outsideFile = null;
+    try {
+      File localDir = new File(context.localDirs[0]);
+      File parentOfLocalDir = localDir.getParentFile();
+      outsideFile = new File(parentOfLocalDir, "outside-" + UUID.randomUUID() + ".txt");
+      Files.writeString(outsideFile.toPath(), "should not be deleted", StandardCharsets.UTF_8);
+      assertTrue(outsideFile.exists(), "precondition: file was created");
+
+      ExternalShuffleBlockResolver resolver = new ExternalShuffleBlockResolver(conf, null);
+      resolver.registerExecutor("app0", "0", context.createExecutorInfo(SORT_MANAGER));
+
+      String blockId = ".." + File.separator + ".." + File.separator + outsideFile.getName();
+      int removed = resolver.removeBlocks("app0", "0", new String[] { blockId });
+
+      assertEquals(0, removed, "A block id with parent-dir segments must not remove a block");
+      assertTrue(outsideFile.exists(),
+        "A file outside the local directory must not be removed (" +
+          outsideFile.getAbsolutePath() + ")");
+    } finally {
+      if (outsideFile != null && outsideFile.exists()) {
+        assertTrue(outsideFile.delete() || !outsideFile.exists());
+      }
+      context.cleanup();
+    }
+  }
+
+  /**
+   * A block id containing a path separator is not a plain file name and is skipped, so removeBlocks
+   * reports zero removed blocks.
+   */
+  @Test
+  public void removeBlocksIgnoresBlockIdWithSeparator() throws IOException {
+    TestShuffleDataContext context = new TestShuffleDataContext(1, 5);
+    context.create();
+    try {
+      ExternalShuffleBlockResolver resolver = new ExternalShuffleBlockResolver(conf, null);
+      resolver.registerExecutor("app0", "0", context.createExecutorInfo(SORT_MANAGER));
+
+      String blockId = "sub" + File.separator + "child";
+      int removed = resolver.removeBlocks("app0", "0", new String[] { blockId });
+
+      assertEquals(0, removed, "A block id containing a path separator must be skipped");
+    } finally {
+      context.cleanup();
+    }
+  }
+
+  /**
+   * In a batched call, an invalid block id does not abort the whole operation: a valid block id in
+   * the same batch is still removed while the invalid one is skipped.
+   */
+  @Test
+  public void removeBlocksContinuesAfterIgnoringInvalidBlockId() throws IOException {
+    TestShuffleDataContext context = new TestShuffleDataContext(1, 5);
+    context.create();
+    File outsideFile = null;
+    try {
+      File localDir = new File(context.localDirs[0]);
+      File parentOfLocalDir = localDir.getParentFile();
+      outsideFile = new File(parentOfLocalDir, "outside-" + UUID.randomUUID() + ".txt");
+      Files.writeString(outsideFile.toPath(), "should not be deleted", StandardCharsets.UTF_8);
+
+      // Plant a real, legitimately-named block in the localDir so the valid removal has something
+      // to actually remove.
+      context.insertCachedRddData(7, 0, new byte[] { 1, 2, 3 });
+      String validBlockId = "rdd_7_0";
+      File validFile = new File(ExecutorDiskUtils.getFilePath(
+        context.localDirs, context.subDirsPerLocalDir, validBlockId));
+      assertTrue(validFile.exists(), "precondition: block file was created");
+
+      ExternalShuffleBlockResolver resolver = new ExternalShuffleBlockResolver(conf, null);
+      resolver.registerExecutor("app0", "0", context.createExecutorInfo(SORT_MANAGER));
+
+      String invalidBlockId =
+        ".." + File.separator + ".." + File.separator + outsideFile.getName();
+      int removed = resolver.removeBlocks(
+        "app0", "0", new String[] { invalidBlockId, validBlockId });
+
+      assertEquals(1, removed, "The valid block id in the batch should still be removed");
+      assertTrue(outsideFile.exists(),
+        "The invalid entry must not remove a file outside the local directory");
+      assertFalse(validFile.exists(), "The valid block file should have been removed by the batch");
+    } finally {
+      if (outsideFile != null && outsideFile.exists()) {
+        assertTrue(outsideFile.delete() || !outsideFile.exists());
+      }
+      context.cleanup();
+    }
+  }
+
 }