Security fix to prevent allowed_urls Filter Bypass (#3082)

* allowed url test fix * allowed url test fix * changing test name * fix format --------- Co-authored-by: Mark Saroufim <marksaroufim@fb.com>
pytorch · Apr 11, 2024 · cdba0fd · cdba0fd
1 parent 440fca0
commit cdba0fd
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 6 deletions.
diff --git a/frontend/archive/src/main/java/org/pytorch/serve/archive/model/ModelArchive.java b/frontend/archive/src/main/java/org/pytorch/serve/archive/model/ModelArchive.java
@@ -56,17 +56,18 @@ public static ModelArchive downloadModel(
 
         String marFileName = ArchiveUtils.getFilenameFromUrl(url);
         File modelLocation = new File(modelStore, marFileName);
+
+        if (url.contains("..")) {
+            throw new ModelNotFoundException("Relative path is not allowed in url: " + url);
+        }
+
         try {
             ArchiveUtils.downloadArchive(
                     allowedUrls, modelLocation, marFileName, url, s3SseKmsEnabled);
         } catch (InvalidArchiveURLException e) {
             throw new ModelNotFoundException(e.getMessage()); // NOPMD
         }
 
-        if (url.contains("..")) {
-            throw new ModelNotFoundException("Relative path is not allowed in url: " + url);
-        }
-
         if (modelLocation.isFile()) {
             try (InputStream is = Files.newInputStream(modelLocation.toPath())) {
                 File unzipDir;

diff --git a/frontend/archive/src/test/java/org/pytorch/serve/archive/model/ModelArchiveTest.java b/frontend/archive/src/test/java/org/pytorch/serve/archive/model/ModelArchiveTest.java
@@ -158,11 +158,11 @@ public void archiveTest() throws ModelException, IOException, DownloadArchiveExc
 
     @Test(expectedExceptions = DownloadArchiveException.class)
     public void testMalformedURL() throws ModelException, IOException, DownloadArchiveException {
-        String modelStore = "src/test/resources/models";
+        String modelStore = "src/test/resources";
         ModelArchive.downloadModel(
                 ALLOWED_URLS_LIST,
                 modelStore,
-                "https://../model-server/models/squeezenet_v1.1/squeezenet_v1.1.mod");
+                "https://model-server/models/squeezenet_v1.1/squeezenet_v1.1.mod");
     }
 
     @Test(
@@ -174,6 +174,37 @@ public void testRelativePath() throws ModelException, IOException, DownloadArchi
         ModelArchive.downloadModel(ALLOWED_URLS_LIST, modelStore, "../mnist.mar");
     }
 
+    @Test
+    public void testRelativePathFileExists()
+            throws ModelException, IOException, DownloadArchiveException {
+        String modelStore = "src/test/resources/models";
+        String curDir = System.getProperty("user.dir");
+        File curDirFile = new File(curDir);
+        String parent = curDirFile.getParent();
+
+        // Setup: This test needs mar file in local path. Copying mnist.mar from model folder.
+        String source = modelStore + "/mnist.mar";
+        String destination = parent + "/archive/mnist1.mar";
+        File sourceFile = new File(source);
+        File destinationFile = new File(destination);
+        FileUtils.copyFile(sourceFile, destinationFile);
+
+        String fileUrl = "file:///" + parent + "/archive/../archive/mnist1.mar";
+        try {
+            ModelArchive archive =
+                    ModelArchive.downloadModel(ALLOWED_URLS_LIST, modelStore, fileUrl);
+        } catch (ModelNotFoundException e) {
+            String expectedMessagePattern = "Relative path is not allowed in url: " + fileUrl;
+            Assert.assertTrue(
+                    e.getMessage().matches(expectedMessagePattern),
+                    "Exception message does not match the expected pattern.");
+        }
+
+        // Verify the file doesn't exist
+        File modelLocation = new File(modelStore + "/mnist1.mar");
+        Assert.assertFalse(modelLocation.exists());
+    }
+
     @Test(
             expectedExceptions = ModelNotFoundException.class,
             expectedExceptionsMessageRegExp = "Model store has not been configured\\.")