Python: type-track through dict-updates

RasmusWL · RasmusWL · commit dac2b57bb029 · 2024-03-14T10:42:38.000+01:00
diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPrivate.qll b/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPrivate.qll
@@ -648,6 +648,8 @@ predicate storeStepCommon(Node nodeFrom, ContentSet c, Node nodeTo) {
   tupleStoreStep(nodeFrom, c, nodeTo)
   or
   dictStoreStep(nodeFrom, c, nodeTo)
+  or
+  moreDictStoreSteps(nodeFrom, c, nodeTo)
 }
 
 /**
@@ -661,8 +663,6 @@ predicate storeStep(Node nodeFrom, ContentSet c, Node nodeTo) {
   or
   setStoreStep(nodeFrom, c, nodeTo)
   or
-  moreDictStoreSteps(nodeFrom, c, nodeTo)
-  or
   comprehensionStoreStep(nodeFrom, c, nodeTo)
   or
   iterableUnpackingStoreStep(nodeFrom, c, nodeTo)
diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/TypeTrackingImpl.qll b/python/ql/lib/semmle/python/dataflow/new/internal/TypeTrackingImpl.qll
@@ -175,7 +175,18 @@ module TypeTrackingInput implements Shared::TypeTrackingInput {
       nodeTo = a.getObject()
     )
     or
-    DataFlowPrivate::storeStepCommon(nodeFrom, content, nodeTo)
+    // type-tracking doesn't really handle PostUpdateNodes, so for some assignment steps
+    // like `my_dict["foo"] = foo` the data-flow step targets the PostUpdateNode for
+    // `my_dict`, where we want to translate that into a type-tracking step that targets
+    // the normal/non-PostUpdateNode for `my_dict`.
+    exists(DataFlowPublic::Node storeTarget |
+      DataFlowPrivate::storeStepCommon(nodeFrom, content, storeTarget)
+    |
+      not storeTarget instanceof DataFlowPrivate::SyntheticPostUpdateNode and
+      nodeTo = storeTarget
+      or
+      nodeTo = storeTarget.(DataFlowPrivate::SyntheticPostUpdateNode).getPreUpdateNode()
+    )
     or
     TypeTrackerSummaryFlow::basicStoreStep(nodeFrom, nodeTo, content)
   }
diff --git a/python/ql/test/experimental/library-tests/CallGraph/InlineCallGraphTest.expected b/python/ql/test/experimental/library-tests/CallGraph/InlineCallGraphTest.expected
@@ -39,6 +39,7 @@ typeTracker_found_pointsTo_notFound
 | code/class_super.py:108:1:108:8 | ControlFlowNode for Attribute() | Z.foo |
 | code/def_in_function.py:22:5:22:11 | ControlFlowNode for Attribute() | test.A.foo |
 | code/func_ref_in_content.py:29:1:29:4 | ControlFlowNode for f4() | func |
+| code/func_ref_in_content.py:40:1:40:4 | ControlFlowNode for f5() | func |
 | code/isinstance.py:9:13:9:22 | ControlFlowNode for Attribute() | A.foo |
 | code/isinstance.py:9:13:9:22 | ControlFlowNode for Attribute() | ASub.foo |
 | code/isinstance.py:14:13:14:22 | ControlFlowNode for Attribute() | A.foo |
diff --git a/python/ql/test/experimental/library-tests/CallGraph/code/func_ref_in_content.py b/python/ql/test/experimental/library-tests/CallGraph/code/func_ref_in_content.py
@@ -37,7 +37,7 @@ def return_func_in_dict_update():
 dct2 = return_func_in_dict_update() # $ pt,tt=return_func_in_dict_update
 
 f5 = dct2['func']
-f5() # $ MISSING: tt=func
+f5() # $ tt=func
 
 
 def return_func_in_list():
