Uncontrolled input fix for ModelManager (#2635)

* Uncontrolled input fix for ModelManager * update * update
pytorch · Sep 29, 2023 · 6e5b8cd · 6e5b8cd
1 parent eaad8a4
commit 6e5b8cd
Showing 1 changed file with 23 additions and 13 deletions.
diff --git a/frontend/server/src/main/java/org/pytorch/serve/wlm/ModelManager.java b/frontend/server/src/main/java/org/pytorch/serve/wlm/ModelManager.java
@@ -9,6 +9,7 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -217,26 +218,35 @@ private void setupModelDependencies(Model model)
             if (Files.isSymbolicLink(dependencyPath.toPath())) {
                 dependencyPath = dependencyPath.getParentFile();
             }
-            String packageInstallCommand =
-                    pythonRuntime
-                            + " -m pip install -U -t "
-                            + dependencyPath.getAbsolutePath()
-                            + " -r "
-                            + requirementsFilePath; // NOPMD
+
+            List<String> commandParts = new ArrayList<>();
+
+            commandParts.add(pythonRuntime);
+            commandParts.add("-m");
+            commandParts.add("pip");
+            commandParts.add("install");
+            commandParts.add("-U");
+            commandParts.add("-t");
+            commandParts.add(dependencyPath.getAbsolutePath());
+            commandParts.add("-r");
+            commandParts.add(requirementsFilePath.toString());
 
             String[] envp =
                     EnvironmentUtils.getEnvString(
                             configManager.getModelServerHome(),
                             model.getModelDir().getAbsolutePath(),
                             null);
 
-            Process process =
-                    Runtime.getRuntime()
-                            .exec(
-                                    packageInstallCommand,
-                                    envp,
-                                    model.getModelDir().getAbsoluteFile());
-
+            ProcessBuilder processBuilder = new ProcessBuilder(commandParts);
+            processBuilder.directory(model.getModelDir().getAbsoluteFile());
+            Map<String, String> environment = processBuilder.environment();
+            for (String envVar : envp) {
+                String[] parts = envVar.split("=", 2);
+                if (parts.length == 2) {
+                    environment.put(parts[0], parts[1]);
+                }
+            }
+            Process process = processBuilder.start();
             int exitCode = process.waitFor();
 
             if (exitCode != 0) {