From 0da60ca1a4754f858cf7c52dd4f0c97ae0e1b0cb Mon Sep 17 00:00:00 2001
From: Anupam Aggarwal <aaggarwal@confluent.io>
Date: Wed, 13 Mar 2024 20:31:09 +0530
Subject: [PATCH] [FLINK-29122][core] Improve robustness of
 FileUtils.expandDirectory() (#24307)

---
 .../java/org/apache/flink/util/FileUtils.java | 14 ++++++
 .../org/apache/flink/util/FileUtilsTest.java  | 49 +++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/flink-core/src/main/java/org/apache/flink/util/FileUtils.java b/flink-core/src/main/java/org/apache/flink/util/FileUtils.java
index e09ec67f98753..1571906a00683 100644
--- a/flink-core/src/main/java/org/apache/flink/util/FileUtils.java
+++ b/flink-core/src/main/java/org/apache/flink/util/FileUtils.java
@@ -515,12 +515,22 @@ private static void addToZip(
         }
     }
 
+    /**
+     * Un-archives files inside the target directory. Illegal fs access outside target directory is
+     * not permitted.
+     *
+     * @param file path to zipped/archived file
+     * @param targetDirectory directory path where file needs to be unarchived
+     * @return path to folder with unarchived contents
+     * @throws IOException if file open fails or in case of unsafe access outside target directory
+     */
     public static Path expandDirectory(Path file, Path targetDirectory) throws IOException {
         FileSystem sourceFs = file.getFileSystem();
         FileSystem targetFs = targetDirectory.getFileSystem();
         Path rootDir = null;
         try (ZipInputStream zis = new ZipInputStream(sourceFs.open(file))) {
             ZipEntry entry;
+            String targetDirStr = targetDirectory.toString();
             while ((entry = zis.getNextEntry()) != null) {
                 Path relativePath = new Path(entry.getName());
                 if (rootDir == null) {
@@ -529,6 +539,10 @@ public static Path expandDirectory(Path file, Path targetDirectory) throws IOExc
                 }
 
                 Path newFile = new Path(targetDirectory, relativePath);
+                if (!newFile.toString().startsWith(targetDirStr)) {
+                    throw new IOException("Illegal escape from target directory");
+                }
+
                 if (entry.isDirectory()) {
                     targetFs.mkdirs(newFile);
                 } else {
diff --git a/flink-core/src/test/java/org/apache/flink/util/FileUtilsTest.java b/flink-core/src/test/java/org/apache/flink/util/FileUtilsTest.java
index 0a9953c060bd4..1576305d0aa9a 100644
--- a/flink-core/src/test/java/org/apache/flink/util/FileUtilsTest.java
+++ b/flink-core/src/test/java/org/apache/flink/util/FileUtilsTest.java
@@ -51,6 +51,8 @@
 import java.util.Random;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipOutputStream;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -517,6 +519,53 @@ private void verifyDirectoryCompression(
         assertDirEquals(compressDir.resolve(originalDir), extractDir.resolve(originalDir));
     }
 
+    /**
+     * Generates zip archive, containing path to file/dir passed in, un-archives the generated zip
+     * using {@link org.apache.flink.util.FileUtils#expandDirectory(org.apache.flink.core.fs.Path,
+     * org.apache.flink.core.fs.Path)} and returns path to expanded folder.
+     *
+     * @param prefix prefix to use for creating source and destination folders
+     * @param path path to contents of zip
+     * @return Path to folder containing unzipped contents
+     * @throws IOException if I/O error in file creation, un-archiving detects unsafe access outside
+     *     target folder
+     */
+    private Path writeZipAndFetchExpandedPath(String prefix, String path) throws IOException {
+        // random source folder
+        String sourcePath = prefix + "src";
+        String dstPath = prefix + "dst";
+        java.nio.file.Path srcFolder = TempDirUtils.newFolder(temporaryFolder, sourcePath).toPath();
+        java.nio.file.Path zippedFile = srcFolder.resolve("file.zip");
+        ZipOutputStream out = new ZipOutputStream(Files.newOutputStream(zippedFile));
+        ZipEntry e1 = new ZipEntry(path);
+        out.putNextEntry(e1);
+        out.close();
+        java.nio.file.Path dstFolder = TempDirUtils.newFolder(temporaryFolder, dstPath).toPath();
+        return FileUtils.expandDirectory(
+                new Path(zippedFile.toString()), new Path(dstFolder.toString()));
+    }
+
+    @Test
+    public void testExpandDirWithValidPaths() {
+        Assertions.assertDoesNotThrow(() -> writeZipAndFetchExpandedPath("t0", "/level1/level2/"));
+        Assertions.assertDoesNotThrow(
+                () -> writeZipAndFetchExpandedPath("t1", "/level1/level2/file.txt"));
+        Assertions.assertDoesNotThrow(
+                () -> writeZipAndFetchExpandedPath("t2", "/level1/../level1/file.txt"));
+        Assertions.assertDoesNotThrow(
+                () -> writeZipAndFetchExpandedPath("t3", "/level1/level2/level3/../"));
+        Assertions.assertThrows(
+                IOException.class,
+                () -> writeZipAndFetchExpandedPath("t2", "/level1/level2/../../pass"));
+    }
+
+    @Test
+    public void testExpandDirWithForbiddenEscape() {
+        Assertions.assertThrows(
+                IOException.class, () -> writeZipAndFetchExpandedPath("t1", "/../../"));
+        Assertions.assertThrows(IOException.class, () -> writeZipAndFetchExpandedPath("t2", "../"));
+    }
+
     /**
      * Generates a random content file.
      *