Merge pull request #17663 from aschackmull/dataflow/speculative-flow

Dataflow: Add support for speculative taint flow.

shared/dataflow/codeql/dataflow/DataFlow.qll

@@ -676,6 +676,12 @@ module DataFlowMake<LocationSig Location, InputSig<Location> Lang> {
       predicate isAdditionalFlowStep(Node node1, Node node2, string model) {
         Config::isAdditionalFlowStep(node1, node2) and model = "Config"
       }
+
+      predicate isAdditionalFlowStep(
+        Node node1, FlowState state1, Node node2, FlowState state2, string model
+      ) {
+        Config::isAdditionalFlowStep(node1, state1, node2, state2) and model = "Config"
+      }
     }
 
     import Impl<C>

shared/dataflow/codeql/dataflow/TaintTracking.qll

@@ -29,6 +29,12 @@ signature module InputSig<LocationSig Location, DF::InputSig<Location> Lang> {
    */
   bindingset[node]
   predicate defaultImplicitTaintRead(Lang::Node node, Lang::ContentSet c);
+
+  /**
+   * Holds if the additional step from `src` to `sink` should be considered in
+   * speculative taint flow exploration.
+   */
+  predicate speculativeTaintStep(Lang::Node src, Lang::Node sink);
 }
 
 /**
@@ -63,7 +69,7 @@ module TaintFlowMake<
         Config::isSink(node) or
         Config::isSink(node, _) or
         Config::isAdditionalFlowStep(node, _, _) or
-        Config::isAdditionalFlowStep(node, _, _, _)
+        Config::isAdditionalFlowStep(node, _, _, _, _)
       ) and
       defaultImplicitTaintRead(node, c)
     }
@@ -108,6 +114,13 @@ module TaintFlowMake<
       ) {
         Config::isAdditionalFlowStep(node1, node2) and model = "Config"
       }
+
+      predicate isAdditionalFlowStep(
+        DataFlowLang::Node node1, FlowState state1, DataFlowLang::Node node2, FlowState state2,
+        string model
+      ) {
+        Config::isAdditionalFlowStep(node1, state1, node2, state2) and model = "Config"
+      }
     }
 
     private module C implements DataFlowInternal::FullStateConfigSig {
@@ -123,4 +136,132 @@ module TaintFlowMake<
   {
     import GlobalWithState<Config>
   }
+
+  signature int speculationLimitSig();
+
+  private module AddSpeculativeTaintSteps<
+    DataFlowInternal::FullStateConfigSig Config, speculationLimitSig/0 speculationLimit> implements
+    DataFlowInternal::FullStateConfigSig
+  {
+    import Config
+
+    private predicate relevantState(Config::FlowState state) {
+      Config::isSource(_, state)
+      or
+      exists(Config::FlowState state0 |
+        relevantState(state0) and Config::isAdditionalFlowStep(_, state0, _, state, _)
+      )
+    }
+
+    private newtype TFlowState =
+      TMkFlowState(Config::FlowState state, int spec) {
+        relevantState(state) and spec = [0 .. speculationLimit()]
+      }
+
+    class FlowState extends TFlowState {
+      private Config::FlowState state;
+      private int spec;
+
+      FlowState() { this = TMkFlowState(state, spec) }
+
+      string toString() { result = "FlowState" }
+
+      Config::FlowState getState() { result = state }
+
+      int getSpec() { result = spec }
+    }
+
+    predicate isSource(DataFlowLang::Node source, FlowState state) {
+      Config::isSource(source, state.getState()) and state.getSpec() = 0
+    }
+
+    predicate isSink(DataFlowLang::Node sink, FlowState state) {
+      Config::isSink(sink, state.getState())
+    }
+
+    predicate isBarrier(DataFlowLang::Node node, FlowState state) {
+      Config::isBarrier(node, state.getState())
+    }
+
+    predicate isBarrierIn(DataFlowLang::Node node, FlowState state) {
+      Config::isBarrierIn(node, state.getState())
+    }
+
+    predicate isBarrierOut(DataFlowLang::Node node, FlowState state) {
+      Config::isBarrierOut(node, state.getState())
+    }
+
+    predicate isAdditionalFlowStep(
+      DataFlowLang::Node node1, FlowState state1, DataFlowLang::Node node2, FlowState state2,
+      string model
+    ) {
+      Config::isAdditionalFlowStep(node1, state1.getState(), node2, state2.getState(), model) and
+      state1.getSpec() = state2.getSpec()
+      or
+      speculativeTaintStep(node1, node2) and
+      not defaultAdditionalTaintStep(node1, node2, _) and
+      not Config::isAdditionalFlowStep(node1, _, node2, _, _) and
+      not Config::isAdditionalFlowStep(node1, node2, _) and
+      model = "Speculative" and
+      state1.getSpec() + 1 = state2.getSpec() and
+      state1.getState() = state2.getState()
+    }
+  }
+
+  /**
+   * Constructs a global taint tracking computation that also allows a given
+   * maximum number of speculative taint steps.
+   */
+  module SpeculativeGlobal<DataFlow::ConfigSig Config, speculationLimitSig/0 speculationLimit>
+    implements DataFlow::GlobalFlowSig
+  {
+    private module Config0 implements DataFlowInternal::FullStateConfigSig {
+      import DataFlowInternal::DefaultState<Config>
+      import Config
+
+      predicate isAdditionalFlowStep(
+        DataFlowLang::Node node1, DataFlowLang::Node node2, string model
+      ) {
+        Config::isAdditionalFlowStep(node1, node2) and model = "Config"
+      }
+    }
+
+    private module C implements DataFlowInternal::FullStateConfigSig {
+      import AddTaintDefaults<AddSpeculativeTaintSteps<Config0, speculationLimit/0>>
+    }
+
+    import DataFlowInternal::Impl<C>
+  }
+
+  /**
+   * Constructs a global taint tracking computation using flow state that also
+   * allows a given maximum number of speculative taint steps.
+   */
+  module SpeculativeGlobalWithState<
+    DataFlow::StateConfigSig Config, speculationLimitSig/0 speculationLimit> implements
+    DataFlow::GlobalFlowSig
+  {
+    private module Config0 implements DataFlowInternal::FullStateConfigSig {
+      import Config
+
+      predicate isAdditionalFlowStep(
+        DataFlowLang::Node node1, DataFlowLang::Node node2, string model
+      ) {
+        Config::isAdditionalFlowStep(node1, node2) and model = "Config"
+      }
+
+      predicate isAdditionalFlowStep(
+        DataFlowLang::Node node1, FlowState state1, DataFlowLang::Node node2, FlowState state2,
+        string model
+      ) {
+        Config::isAdditionalFlowStep(node1, state1, node2, state2) and model = "Config"
+      }
+    }
+
+    private module C implements DataFlowInternal::FullStateConfigSig {
+      import AddTaintDefaults<AddSpeculativeTaintSteps<Config0, speculationLimit/0>>
+    }
+
+    import DataFlowInternal::Impl<C>
+  }
 }

shared/dataflow/codeql/dataflow/internal/DataFlowImpl.qll

@@ -73,7 +73,9 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
      * Holds if data may flow from `node1` to `node2` in addition to the normal data-flow steps.
      * This step is only applicable in `state1` and updates the flow state to `state2`.
      */
-    predicate isAdditionalFlowStep(Node node1, FlowState state1, Node node2, FlowState state2);
+    predicate isAdditionalFlowStep(
+      Node node1, FlowState state1, Node node2, FlowState state2, string model
+    );
 
     /**
      * Holds if an arbitrary number of implicit read steps of content `c` may be
@@ -153,7 +155,9 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
 
     predicate isBarrierOut(Node node, FlowState state) { none() }
 
-    predicate isAdditionalFlowStep(Node node1, FlowState state1, Node node2, FlowState state2) {
+    predicate isAdditionalFlowStep(
+      Node node1, FlowState state1, Node node2, FlowState state2, string model
+    ) {
       none()
     }
   }
@@ -364,12 +368,13 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
     }
 
     private predicate additionalLocalStateStep(
-      NodeEx node1, FlowState s1, NodeEx node2, FlowState s2
+      NodeEx node1, FlowState s1, NodeEx node2, FlowState s2, string model
     ) {
       exists(Node n1, Node n2 |
         node1.asNodeOrImplicitRead() = n1 and
         node2.asNode() = n2 and
-        Config::isAdditionalFlowStep(pragma[only_bind_into](n1), s1, pragma[only_bind_into](n2), s2) and
+        Config::isAdditionalFlowStep(pragma[only_bind_into](n1), s1, pragma[only_bind_into](n2), s2,
+          model) and
         getNodeEnclosingCallable(n1) = getNodeEnclosingCallable(n2) and
         stateStepFilter(node1, s1, node2, s2)
       )
@@ -402,11 +407,14 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
       )
     }
 
-    private predicate additionalJumpStateStep(NodeEx node1, FlowState s1, NodeEx node2, FlowState s2) {
+    private predicate additionalJumpStateStep(
+      NodeEx node1, FlowState s1, NodeEx node2, FlowState s2, string model
+    ) {
       exists(Node n1, Node n2 |
         node1.asNodeOrImplicitRead() = n1 and
         node2.asNode() = n2 and
-        Config::isAdditionalFlowStep(pragma[only_bind_into](n1), s1, pragma[only_bind_into](n2), s2) and
+        Config::isAdditionalFlowStep(pragma[only_bind_into](n1), s1, pragma[only_bind_into](n2), s2,
+          model) and
         getNodeEnclosingCallable(n1) != getNodeEnclosingCallable(n2) and
         stateStepFilter(node1, s1, node2, s2) and
         not Config::getAFeature() instanceof FeatureEqualSourceSinkCallContext
@@ -537,13 +545,13 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
         exists(NodeEx mid | fwdFlow(mid, cc) |
           localFlowStepEx(mid, node, _) or
           additionalLocalFlowStep(mid, node, _) or
-          additionalLocalStateStep(mid, _, node, _)
+          additionalLocalStateStep(mid, _, node, _, _)
         )
         or
         exists(NodeEx mid | fwdFlow(mid, _) and cc = false |
           jumpStepEx(mid, node) or
           additionalJumpStep(mid, node, _) or
-          additionalJumpStateStep(mid, _, node, _)
+          additionalJumpStateStep(mid, _, node, _, _)
         )
         or
         // store
@@ -684,8 +692,8 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
 
       private predicate stateStepFwd(FlowState state1, FlowState state2) {
         exists(NodeEx node1 |
-          additionalLocalStateStep(node1, state1, _, state2) or
-          additionalJumpStateStep(node1, state1, _, state2)
+          additionalLocalStateStep(node1, state1, _, state2, _) or
+          additionalJumpStateStep(node1, state1, _, state2, _)
         |
           fwdFlow(node1)
         )
@@ -730,13 +738,13 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
         exists(NodeEx mid | revFlow(mid, toReturn) |
           localFlowStepEx(node, mid, _) or
           additionalLocalFlowStep(node, mid, _) or
-          additionalLocalStateStep(node, _, mid, _)
+          additionalLocalStateStep(node, _, mid, _, _)
         )
         or
         exists(NodeEx mid | revFlow(mid, _) and toReturn = false |
           jumpStepEx(node, mid) or
           additionalJumpStep(node, mid, _) or
-          additionalJumpStateStep(node, _, mid, _)
+          additionalJumpStateStep(node, _, mid, _, _)
         )
         or
         // store
@@ -854,8 +862,8 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
 
       private predicate stateStepRev(FlowState state1, FlowState state2) {
         exists(NodeEx node1, NodeEx node2 |
-          additionalLocalStateStep(node1, state1, node2, state2) or
-          additionalJumpStateStep(node1, state1, node2, state2)
+          additionalLocalStateStep(node1, state1, node2, state2, _) or
+          additionalJumpStateStep(node1, state1, node2, state2, _)
         |
           revFlow(node1, _) and
           revFlow(node2, _) and
@@ -1075,8 +1083,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
     ) {
       Stage1::revFlow(node1) and
       Stage1::revFlow(node2) and
-      additionalLocalStateStep(node1, state1, node2, state2) and
-      label = "Config" and
+      additionalLocalStateStep(node1, state1, node2, state2, label) and
       t = node2.getDataFlowType() and
       lcc.relevantFor(node1.getEnclosingCallable()) and
       not isUnreachableInCall1(node1, lcc) and
@@ -1628,7 +1635,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
           or
           exists(NodeEx mid, FlowState state0 |
             fwdFlow(mid, state0, _, _, _, ap, apa) and
-            additionalJumpStateStep(mid, state0, node, state) and
+            additionalJumpStateStep(mid, state0, node, state, _) and
             t = getNodeTyp(node) and
             ap instanceof ApNil
           )
@@ -2293,7 +2300,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
           )
           or
           exists(NodeEx mid, FlowState state0 |
-            additionalJumpStateStep(node, state, mid, state0) and
+            additionalJumpStateStep(node, state, mid, state0, _) and
             revFlow(pragma[only_bind_into](mid), pragma[only_bind_into](state0), _, _, ap) and
             ap instanceof ApNil
           )
@@ -2658,7 +2665,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
               or
               additionalJumpStep(_, node, _)
               or
-              additionalJumpStateStep(_, _, node, state)
+              additionalJumpStateStep(_, _, node, state, _)
               or
               node instanceof ParamNodeEx
               or
@@ -2706,9 +2713,9 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
               exists(NodeEx next, FlowState s |
                 revFlow(next, s, pragma[only_bind_into](ap)) and ap instanceof ApNil
               |
-                additionalJumpStateStep(node, state, next, s)
+                additionalJumpStateStep(node, state, next, s, _)
                 or
-                additionalLocalStateStep(node, state, next, s) and
+                additionalLocalStateStep(node, state, next, s, _) and
                 s != state
               )
               or
@@ -3253,10 +3260,9 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
               t = getNodeTyp(node) and
               ap instanceof ApNil
               or
-              additionalJumpStateStep(mid, state0, node, state) and
+              additionalJumpStateStep(mid, state0, node, state, label) and
               t = getNodeTyp(node) and
-              ap instanceof ApNil and
-              label = "Config"
+              ap instanceof ApNil
             )
             or
             // flow into a callable
@@ -4817,7 +4823,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
           or
           additionalJumpStep(node1, node2, _)
           or
-          additionalJumpStateStep(node1, _, node2, _)
+          additionalJumpStateStep(node1, _, node2, _, _)
           or
           // flow into callable
           viableParamArgEx(_, node2, node1)
@@ -4931,10 +4937,10 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
       private predicate relevantState(FlowState state) {
         sourceNode(_, state) or
         sinkNodeWithState(_, state) or
-        additionalLocalStateStep(_, state, _, _) or
-        additionalLocalStateStep(_, _, _, state) or
-        additionalJumpStateStep(_, state, _, _) or
-        additionalJumpStateStep(_, _, _, state)
+        additionalLocalStateStep(_, state, _, _, _) or
+        additionalLocalStateStep(_, _, _, state, _) or
+        additionalJumpStateStep(_, state, _, _, _) or
+        additionalJumpStateStep(_, _, _, state, _)
       }
 
       private predicate revSinkNode(NodeEx node, FlowState state) {
@@ -5270,7 +5276,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
           t = node.getDataFlowType() and
           ap = TPartialNil()
           or
-          additionalLocalStateStep(mid.getNodeEx(), mid.getState(), node, state) and
+          additionalLocalStateStep(mid.getNodeEx(), mid.getState(), node, state, _) and
           cc = mid.getCallContext() and
           sc1 = mid.getSummaryCtx1() and
           sc2 = mid.getSummaryCtx2() and
@@ -5305,7 +5311,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
         ap = TPartialNil() and
         isStoreStep = false
         or
-        additionalJumpStateStep(mid.getNodeEx(), mid.getState(), node, state) and
+        additionalJumpStateStep(mid.getNodeEx(), mid.getState(), node, state, _) and
         cc = callContextNone() and
         sc1 = TSummaryCtx1None() and
         sc2 = TSummaryCtx2None() and
@@ -5577,7 +5583,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
         ap = TPartialNil() and
         isStoreStep = false
         or
-        additionalLocalStateStep(node, state, mid.getNodeEx(), mid.getState()) and
+        additionalLocalStateStep(node, state, mid.getNodeEx(), mid.getState(), _) and
         sc1 = mid.getSummaryCtx1() and
         sc2 = mid.getSummaryCtx2() and
         sc3 = mid.getSummaryCtx3() and
@@ -5602,7 +5608,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
         ap = TPartialNil() and
         isStoreStep = false
         or
-        additionalJumpStateStep(node, state, mid.getNodeEx(), mid.getState()) and
+        additionalJumpStateStep(node, state, mid.getNodeEx(), mid.getState(), _) and
         sc1 = TRevSummaryCtx1None() and
         sc2 = TRevSummaryCtx2None() and
         sc3 = TRevSummaryCtx3None() and

shared/dataflow/codeql/dataflow/internal/DataFlowImplConsistency.qll

@@ -323,4 +323,18 @@ module MakeConsistency<
     lambdaCall(call, _, receiver) and
     not nodeGetEnclosingCallable(receiver) = call.getEnclosingCallable()
   }
+
+  query predicate speculativeStepAlreadyHasModel(Node n1, Node n2, string model) {
+    speculativeTaintStep(n1, n2) and
+    not defaultAdditionalTaintStep(n1, n2, _) and
+    (
+      simpleLocalFlowStep(n1, n2, _) and model = "SimpleLocalFlowStep"
+      or
+      exists(DataFlowCall call |
+        exists(viableCallable(call)) and
+        isArgumentNode(n1, call, _) and
+        model = "dispatch"
+      )
+    )
+  }
 }