From f669a4f3bf16c34dffdb6b4d17e3ba85e8b5469e Mon Sep 17 00:00:00 2001
From: Rasmus Lerchedahl Petersen <yoff@github.com>
Date: Mon, 11 Nov 2024 18:43:21 +0100
Subject: [PATCH] Python: Make sure all imprecise taint bubbles up

---
 python/ql/lib/semmle/python/frameworks/Stdlib.qll     | 11 +++++++++--
 .../defaultAdditionalTaintStep/test_collections.py    | 10 ++++------
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/python/ql/lib/semmle/python/frameworks/Stdlib.qll b/python/ql/lib/semmle/python/frameworks/Stdlib.qll
index 5d3b994880a..7e5d0d8ab06 100644
--- a/python/ql/lib/semmle/python/frameworks/Stdlib.qll
+++ b/python/ql/lib/semmle/python/frameworks/Stdlib.qll
@@ -4244,8 +4244,15 @@ module StdlibPrivate {
         )
         // TODO: Once we have DictKeyContent, we need to transform that into ListElementContent
       ) and
-      output = "ReturnValue.ListElement" and
-      preservesValue = true
+      (
+        //Element content is mutated into list element content
+        output = "ReturnValue.ListElement" and
+        preservesValue = true
+        or
+        // Since list content is imprecise, we also taint the list.
+        output = "ReturnValue" and
+        preservesValue = false
+      )
       or
       input = "Argument[0]" and
       output = "ReturnValue" and
diff --git a/python/ql/test/library-tests/dataflow/tainttracking/defaultAdditionalTaintStep/test_collections.py b/python/ql/test/library-tests/dataflow/tainttracking/defaultAdditionalTaintStep/test_collections.py
index a67438316f4..b9fa1ebffd4 100644
--- a/python/ql/test/library-tests/dataflow/tainttracking/defaultAdditionalTaintStep/test_collections.py
+++ b/python/ql/test/library-tests/dataflow/tainttracking/defaultAdditionalTaintStep/test_collections.py
@@ -27,14 +27,11 @@ def test_construction():
         tainted_dict, # $ tainted
     )
 
-    # There are no implicit reads for list content as it is imprecise
-    # Therefore, list content stemming from precise content does not end up on the list itself. 
     ensure_tainted(
         list(tainted_list), # $ tainted
-        list(tainted_tuple)[0], # $ tainted
+        list(tainted_tuple), # $ tainted
         list(tainted_set), # $ tainted
-        list(tainted_dict.values())[0], # $ tainted
-        list(tainted_dict.items())[0], # $ tainted
+        list(tainted_dict.values()), # $ tainted
 
         tuple(tainted_list), # $ tainted
         set(tainted_list), # $ tainted
@@ -46,7 +43,8 @@ def test_construction():
     )
 
     ensure_not_tainted(
-        dict(k = tainted_string)["k1"]
+        dict(k = tainted_string)["k1"],
+        list(tainted_dict.items()),
     )