Load: Harden LOAD TSFILE authorization and source path validation

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java

@@ -298,6 +298,8 @@ public class IoTDBConfig {
     tierDataDirs[0][0] + File.separator + IoTDBConstant.LOAD_TSFILE_FOLDER_NAME
   };
 
+  private String[] loadTsFileAllowedDirs = new String[0];
+
   /** Strategy of multiple directories. */
   private String multiDirStrategyClassName = null;
 
@@ -1355,6 +1357,9 @@ private void formulateFolders() {
     for (int i = 0; i < loadActiveListeningDirs.length; i++) {
       loadActiveListeningDirs[i] = addDataHomeDir(loadActiveListeningDirs[i]);
     }
+    for (int i = 0; i < loadTsFileAllowedDirs.length; i++) {
+      loadTsFileAllowedDirs[i] = addDataHomeDir(loadTsFileAllowedDirs[i]);
+    }
     loadActiveListeningPipeDir = addDataHomeDir(loadActiveListeningPipeDir);
     loadActiveListeningFailDir = addDataHomeDir(loadActiveListeningFailDir);
     udfDir = addDataHomeDir(udfDir);
@@ -1560,6 +1565,19 @@ public String[] getLoadTsFileDirs() {
     return this.loadTsFileDirs;
   }
 
+  public String[] getLoadTsFileAllowedDirs() {
+    return this.loadTsFileAllowedDirs.length == 0
+        ? getLoadTsFileDirs()
+        : this.loadTsFileAllowedDirs;
+  }
+
+  public void setLoadTsFileAllowedDirs(String[] loadTsFileAllowedDirs) {
+    for (int i = 0; i < loadTsFileAllowedDirs.length; i++) {
+      loadTsFileAllowedDirs[i] = addDataHomeDir(loadTsFileAllowedDirs[i]);
+    }
+    this.loadTsFileAllowedDirs = loadTsFileAllowedDirs;
+  }
+
   public void formulateLoadTsFileDirs(String[][] tierDataDirs) {
     if (tierDataDirs.length < 1) {
       logger.warn("No data directory is set. loadTsFileDirs is kept as the default value.");

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java

@@ -2456,6 +2456,11 @@
                 "load_write_throughput_bytes_per_second",
                 String.valueOf(conf.getLoadWriteThroughputBytesPerSecond()))));
 
+    conf.setLoadTsFileAllowedDirs(
+        Arrays.stream(properties.getProperty("load_tsfile_allowed_dirs", "").trim().split(","))
+            .filter(dir -> !dir.isEmpty())
+            .toArray(String[]::new));
+
     conf.setLoadTabletConversionThresholdBytes(
         Long.parseLong(
             properties.getProperty(
@@ -2573,6 +2578,18 @@
                 ConfigurationFileUtils.getConfigurationDefaultValue(
                     "load_write_throughput_bytes_per_second"))));
 
+    conf.setLoadTsFileAllowedDirs(
+        Arrays.stream(
+                properties
+                    .getProperty(
+                        "load_tsfile_allowed_dirs",
+                        ConfigurationFileUtils.getConfigurationDefaultValue(
+                            "load_tsfile_allowed_dirs"))
+                    .trim()
+                    .split(","))
+            .filter(dir -> !dir.isEmpty())
+            .toArray(String[]::new));
+
     conf.setLoadActiveListeningEnable(
         Boolean.parseBoolean(
             properties.getProperty(

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/legacy/loader/TsFileLoader.java

@@ -55,7 +55,7 @@ public TsFileLoader(File tsFile, String database) {
   @Override
   public void load() {
     try {
-      LoadTsFileStatement statement = new LoadTsFileStatement(tsFile.getAbsolutePath());
+      LoadTsFileStatement statement = LoadTsFileStatement.createUnchecked(tsFile.getAbsolutePath());
       statement.setDeleteAfterLoad(true);
       statement.setConvertOnTypeMismatch(true);
       statement.setDatabaseLevel(parseSgLevel());

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java

@@ -585,7 +585,7 @@ private TSStatus loadTsFileAsync(final String dataBaseName, final List<String> a
 
   private TSStatus loadTsFileSync(final String dataBaseName, final String fileAbsolutePath)
       throws FileNotFoundException {
-    final LoadTsFileStatement statement = new LoadTsFileStatement(fileAbsolutePath);
+    final LoadTsFileStatement statement = LoadTsFileStatement.createUnchecked(fileAbsolutePath);
     statement.setDeleteAfterLoad(true);
     statement.setConvertOnTypeMismatch(true);
     statement.setVerifySchema(validateTsFile.get());

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java

@@ -443,14 +443,15 @@ private boolean handleSingleMiniFile(final int i) throws FileNotFoundException {
           isTableModelTsFile.get(i)
               ? loadTsFileDataTypeConverter
                   .convertForTableModel(
-                      new LoadTsFile(null, tsFiles.get(i).getPath(), Collections.emptyMap())
+                      LoadTsFile.createUnchecked(
+                              null, tsFiles.get(i).getPath(), Collections.emptyMap())
                           .setDatabase(databaseForTableData)
                           .setDeleteAfterLoad(isDeleteAfterLoad)
                           .setConvertOnTypeMismatch(isConvertOnTypeMismatch))
                   .orElse(null)
               : loadTsFileDataTypeConverter
                   .convertForTreeModel(
-                      new LoadTsFileStatement(tsFiles.get(i).getPath())
+                      LoadTsFileStatement.createUnchecked(tsFiles.get(i).getPath())
                           .setDeleteAfterLoad(isDeleteAfterLoad)
                           .setConvertOnTypeMismatch(isConvertOnTypeMismatch))
                   .orElse(null);
@@ -711,14 +712,15 @@ private void executeTabletConversionOnException(
             isTableModelTsFile.get(i)
                 ? loadTsFileDataTypeConverter
                     .convertForTableModel(
-                        new LoadTsFile(null, tsFiles.get(i).getPath(), Collections.emptyMap())
+                        LoadTsFile.createUnchecked(
+                                null, tsFiles.get(i).getPath(), Collections.emptyMap())
                             .setDatabase(databaseForTableData)
                             .setDeleteAfterLoad(isDeleteAfterLoad)
                             .setConvertOnTypeMismatch(isConvertOnTypeMismatch))
                     .orElse(null)
                 : loadTsFileDataTypeConverter
                     .convertForTreeModel(
-                        new LoadTsFileStatement(tsFiles.get(i).getPath())
+                        LoadTsFileStatement.createUnchecked(tsFiles.get(i).getPath())
                             .setDeleteAfterLoad(isDeleteAfterLoad)
                             .setConvertOnTypeMismatch(isConvertOnTypeMismatch))
                     .orElse(null);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/LoadTsFile.java

@@ -72,6 +72,19 @@ public class LoadTsFile extends Statement {
   private boolean needDecode4TimeColumn;
 
   public LoadTsFile(NodeLocation location, String filePath, Map<String, String> loadAttributes) {
+    this(location, filePath, loadAttributes, true);
+  }
+
+  public static LoadTsFile createUnchecked(
+      NodeLocation location, String filePath, Map<String, String> loadAttributes) {
+    return new LoadTsFile(location, filePath, loadAttributes, false);
+  }
+
+  private LoadTsFile(
+      NodeLocation location,
+      String filePath,
+      Map<String, String> loadAttributes,
+      boolean validateSourcePath) {
     super(location);
     this.filePath = requireNonNull(filePath, "filePath is null");
 
@@ -89,7 +102,7 @@ public LoadTsFile(NodeLocation location, String filePath, Map<String, String> lo
     try {
       this.tsFiles =
           org.apache.iotdb.db.queryengine.plan.statement.crud.LoadTsFileStatement.processTsFile(
-              new File(filePath));
+              new File(filePath), validateSourcePath);
       this.resources = new ArrayList<>();
       this.writePointCountList = new ArrayList<>();
       this.isTableModel = new ArrayList<>(Collections.nCopies(this.tsFiles.size(), true));
@@ -283,7 +296,7 @@ public List<LoadTsFile> getSubStatements() {
       final Map<String, String> properties = this.loadAttributes;
 
       final LoadTsFile subStatement =
-          new LoadTsFile(getLocation().orElse(null), filePath, properties);
+          LoadTsFile.createUnchecked(getLocation().orElse(null), filePath, properties);
 
       // Copy all configuration properties
       subStatement.databaseLevel = this.databaseLevel;

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java

@@ -169,7 +169,7 @@
   }
 
   @Override
   public void start() {
     try {
       stateMachine.transitionToRunning();
       int tsFileNodeListSize = tsFileNodeList.size();
@@ -586,14 +586,14 @@
             failedNode.isTableModel()
                 ? loadTsFileDataTypeConverter
                     .convertForTableModel(
-                        new LoadTsFile(null, filePath, Collections.emptyMap())
+                        LoadTsFile.createUnchecked(null, filePath, Collections.emptyMap())
                             .setDatabase(failedNode.getDatabase())
                             .setDeleteAfterLoad(failedNode.isDeleteAfterLoad())
                             .setConvertOnTypeMismatch(true))
                     .orElse(null)
                 : loadTsFileDataTypeConverter
                     .convertForTreeModel(
-                        new LoadTsFileStatement(filePath)
+                        LoadTsFileStatement.createUnchecked(filePath)
                             .setDeleteAfterLoad(failedNode.isDeleteAfterLoad())
                             .setConvertOnTypeMismatch(true))
                     .orElse(null);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/LoadTsFileStatement.java

@@ -35,7 +35,9 @@
 import java.io.File;
 import java.io.FileNotFoundException;
 import java.io.IOException;
+import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -71,6 +73,15 @@ public class LoadTsFileStatement extends Statement {
   private boolean needDecode4TimeColumn;
 
   public LoadTsFileStatement(String filePath) throws FileNotFoundException {
+    this(filePath, true);
+  }
+
+  public static LoadTsFileStatement createUnchecked(String filePath) throws FileNotFoundException {
+    return new LoadTsFileStatement(filePath, false);
+  }
+
+  private LoadTsFileStatement(String filePath, boolean validateSourcePath)
+      throws FileNotFoundException {
     this.file = new File(filePath).getAbsoluteFile();
     this.databaseLevel = IoTDBDescriptor.getInstance().getConfig().getDefaultDatabaseLevel();
     this.verifySchema = true;
@@ -80,14 +91,23 @@ public LoadTsFileStatement(String filePath) throws FileNotFoundException {
         IoTDBDescriptor.getInstance().getConfig().getLoadTabletConversionThresholdBytes();
     this.autoCreateDatabase = IoTDBDescriptor.getInstance().getConfig().isAutoCreateSchemaEnabled();
 
-    this.tsFiles = processTsFile(file);
+    this.tsFiles = processTsFile(file, validateSourcePath);
     this.resources = new ArrayList<>();
     this.writePointCountList = new ArrayList<>();
     this.isTableModel = new ArrayList<>(Collections.nCopies(this.tsFiles.size(), false));
     this.statementType = StatementType.MULTI_BATCH_INSERT;
   }
 
   public static List<File> processTsFile(final File file) throws FileNotFoundException {
+    return processTsFile(file, true);
+  }
+
+  public static List<File> processTsFile(final File file, final boolean validateSourcePath)
+      throws FileNotFoundException {
+    if (validateSourcePath) {
+      validateLoadSourcePath(file);
+    }
+
     final List<File> tsFiles = new ArrayList<>();
     if (file.isFile()) {
       tsFiles.add(file);
@@ -98,7 +118,7 @@ public static List<File> processTsFile(final File file) throws FileNotFoundExcep
                 "Can not find %s on this machine, notice that load can only handle files on this machine.",
                 file.getPath()));
       }
-      tsFiles.addAll(findAllTsFile(file));
+      tsFiles.addAll(findAllTsFile(file, validateSourcePath));
     }
     sortTsFiles(tsFiles);
     return tsFiles;
@@ -120,23 +140,59 @@ protected LoadTsFileStatement() {
     this.statementType = StatementType.MULTI_BATCH_INSERT;
   }
 
-  private static List<File> findAllTsFile(File file) {
+  private static List<File> findAllTsFile(File file, boolean validateSourcePath)
+      throws FileNotFoundException {
     final File[] files = file.listFiles();
     if (files == null) {
       return Collections.emptyList();
     }
 
     final List<File> tsFiles = new ArrayList<>();
     for (File nowFile : files) {
+      if (validateSourcePath) {
+        validateLoadSourcePath(nowFile);
+      }
       if (nowFile.getName().endsWith(TsFileConstant.TSFILE_SUFFIX)) {
         tsFiles.add(nowFile);
       } else if (nowFile.isDirectory()) {
-        tsFiles.addAll(findAllTsFile(nowFile));
+        tsFiles.addAll(findAllTsFile(nowFile, validateSourcePath));
       }
     }
     return tsFiles;
   }
 
+  public static void validateLoadSourcePath(final String filePath) throws FileNotFoundException {
+    validateLoadSourcePath(new File(filePath));
+  }
+
+  private static void validateLoadSourcePath(final File file) throws FileNotFoundException {
+    final Path sourcePath = canonicalPath(file);
+    final String[] allowedDirs =
+        IoTDBDescriptor.getInstance().getConfig().getLoadTsFileAllowedDirs();
+
+    for (final String allowedDir : allowedDirs) {
+      if (sourcePath.startsWith(canonicalPath(new File(allowedDir)))) {
+        return;
+      }
+    }
+
+    throw new FileNotFoundException(
+        String.format(
+            "Load TsFile source path %s is outside allowed directories %s.",
+            sourcePath, Arrays.toString(allowedDirs)));
+  }
+
+  private static Path canonicalPath(final File file) throws FileNotFoundException {
+    try {
+      return file.getCanonicalFile().toPath();
+    } catch (final IOException e) {
+      throw new FileNotFoundException(
+          String.format(
+              "Failed to resolve canonical path for Load TsFile source %s: %s",
+              file.getPath(), e.getMessage()));
+    }
+  }
+
   private static void sortTsFiles(List<File> files) {
     files.sort(
         (o1, o2) -> {
@@ -389,7 +445,7 @@ public List<PartialPath> getPaths() {
       loadAttributes.put(PIPE_GENERATED_KEY, String.valueOf(true));
     }
 
-    return new LoadTsFile(null, file.getAbsolutePath(), loadAttributes);
+    return LoadTsFile.createUnchecked(null, file.getAbsolutePath(), loadAttributes);
   }
 
   @Override

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/active/ActiveLoadTsFileLoader.java

@@ -219,7 +219,8 @@ private TSStatus loadTsFile(
       final ActiveLoadPendingQueue.ActiveLoadEntry entry, final IClientSession session)
       throws FileNotFoundException {
     final File tsFile = new File(entry.getFile());
-    final LoadTsFileStatement statement = new LoadTsFileStatement(tsFile.getAbsolutePath());
+    final LoadTsFileStatement statement =
+        LoadTsFileStatement.createUnchecked(tsFile.getAbsolutePath());
     final List<File> files = statement.getTsFiles();
 
     statement.setDeleteAfterLoad(true);

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/statement/crud/LoadTsFileStatementTest.java

@@ -25,6 +25,7 @@
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -38,10 +39,12 @@ public class LoadTsFileStatementTest {
   public void testSubStatementsKeepDatabase() throws Exception {
     final IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig();
     final int originalBatchSize = config.getLoadTsFileSubStatementBatchSize();
+    final String[] originalAllowedDirs = config.getLoadTsFileAllowedDirs().clone();
     final Path tempDir = Files.createTempDirectory("load-tsfile-sub-statements");
 
     try {
       config.setLoadTsFileSubStatementBatchSize(1);
+      config.setLoadTsFileAllowedDirs(new String[] {tempDir.toString()});
       Files.createFile(tempDir.resolve("a.tsfile"));
       Files.createFile(tempDir.resolve("b.tsfile"));
 
@@ -54,10 +57,42 @@ public void testSubStatementsKeepDatabase() throws Exception {
           subStatement -> Assert.assertEquals("test_db", subStatement.getDatabase()));
     } finally {
       config.setLoadTsFileSubStatementBatchSize(originalBatchSize);
+      config.setLoadTsFileAllowedDirs(originalAllowedDirs);
       deleteRecursively(tempDir);
     }
   }
 
+  @Test
+  public void testLoadSourcePathMustBeInAllowedDirs() throws Exception {
+    final IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig();
+    final String[] originalAllowedDirs = config.getLoadTsFileAllowedDirs().clone();
+    final Path allowedDir = Files.createTempDirectory("load-tsfile-allowed");
+    final Path deniedDir = Files.createTempDirectory("load-tsfile-denied");
+
+    try {
+      config.setLoadTsFileAllowedDirs(new String[] {allowedDir.toString()});
+      final Path deniedTsFile = Files.createFile(deniedDir.resolve("denied.tsfile"));
+      final Path traversalTsFile =
+          allowedDir.resolve("..").resolve(deniedDir.getFileName()).resolve("denied.tsfile");
+
+      assertLoadSourcePathRejected(deniedTsFile);
+      assertLoadSourcePathRejected(traversalTsFile);
+    } finally {
+      config.setLoadTsFileAllowedDirs(originalAllowedDirs);
+      deleteRecursively(allowedDir);
+      deleteRecursively(deniedDir);
+    }
+  }
+
+  private static void assertLoadSourcePathRejected(final Path sourcePath) {
+    try {
+      new LoadTsFileStatement(sourcePath.toString());
+      Assert.fail("Expected disallowed LOAD TSFILE source path to be rejected.");
+    } catch (final FileNotFoundException e) {
+      Assert.assertTrue(e.getMessage().contains("outside allowed directories"));
+    }
+  }
+
   private static void deleteRecursively(final Path path) throws IOException {
     if (path == null || !Files.exists(path)) {
       return;

iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template

@@ -2182,6 +2182,12 @@ load_clean_up_task_execution_delay_time_seconds=1800
 # Datatype: int
 load_write_throughput_bytes_per_second=-1
 
+# Comma-separated list of directories from which user-issued LOAD TSFILE statements can read.
+# If empty, IoTDB only permits LOAD sources under the internal load TsFile directories.
+# effectiveMode: hot_reload
+# Datatype: String
+load_tsfile_allowed_dirs=
+
 # Whether to enable the active listening mode for tsfile loading.
 # effectiveMode: hot_reload
 # Datatype: Boolean