From 709649e9df4cc46a640a23156c7e65c86ea62080 Mon Sep 17 00:00:00 2001
From: Ed Minnix <egregius313@github.com>
Date: Wed, 13 Dec 2023 12:54:53 -0500
Subject: [PATCH] Model `replace` and `putIfAbsent`

---
 .../security/TaintedEnvironmentVariableQuery.qll   | 14 ++++++++------
 .../security/CWE-078/ExecRelative.expected         |  2 +-
 .../security/CWE-078/TaintedEnvironment.java       | 11 +++++++++++
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/java/ql/lib/semmle/code/java/security/TaintedEnvironmentVariableQuery.qll b/java/ql/lib/semmle/code/java/security/TaintedEnvironmentVariableQuery.qll
index f76e92a49bd..d590161b20a 100644
--- a/java/ql/lib/semmle/code/java/security/TaintedEnvironmentVariableQuery.qll
+++ b/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())
     )
   }
diff --git a/java/ql/test/query-tests/security/CWE-078/ExecRelative.expected b/java/ql/test/query-tests/security/CWE-078/ExecRelative.expected
index be2194d5046..e5b842b4395 100644
--- a/java/ql/test/query-tests/security/CWE-078/ExecRelative.expected
+++ b/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. |
diff --git a/java/ql/test/query-tests/security/CWE-078/TaintedEnvironment.java b/java/ql/test/query-tests/security/CWE-078/TaintedEnvironment.java
index b7b1cc05f83..df4812eed2c 100644
--- a/java/ql/test/query-tests/security/CWE-078/TaintedEnvironment.java
+++ b/java/ql/test/query-tests/security/CWE-078/TaintedEnvironment.java
@@ -15,6 +15,17 @@ public class TaintedEnvironment {
 
         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