Model replace and putIfAbsent

Author: egregius313

java/ql/lib/semmle/code/java/security/TaintedEnvironmentVariableQuery.qll

@@ -5,10 +5,10 @@ private import semmle.code.java.dataflow.FlowSources
 private import semmle.code.java.Maps
 private import semmle.code.java.JDK
 
-private class MapPutOrRemove extends MethodCall {
-  MapPutOrRemove() {
-    this.getMethod() instanceof MapMutator and
-    this.getMethod().getName().matches(["put%", "remove"])
+private class MapUpdateWithKeyOrValue extends MethodCall {
+  MapUpdateWithKeyOrValue() {
+    this.getMethod() instanceof MapMethod and
+    this.getMethod().getName().matches(["put%", "remove", "replace"])
   }
 }
 
@@ -19,7 +19,9 @@ private module ProcessBuilderEnvironmentConfig implements DataFlow::ConfigSig {
     )
   }
 
-  predicate isSink(DataFlow::Node sink) { sink.asExpr() = any(MapPutOrRemove mm).getQualifier() }
+  predicate isSink(DataFlow::Node sink) {
+    sink.asExpr() = any(MapUpdateWithKeyOrValue mm).getQualifier()
+  }
 }
 
 private module ProcessBuilderEnvironmentFlow = DataFlow::Global<ProcessBuilderEnvironmentConfig>;
@@ -41,7 +43,7 @@ module ExecTaintedEnvironmentConfig implements DataFlow::ConfigSig {
     sinkNode(sink, "environment-injection")
     or
     // sink is a key or value added to a `ProcessBuilder::environment` map.
-    exists(MapPutOrRemove mm | mm.getAnArgument() = sink.asExpr() |
+    exists(MapUpdateWithKeyOrValue mm | mm.getAnArgument() = sink.asExpr() |
       ProcessBuilderEnvironmentFlow::flowToExpr(mm.getQualifier())
     )
   }

java/ql/test/query-tests/security/CWE-078/ExecRelative.expected

@@ -1,2 +1,2 @@
-| TaintedEnvironment.java:28:35:28:55 | new String[] | Command with a relative path 'ls' is executed. |
+| TaintedEnvironment.java:39:35:39:55 | new String[] | Command with a relative path 'ls' is executed. |
 | Test.java:50:46:50:49 | "ls" | Command with a relative path 'ls' is executed. |

java/ql/test/query-tests/security/CWE-078/TaintedEnvironment.java

@@ -15,6 +15,17 @@ public void buildProcess() throws java.io.IOException {
 
         pb.environment().put(s, "foo"); // $hasTaintFlow
 
+        Map<String, String> extra = Map.of("USER", s);
+
+        pb.environment().putAll(extra); // $hasTaintFlow
+
+        pb.environment().putIfAbsent("foo", s); // $hasTaintFlow
+        pb.environment().putIfAbsent(s, "foo"); // $hasTaintFlow
+
+        pb.environment().replace("foo", s); // $hasTaintFlow
+        pb.environment().replace(s, "foo"); // $hasTaintFlow
+        pb.environment().replace("foo", "bar", s); // $hasTaintFlow
+
         Map<String, String> env = pb.environment();
 
         env.put("foo", s); // $hasTaintFlow