Merge pull request #15781 from RasmusWL/dict-update

Python: Fix missing DictionaryElementContents

python/ql/lib/change-notes/2024-03-01-dict-update-content.md

@@ -0,0 +1,4 @@
+---
+category: minorAnalysis
+---
+* Fixed missing flow for dictionary updates (`d[<key>] = ...`) when `<key>` is a string constant not used in dictionary literals or as name of keyword-argument.

python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPrivate.qll

@@ -809,6 +809,8 @@ predicate dictStoreStep(CfgNode nodeFrom, DictionaryElementContent c, Node nodeT
  * TODO: Once TaintTracking no longer uses `dictStoreStep`, unify the two predicates.
  */
 private predicate moreDictStoreSteps(CfgNode nodeFrom, DictionaryElementContent c, Node nodeTo) {
+  // NOTE: It's important to add logic to the newtype definition of
+  // DictionaryElementContent if you add new cases here.
   exists(SubscriptNode subscript |
     nodeTo.(PostUpdateNode).getPreUpdateNode().asCfgNode() = subscript.getObject() and
     nodeFrom.asCfgNode() = subscript.(DefinitionNode).getValue() and

python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPublic.qll

@@ -605,9 +605,19 @@ newtype TContent =
   } or
   /** An element of a dictionary under a specific key. */
   TDictionaryElementContent(string key) {
-    key = any(KeyValuePair kvp).getKey().(StrConst).getS()
+    // {"key": ...}
+    key = any(KeyValuePair kvp).getKey().(StrConst).getText()
     or
+    // func(key=...)
     key = any(Keyword kw).getArg()
+    or
+    // d["key"] = ...
+    key = any(SubscriptNode sub | sub.isStore() | sub.getIndex().getNode().(StrConst).getText())
+    or
+    // d.setdefault("key", ...)
+    exists(CallNode call | call.getFunction().(AttrNode).getName() = "setdefault" |
+      key = call.getArg(0).getNode().(StrConst).getText()
+    )
   } or
   /** An element of a dictionary under any key. */
   TDictionaryElementAnyContent() or

python/ql/test/experimental/dataflow/fieldflow/test_dict.py

@@ -45,6 +45,15 @@ def test_dict_update():
     SINK(d["key"])  # $ flow="SOURCE, l:-1 -> d['key']"
     SINK(d.get("key"))  # $ flow="SOURCE, l:-2 -> d.get(..)"
 
+
+def test_dict_update_fresh_key():
+    # we had a regression where we did not create a dictionary element content
+    # for keys used in "inline update" like this
+    d = {}
+    d["fresh_key"] = SOURCE
+    SINK(d["fresh_key"])  # $ flow="SOURCE, l:-1 -> d['fresh_key']"
+
+
 @expects(3) # $ unresolved_call=expects(..) unresolved_call=expects(..)(..)
 def test_dict_setdefault():
     d = {}