[secure-mode] Allowlist option.pythonExecutable, but only for recogni…

…zed paths (#2715)
deepjavalibrary · Feb 4, 2025 · 509b7bd · 509b7bd
1 parent b437a9e
commit 509b7bd
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 1 deletion.
diff --git a/plugins/secure-mode/src/main/java/ai/djl/serving/plugins/securemode/SecureModeAllowList.java b/plugins/secure-mode/src/main/java/ai/djl/serving/plugins/securemode/SecureModeAllowList.java
@@ -107,5 +107,12 @@ interface SecureModeAllowList {
                     "option.enable_prefix_caching",
                     "option.disable_sliding_window",
                     "option.enable_streaming",
-                    "option.tgi_compat");
+                    "option.tgi_compat",
+                    "option.pythonExecutable");
+
+    public static final Set<String> PYTHON_EXECUTABLE_ALLOWLIST =
+            Set.of(
+                    "/opt/djl/lmi_dist_venv/bin/python",
+                    "/opt/djl/vllm_venv/bin/python",
+                    "/usr/bin/python3");
 }
diff --git a/plugins/secure-mode/src/main/java/ai/djl/serving/plugins/securemode/SecureModeUtils.java b/plugins/secure-mode/src/main/java/ai/djl/serving/plugins/securemode/SecureModeUtils.java
@@ -150,6 +150,14 @@ private static void checkOptions(ModelInfo<?, ?> modelInfo, Set<String> security
                         "Installing additional dependencies is prohibited in Secure Mode.");
             }
         }
+        String pythonExecutable = prop.getProperty("option.pythonExecutable");
+        if (pythonExecutable != null
+                && !SecureModeAllowList.PYTHON_EXECUTABLE_ALLOWLIST.contains(pythonExecutable)) {
+            throw new IllegalConfigurationException(
+                    "Custom Python executable path is prohibited in Secure Mode. "
+                            + "Only the following paths are allowed: "
+                            + SecureModeAllowList.PYTHON_EXECUTABLE_ALLOWLIST);
+        }
     }
 
     /**

diff --git a/...ins/secure-mode/src/test/java/ai/djl/serving/plugins/securemode/SecureModePluginTest.java b/...ins/secure-mode/src/test/java/ai/djl/serving/plugins/securemode/SecureModePluginTest.java
@@ -272,6 +272,22 @@ void testPropertiesAllowlist() throws IOException, ModelException {
                 "foo", TEST_MODEL_DIR.resolve("serving.properties"), "option.not_allowlisted=foo");
     }
 
+    @Test(expectedExceptions = IllegalConfigurationException.class)
+    void testInvalidPythonExecutablePath() throws IOException, ModelException {
+        mockSecurityEnv(
+                "foo",
+                TEST_MODEL_DIR.resolve("serving.properties"),
+                "option.pythonExecutable=/foo/bar/python3");
+    }
+
+    @Test
+    void testAllowedPythonExecutablePath() throws IOException, ModelException {
+        mockSecurityEnv(
+                "foo",
+                TEST_MODEL_DIR.resolve("serving.properties"),
+                "option.pythonExecutable=/opt/djl/lmi_dist_venv/bin/python");
+    }
+
     private void createFileWithContent(Path file, String content) throws IOException {
         if (Files.exists(file)) {
             return;