Python taint-tracking. Rework handling of sequences and dicts of taint for performance.

markshannon · markshannon · commit 3872c7a1f992 · 2019-08-29T10:31:50.000+01:00
diff --git a/python/ql/src/semmle/python/dataflow/Implementation.qll b/python/ql/src/semmle/python/dataflow/Implementation.qll
@@ -361,7 +361,19 @@ class TaintTrackingImplementation extends string {
             |
             kind = srckind.getTaintForFlowStep(srcnode.asCfgNode(), node.asCfgNode(), edgeLabel)
             or
-            kind.isResultOfStep(srckind, srcnode.asCfgNode(), node.asCfgNode(), edgeLabel)
+            kind = srckind.(CollectionKind).getMember() and
+            srckind.(CollectionKind).flowToMember(srcnode, node) and
+            edgeLabel = "to member"
+            or
+            srckind = kind.(CollectionKind).getMember() and
+            kind.(CollectionKind).flowFromMember(srcnode, node) and
+            edgeLabel = "from member"
+            or
+            kind = srckind and srckind.flowStep(srcnode, node, edgeLabel)
+            or
+            kind = srckind and srckind instanceof DictKind and DictKind::flowStep(srcnode.asCfgNode(), node.asCfgNode(), edgeLabel)
+            or
+            kind = srckind and srckind instanceof SequenceKind and SequenceKind::flowStep(srcnode.asCfgNode(), node.asCfgNode(), edgeLabel)
         )
     }
 
@@ -755,11 +767,13 @@ private class EssaTaintTracking extends string {
         or
         result = testEvaluates(not_operand(test), use, kind).booleanNot()
         or
+        kind.taints(use) and
         exists(ControlFlowNode const |
             Filters::equality_test(test, use, result.booleanNot(), const) and
             const.getNode() instanceof ImmutableLiteral
         )
         or
+        kind.taints(use) and
         exists(ControlFlowNode c, ClassValue cls |
             Filters::isinstance(test, c, use) and
             c.pointsTo(cls)
diff --git a/python/ql/src/semmle/python/security/TaintTracking.qll b/python/ql/src/semmle/python/security/TaintTracking.qll
@@ -126,12 +126,6 @@ abstract class TaintKind extends string {
         result = this.getTaintForFlowStep(fromnode, tonode) and edgeLabel = "custom taint flow step for " + this
     }
 
-    predicate isResultOfStep(TaintKind fromkind, ControlFlowNode fromnode, ControlFlowNode tonode) { none() }
-
-    predicate isResultOfStep(TaintKind fromkind, ControlFlowNode fromnode, ControlFlowNode tonode, string edgeLabel) {
-        this.isResultOfStep(fromkind, fromnode, tonode) and edgeLabel = "custom taint reverse flow step for " + this
-    }
-
     /** DEPRECATED -- Use `TaintFlow.additionalFlowStepVar(EssaVariable fromvar, EssaVariable tovar, TaintKind kind)` instead.
      *
      * Holds if this kind of taint passes from variable `fromvar` to  variable `tovar`
@@ -179,6 +173,10 @@ abstract class TaintKind extends string {
      */
     TaintKind getTaintForIteration() { none() }
 
+    predicate flowStep(DataFlow::Node fromnode, DataFlow::Node tonode, string edgeLabel) {
+        this.additionalFlowStepVar(fromnode.asVariable(), tonode.asVariable()) and
+        edgeLabel = "custom taint variable step"
+    }
 }
 
 /**
@@ -202,6 +200,11 @@ abstract class CollectionKind extends TaintKind {
         not this.charAt(2) = "[" and not this.charAt(2) = "{"
     }
 
+    abstract TaintKind getMember();
+
+    abstract predicate flowFromMember(DataFlow::Node fromnode, DataFlow::Node tonode);
+
+    abstract predicate flowToMember(DataFlow::Node fromnode, DataFlow::Node tonode);
 }
 
 /** A taint kind representing a flat collections of kinds.
@@ -227,8 +230,6 @@ class SequenceKind extends CollectionKind {
             result = this.getItem() and
             result.getType() = ObjectInternal::builtin("str")
         )
-        or
-        result = this.getItem() and SequenceKind::itemFlowStep(fromnode, tonode)
     }
 
     override TaintKind getTaintOfMethodResult(string name) {
@@ -243,25 +244,32 @@ class SequenceKind extends CollectionKind {
         result = itemKind
     }
 
-    override predicate isResultOfStep(TaintKind fromkind, ControlFlowNode fromnode, ControlFlowNode tonode) {
-        SequenceKind::flowStep(fromnode, tonode) and this = fromkind
-        or
-        sequence_construct(fromnode, tonode) and this.getItem() = fromkind
+    override TaintKind getMember() {
+        result = itemKind
+    }
+
+    override predicate flowFromMember(DataFlow::Node fromnode, DataFlow::Node tonode) {
+        sequence_construct(fromnode.asCfgNode(), tonode.asCfgNode())
+    }
+
+    override predicate flowToMember(DataFlow::Node fromnode, DataFlow::Node tonode) {
+        SequenceKind::itemFlowStep(fromnode.asCfgNode(), tonode.asCfgNode())
     }
 
 }
 
 
 module SequenceKind {
 
-    predicate flowStep(ControlFlowNode fromnode, ControlFlowNode tonode) {
-        tonode.(BinaryExprNode).getAnOperand() = fromnode
+    predicate flowStep(ControlFlowNode fromnode, ControlFlowNode tonode, string edgeLabel) {
+        tonode.(BinaryExprNode).getAnOperand() = fromnode and edgeLabel = "binary operation"
         or
-        Implementation::copyCall(fromnode, tonode)
+        Implementation::copyCall(fromnode, tonode) and
+        edgeLabel = "dict copy"
         or
-        sequence_call(fromnode, tonode)
+        sequence_call(fromnode, tonode) and edgeLabel = "sequence construction"
         or
-        subscript_slice(fromnode, tonode)
+        subscript_slice(fromnode, tonode) and edgeLabel = "slicing"
     }
 
     predicate itemFlowStep(ControlFlowNode fromnode, ControlFlowNode tonode) {
@@ -270,6 +278,17 @@ module SequenceKind {
 
 }
 
+module DictKind {
+     predicate flowStep(ControlFlowNode fromnode, ControlFlowNode tonode, string edgeLabel) {
+        Implementation::copyCall(fromnode, tonode) and
+        edgeLabel = "dict copy"
+        or
+        tonode.(CallNode).getFunction().pointsTo(ObjectInternal::builtin("dict")) and
+        tonode.(CallNode).getArg(0) = fromnode and
+        edgeLabel = "dict() call"
+    }
+}
+
 
 /* Helper for sequence flow steps */
 pragma [noinline]
@@ -314,17 +333,16 @@ class DictKind extends CollectionKind {
         result = "dict of " + valueKind
     }
 
-    override predicate isResultOfStep(TaintKind fromkind, ControlFlowNode fromnode, ControlFlowNode tonode) {
-        Implementation::copyCall(fromnode, tonode) and this = fromkind
-        or
-        tonode.(CallNode).getFunction().pointsTo(ObjectInternal::builtin("dict")) and
-        tonode.(CallNode).getArg(0) = fromnode and this = fromkind
-        or
-        dict_construct(fromnode, tonode) and this.getValue() = fromkind
+    override TaintKind getMember() {
+        result = valueKind
     }
 
-    override TaintKind getTaintForFlowStep(ControlFlowNode fromnode, ControlFlowNode tonode) {
-        subscript_index(fromnode, tonode) and result = this.getValue()
+    override predicate flowFromMember(DataFlow::Node fromnode, DataFlow::Node tonode) {
+        dict_construct(fromnode.asCfgNode(), tonode.asCfgNode())
+    }
+
+    override predicate flowToMember(DataFlow::Node fromnode, DataFlow::Node tonode) {
+        subscript_index(fromnode.asCfgNode(), tonode.asCfgNode())
     }
 
 }
