[api] fix issue in Tar/Zip Utils that resulted in incorrect artifact …

…extraction
deepjavalibrary · Nov 27, 2024 · cd00a44 · cd00a44
1 parent 73f7cb0
commit cd00a44
Show file tree

Hide file tree

Showing 10 changed files with 57 additions and 24 deletions.
diff --git a/api/build.gradle.kts b/api/build.gradle.kts
@@ -15,6 +15,7 @@ dependencies {
         exclude("junit", "junit")
     }
     testImplementation(libs.slf4j.simple)
+    testImplementation(project(":testing"))
     testRuntimeOnly(project(":engines:pytorch:pytorch-model-zoo"))
     testRuntimeOnly(project(":engines:pytorch:pytorch-jni"))
 }

diff --git a/api/src/main/java/ai/djl/util/TarUtils.java b/api/src/main/java/ai/djl/util/TarUtils.java
@@ -48,10 +48,8 @@ public static void untar(InputStream is, Path dir, boolean gzip) throws IOExcept
         try (TarArchiveInputStream tis = new TarArchiveInputStream(bis)) {
             TarArchiveEntry entry;
             while ((entry = tis.getNextEntry()) != null) {
-                String entryName = ZipUtils.removeLeadingFileSeparator(entry.getName());
-                if (entryName.contains("..")) {
-                    throw new IOException("Malicious zip entry: " + entryName);
-                }
+                String entryName = entry.getName();
+                ZipUtils.validateArchiveEntry(entryName, dir);
                 Path file = dir.resolve(entryName).toAbsolutePath();
                 if (entry.isDirectory()) {
                     Files.createDirectories(file);

diff --git a/api/src/main/java/ai/djl/util/ZipUtils.java b/api/src/main/java/ai/djl/util/ZipUtils.java
@@ -52,12 +52,10 @@ public static void unzip(InputStream is, Path dest) throws IOException {
         ZipEntry entry;
         Set<String> set = new HashSet<>();
         while ((entry = zis.getNextEntry()) != null) {
-            String name = removeLeadingFileSeparator(entry.getName());
-            if (name.contains("..")) {
-                throw new IOException("Malicious zip entry: " + name);
-            }
-            set.add(name);
-            Path file = dest.resolve(name).toAbsolutePath();
+            String entryName = entry.getName();
+            validateArchiveEntry(entry.getName(), dest);
+            set.add(entryName);
+            Path file = dest.resolve(entryName).toAbsolutePath();
             if (entry.isDirectory()) {
                 Files.createDirectories(file);
             } else {
@@ -121,14 +119,18 @@ private static void addToZip(Path root, Path file, ZipOutputStream zos) throws I
         }
     }
 
-    static String removeLeadingFileSeparator(String name) {
-        int index = 0;
-        for (; index < name.length(); index++) {
-            if (name.charAt(index) != File.separatorChar) {
-                break;
-            }
+    static void validateArchiveEntry(String name, Path destination) throws IOException {
+        if (name.contains("..")) {
+            throw new IOException("Invalid archive entry, contains traversal elements: " + name);
+        }
+        Path expectedOutputPath = destination.resolve(name).toAbsolutePath().normalize();
+        if (!expectedOutputPath.startsWith(destination.normalize())) {
+            throw new IOException(
+                    "Bad archive entry "
+                            + name
+                            + ". Attempted write outside destination "
+                            + destination);
         }
-        return name.substring(index);
     }
 
     private static final class ValidationInputStream extends FilterInputStream {

diff --git a/api/src/test/java/ai/djl/util/ZipUtilsTest.java b/api/src/test/java/ai/djl/util/ZipUtilsTest.java
@@ -47,15 +47,40 @@ public void testEmptyZipFile() throws IOException {
 
     @Test
     public void testOffendingTar() throws IOException {
-        Path path = Paths.get("src/test/resources/offending.tar");
+        String[] offendingTars =
+                new String[] {
+                    "src/test/resources/linux-created-offender-traversal-elements.tar",
+                    "src/test/resources/windows-created-offender-traversal-elements.tar",
+                    "src/test/resources/linux-created-offender-root.tar",
+                    "src/test/resources/windows-created-offender-root.tar",
+                };
         Path output = Paths.get("build/output");
-        Path file = output.resolve("tmp/empty.txt");
-        Utils.deleteQuietly(file);
         Files.createDirectories(output);
-        try (InputStream is = Files.newInputStream(path)) {
-            TarUtils.untar(is, output, false);
+        for (String offendingTar : offendingTars) {
+            Path tarPath = Paths.get(offendingTar);
+            try (InputStream is = Files.newInputStream(tarPath)) {
+                Assert.assertThrows(() -> TarUtils.untar(is, output, false));
+            }
+        }
+    }
+
+    @Test
+    public void testOffendingZip() throws IOException {
+        String[] offendingTars =
+                new String[] {
+                    "src/test/resources/linux-created-offender-traversal-elements.zip",
+                    "src/test/resources/windows-created-offender-traversal-elements.zip",
+                    "src/test/resources/linux-created-offender-root.zip",
+                    "src/test/resources/windows-created-offender-root.zip",
+                };
+        Path output = Paths.get("build/output");
+        Files.createDirectories(output);
+        for (String offendingTar : offendingTars) {
+            Path tarPath = Paths.get(offendingTar);
+            try (InputStream is = Files.newInputStream(tarPath)) {
+                Assert.assertThrows(() -> ZipUtils.unzip(is, output));
+            }
         }
-        Assert.assertTrue(Files.exists(file));
     }
 
     @Test

diff --git a/api/src/test/resources/linux-created-offender-root.tar b/api/src/test/resources/linux-created-offender-root.tar
diff --git a/api/src/test/resources/linux-created-offender-traversal-elements.tar b/api/src/test/resources/linux-created-offender-traversal-elements.tar
diff --git a/api/src/test/resources/windows-created-offender-root.tar b/api/src/test/resources/windows-created-offender-root.tar
diff --git a/api/src/test/resources/windows-created-offender-traversal-elements.tar b/api/src/test/resources/windows-created-offender-traversal-elements.tar
diff --git a/integration/src/test/java/ai/djl/integration/IntegrationTests.java b/integration/src/test/java/ai/djl/integration/IntegrationTests.java
@@ -38,7 +38,7 @@ public void runIntegrationTests() {
             // TODO: windows CPU build is having OOM issue if 3 engines are loaded and running tests
             // together
             if (System.getProperty("os.name").startsWith("Win")) {
-                engines.add("MXNet");
+                engines.add("PyTorch");
             } else if ("aarch64".equals(System.getProperty("os.arch"))) {
                 engines.add("PyTorch");
             } else {

diff --git a/testing/src/main/java/ai/djl/testing/TestRequirements.java b/testing/src/main/java/ai/djl/testing/TestRequirements.java
@@ -96,6 +96,13 @@ public static void notWindows() {
         }
     }
 
+    /** Requires that the test runs on windows, not OSX or linux. */
+    public static void windows() {
+        if (!System.getProperty("os.name").toLowerCase().startsWith("win")) {
+            throw new SkipException("This test requires windows");
+        }
+    }
+
     /** Requires that the test runs on x86_64 arch. */
     public static void notArm() {
         if ("aarch64".equals(System.getProperty("os.arch"))) {