Python: Proper models of json.loads and json.dumps

Author: RasmusWL

python/ql/src/semmle/python/dataflow/new/internal/TaintTrackingPrivate.qll

@@ -31,8 +31,6 @@ predicate localAdditionalTaintStep(DataFlow::Node nodeFrom, DataFlow::Node nodeT
   or
   stringManipulation(nodeFrom, nodeTo)
   or
-  jsonStep(nodeFrom, nodeTo)
-  or
   containerStep(nodeFrom, nodeTo)
   or
   copyStep(nodeFrom, nodeTo)
@@ -135,16 +133,6 @@ predicate stringManipulation(DataFlow::CfgNode nodeFrom, DataFlow::CfgNode nodeT
   // TODO: Handle functions in https://docs.python.org/3/library/binascii.html
 }
 
-/**
- * Holds if taint can flow from `nodeFrom` to `nodeTo` with a step related to JSON encoding/decoding.
- */
-predicate jsonStep(DataFlow::CfgNode nodeFrom, DataFlow::CfgNode nodeTo) {
-  exists(CallNode call | call = nodeTo.getNode() |
-    call.getFunction().(AttrNode).getObject(["load", "loads", "dumps"]).(NameNode).getId() = "json" and
-    call.getArg(0) = nodeFrom.getNode()
-  )
-}
-
 /**
  * Holds if taint can flow from `nodeFrom` to `nodeTo` with a step related to containers
  * (lists/sets/dictionaries): literals, constructor invocation, methods. Note that this

python/ql/src/semmle/python/frameworks/Stdlib.qll

@@ -945,6 +945,93 @@ private module Stdlib {
   private DataFlow::Node io_attr(string attr_name) {
     result = io_attr(DataFlow::TypeTracker::end(), attr_name)
   }
+
+  // ---------------------------------------------------------------------------
+  // json
+  // ---------------------------------------------------------------------------
+  /** Gets a reference to the `json` module. */
+  private DataFlow::Node json(DataFlow::TypeTracker t) {
+    t.start() and
+    result = DataFlow::importNode("json")
+    or
+    exists(DataFlow::TypeTracker t2 | result = json(t2).track(t2, t))
+  }
+
+  /** Gets a reference to the `json` module. */
+  DataFlow::Node json() { result = json(DataFlow::TypeTracker::end()) }
+
+  /**
+   * Gets a reference to the attribute `attr_name` of the `json` module.
+   * WARNING: Only holds for a few predefined attributes.
+   */
+  private DataFlow::Node json_attr(DataFlow::TypeTracker t, string attr_name) {
+    attr_name in ["loads", "dumps"] and
+    (
+      t.start() and
+      result = DataFlow::importNode("json" + "." + attr_name)
+      or
+      t.startInAttr(attr_name) and
+      result = json()
+    )
+    or
+    // Due to bad performance when using normal setup with `json_attr(t2, attr_name).track(t2, t)`
+    // we have inlined that code and forced a join
+    exists(DataFlow::TypeTracker t2 |
+      exists(DataFlow::StepSummary summary |
+        json_attr_first_join(t2, attr_name, result, summary) and
+        t = t2.append(summary)
+      )
+    )
+  }
+
+  pragma[nomagic]
+  private predicate json_attr_first_join(
+    DataFlow::TypeTracker t2, string attr_name, DataFlow::Node res, DataFlow::StepSummary summary
+  ) {
+    DataFlow::StepSummary::step(json_attr(t2, attr_name), res, summary)
+  }
+
+  /**
+   * Gets a reference to the attribute `attr_name` of the `json` module.
+   * WARNING: Only holds for a few predefined attributes.
+   */
+  private DataFlow::Node json_attr(string attr_name) {
+    result = json_attr(DataFlow::TypeTracker::end(), attr_name)
+  }
+
+  /**
+   * A call to `json.loads`
+   * See https://docs.python.org/3/library/json.html#json.loads
+   */
+  private class JsonLoadsCall extends Decoding::Range, DataFlow::CfgNode {
+    override CallNode node;
+
+    JsonLoadsCall() { node.getFunction() = json_attr("loads").asCfgNode() }
+
+    override predicate mayExecuteInput() { none() }
+
+    override DataFlow::Node getAnInput() { result.asCfgNode() = node.getArg(0) }
+
+    override DataFlow::Node getOutput() { result = this }
+
+    override string getFormat() { result = "JSON" }
+  }
+
+  /**
+   * A call to `json.dumps`
+   * See https://docs.python.org/3/library/json.html#json.dumps
+   */
+  private class JsonDumpsCall extends Encoding::Range, DataFlow::CfgNode {
+    override CallNode node;
+
+    JsonDumpsCall() { node.getFunction() = json_attr("dumps").asCfgNode() }
+
+    override DataFlow::Node getAnInput() { result.asCfgNode() = node.getArg(0) }
+
+    override DataFlow::Node getOutput() { result = this }
+
+    override string getFormat() { result = "JSON" }
+  }
 }
 
 // ---------------------------------------------------------------------------

python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/TestTaint.expected

@@ -68,9 +68,9 @@
 | test_json.py:27 | ok   | test | json.loads(..) |
 | test_json.py:34 | fail | test | tainted_filelike |
 | test_json.py:35 | fail | test | json.load(..) |
-| test_json.py:48 | fail | non_syntacical | dumps(..) |
-| test_json.py:49 | fail | non_syntacical | dumps_alias(..) |
-| test_json.py:50 | fail | non_syntacical | loads(..) |
+| test_json.py:48 | ok   | non_syntacical | dumps(..) |
+| test_json.py:49 | ok   | non_syntacical | dumps_alias(..) |
+| test_json.py:50 | ok   | non_syntacical | loads(..) |
 | test_json.py:57 | fail | non_syntacical | tainted_filelike |
 | test_json.py:58 | fail | non_syntacical | load(..) |
 | test_string.py:25 | ok   | str_operations | ts |