Skip to content

Commit

Permalink
[api] fix issue in Tar/Zip Utils that resulted in incorrect artifact …
Browse files Browse the repository at this point in the history
…extraction
  • Loading branch information
siddvenk committed Nov 27, 2024
1 parent 73f7cb0 commit d543ffe
Show file tree
Hide file tree
Showing 7 changed files with 40 additions and 22 deletions.
1 change: 1 addition & 0 deletions api/build.gradle.kts
Original file line number Diff line number Diff line change
Expand Up @@ -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"))
}
Expand Down
6 changes: 2 additions & 4 deletions api/src/main/java/ai/djl/util/TarUtils.java
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
28 changes: 15 additions & 13 deletions api/src/main/java/ai/djl/util/ZipUtils.java
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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("Bad archive entry, traversal element is not allowed: " + name);
}
Path expectedOutputPath = destination.resolve(name).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 {
Expand Down
18 changes: 14 additions & 4 deletions api/src/test/java/ai/djl/util/ZipUtilsTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
*/
package ai.djl.util;

import ai.djl.testing.TestRequirements;

import org.apache.commons.compress.archivers.zip.Zip64Mode;
import org.apache.commons.compress.archivers.zip.ZipArchiveEntry;
import org.apache.commons.compress.archivers.zip.ZipArchiveOutputStream;
Expand Down Expand Up @@ -49,13 +51,21 @@ public void testEmptyZipFile() throws IOException {
public void testOffendingTar() throws IOException {
Path path = Paths.get("src/test/resources/offending.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);
Assert.assertThrows(() -> TarUtils.untar(is, output, false));
}
}

@Test
public void testLinuxCreatedWindowsUsedOffendingTar() throws IOException {
TestRequirements.windows();
Path tarPath = Paths.get("src/test/resources/linux_create_windows_use.tar");
Path output = Paths.get("build/output");
Files.createDirectories(output);
try (InputStream is = Files.newInputStream(tarPath)) {
Assert.assertThrows(() -> TarUtils.untar(is, output, false));
}
Assert.assertTrue(Files.exists(file));
}

@Test
Expand Down
Binary file not shown.
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
7 changes: 7 additions & 0 deletions testing/src/main/java/ai/djl/testing/TestRequirements.java
Original file line number Diff line number Diff line change
Expand Up @@ -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"))) {
Expand Down

0 comments on commit d543ffe

Please sign in to comment.