From cdba0fd449c2fd23dcf37c54c0784035541d5114 Mon Sep 17 00:00:00 2001
From: udaij12 <32673964+udaij12@users.noreply.github.com>
Date: Thu, 11 Apr 2024 10:37:21 -0700
Subject: [PATCH] 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>
---
 .../serve/archive/model/ModelArchive.java     |  9 ++---
 .../serve/archive/model/ModelArchiveTest.java | 35 +++++++++++++++++--
 2 files changed, 38 insertions(+), 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
index fe8b9ee392..9d26ff765f 100644
--- 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,6 +56,11 @@ 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);
@@ -63,10 +68,6 @@ public static ModelArchive downloadModel(
             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
index 4770fff384..1079007ca0 100644
--- 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\\.")