Merge pull request #7031 from yoff/python/taint-through-with

RasmusWL · web-flow · commit 1e31416049d6 · 2021-11-09T14:08:07.000+01:00
Python: Taint through `async with`
diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/TaintTrackingPrivate.qll b/python/ql/lib/semmle/python/dataflow/new/internal/TaintTrackingPrivate.qll
@@ -53,6 +53,8 @@ private module Cached {
     DataFlowPrivate::iterableUnpackingStoreStep(nodeFrom, _, nodeTo)
     or
     awaitStep(nodeFrom, nodeTo)
+    or
+    asyncWithStep(nodeFrom, nodeTo)
   }
 }
 
@@ -211,3 +213,24 @@ predicate copyStep(DataFlow::CfgNode nodeFrom, DataFlow::CfgNode nodeTo) {
 predicate awaitStep(DataFlow::Node nodeFrom, DataFlow::Node nodeTo) {
   nodeTo.asExpr().(Await).getValue() = nodeFrom.asExpr()
 }
+
+/**
+ * Holds if taint can flow from `nodeFrom` to `nodeTo` inside an `async with` statement.
+ *
+ * For example in
+ * ```python
+ * async with open("foo") as f:
+ * ```
+ * the variable `f` is tainted if the result of `open("foo")` is tainted.
+ */
+predicate asyncWithStep(DataFlow::Node nodeFrom, DataFlow::Node nodeTo) {
+  exists(With with, ControlFlowNode contextManager, ControlFlowNode var |
+    nodeFrom.(DataFlow::CfgNode).getNode() = contextManager and
+    nodeTo.(DataFlow::EssaNode).getVar().getDefinition().(WithDefinition).getDefiningNode() = var and
+    // see `with_flow` in `python/ql/src/semmle/python/dataflow/Implementation.qll`
+    with.getContextExpr() = contextManager.getNode() and
+    with.getOptionalVars() = var.getNode() and
+    with.isAsync() and
+    contextManager.strictlyDominates(var)
+  )
+}
diff --git a/python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/test_async.py b/python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/test_async.py
@@ -31,7 +31,7 @@ async def test_async_with():
     ctx = AsyncContext()
     taint(ctx)
     async with ctx as tainted:
-        ensure_tainted(tainted) # $ MISSING: tainted
+        ensure_tainted(tainted) # $ tainted
 
 
 class AsyncIter:
@@ -45,7 +45,7 @@ async def test_async_for():
     iter = AsyncIter()
     taint(iter)
     async for tainted in iter:
-        ensure_tainted(tainted) # $ MISSING: tainted
+        ensure_tainted(tainted) # $ tainted
 
 
 
diff --git a/python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/test_for.py b/python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/test_for.py
@@ -0,0 +1,30 @@
+# Add taintlib to PATH so it can be imported during runtime without any hassle
+import sys; import os; sys.path.append(os.path.dirname(os.path.dirname((__file__))))
+from taintlib import *
+
+# This has no runtime impact, but allows autocomplete to work
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from ..taintlib import *
+
+
+# Actual tests
+
+class Iter:
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        raise StopIteration
+
+def test_for():
+    iter = Iter()
+    taint(iter)
+    for tainted in iter:
+        ensure_tainted(tainted) # $ tainted
+
+
+
+# Make tests runable
+
+test_for()
diff --git a/python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/test_with.py b/python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/test_with.py
@@ -0,0 +1,60 @@
+# Add taintlib to PATH so it can be imported during runtime without any hassle
+import sys; import os; sys.path.append(os.path.dirname(os.path.dirname((__file__))))
+from taintlib import *
+
+# This has no runtime impact, but allows autocomplete to work
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from ..taintlib import *
+
+
+# Actual tests
+
+class Context:
+    def __enter__(self):
+        return ""
+
+    def __exit__(self, exc_type, exc, tb):
+        pass
+
+def test_with():
+    ctx = Context()
+    taint(ctx)
+    with ctx as tainted:
+        ensure_tainted(tainted) # $ tainted
+
+class Context_taint:
+    def __enter__(self):
+        return TAINTED_STRING
+
+    def __exit__(self, exc_type, exc, tb):
+        pass
+
+def test_with_taint():
+    ctx = Context_taint()
+    with ctx as tainted:
+        ensure_tainted(tainted) # $ MISSING: tainted
+
+
+class Context_arg:
+    def __init__(self, arg):
+        self.arg = arg
+
+    def __enter__(self):
+        return self.arg
+
+    def __exit__(self, exc_type, exc, tb):
+        pass
+
+def test_with_arg():
+    ctx = Context_arg(TAINTED_STRING)
+    with ctx as tainted:
+        ensure_tainted(tainted) # $ tainted
+
+
+
+# Make tests runable
+
+test_with()
+test_with_taint()
+test_with_arg()
diff --git a/python/ql/test/experimental/meta/InlineTaintTest.qll b/python/ql/test/experimental/meta/InlineTaintTest.qll
@@ -30,24 +30,36 @@ DataFlow::Node shouldNotBeTainted() {
   )
 }
 
-class TestTaintTrackingConfiguration extends TaintTracking::Configuration {
-  TestTaintTrackingConfiguration() { this = "TestTaintTrackingConfiguration" }
+// this module allows the configuration to be imported in other `.ql` files without the
+// top level query predicates of this file coming into scope.
+module Conf {
+  class TestTaintTrackingConfiguration extends TaintTracking::Configuration {
+    TestTaintTrackingConfiguration() { this = "TestTaintTrackingConfiguration" }
 
-  override predicate isSource(DataFlow::Node source) {
-    source.asCfgNode().(NameNode).getId() in [
-        "TAINTED_STRING", "TAINTED_BYTES", "TAINTED_LIST", "TAINTED_DICT"
-      ]
-    or
-    source instanceof RemoteFlowSource
-  }
+    override predicate isSource(DataFlow::Node source) {
+      source.asCfgNode().(NameNode).getId() in [
+          "TAINTED_STRING", "TAINTED_BYTES", "TAINTED_LIST", "TAINTED_DICT"
+        ]
+      or
+      // User defined sources
+      exists(CallNode call |
+        call.getFunction().(NameNode).getId() = "taint" and
+        source.(DataFlow::CfgNode).getNode() = call.getAnArg()
+      )
+      or
+      source instanceof RemoteFlowSource
+    }
 
-  override predicate isSink(DataFlow::Node sink) {
-    sink = shouldBeTainted()
-    or
-    sink = shouldNotBeTainted()
+    override predicate isSink(DataFlow::Node sink) {
+      sink = shouldBeTainted()
+      or
+      sink = shouldNotBeTainted()
+    }
   }
 }
 
+import Conf
+
 class InlineTaintTest extends InlineExpectationsTest {
   InlineTaintTest() { this = "InlineTaintTest" }
 
diff --git a/python/ql/test/experimental/meta/debug/InlineTaintTestPaths.expected b/python/ql/test/experimental/meta/debug/InlineTaintTestPaths.expected
@@ -0,0 +1,4 @@
+edges
+nodes
+subpaths
+#select
diff --git a/python/ql/test/experimental/meta/debug/InlineTaintTestPaths.ql b/python/ql/test/experimental/meta/debug/InlineTaintTestPaths.ql
@@ -0,0 +1,25 @@
+/**
+ * @kind path-problem
+ */
+
+// This query is for debugging InlineTaintTestFailures.
+// The intended usage is
+// 1. load the database of the failing test
+// 2. run this query to see actual paths
+// 3. if necessary, look at partial paths by (un)commenting appropriate lines
+import python
+import semmle.python.dataflow.new.DataFlow
+import experimental.meta.InlineTaintTest::Conf
+// import DataFlow::PartialPathGraph
+import DataFlow::PathGraph
+
+class Conf extends TestTaintTrackingConfiguration {
+  override int explorationLimit() { result = 5 }
+}
+
+// from Conf config, DataFlow::PartialPathNode source, DataFlow::PartialPathNode sink
+// where config.hasPartialFlow(source, sink, _)
+from Conf config, DataFlow::PathNode source, DataFlow::PathNode sink
+where config.hasFlowPath(source, sink)
+select sink.getNode(), source, sink, "This node receives taint from $@.", source.getNode(),
+  "this source"

-Original file line number
+Diff line change
@@ @@ -0,0 +1,4 @@ @@
 +edges
 +nodes
 +subpaths
 +#select