Data flow: Depth 1 call contexts in store/load matching

Author: hvitved

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

@@ -2572,23 +2572,25 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
 
         pragma[nomagic]
         private predicate revFlowThroughArg(
-          DataFlowCall call, ArgNodeEx arg, FlowState state, ReturnCtx returnCtx, ApOption returnAp,
-          Ap ap
+          DataFlowCall call, ArgNodeEx arg, ParamNodeEx p, FlowState state, ReturnCtx returnCtx,
+          ApOption returnAp, Ap ap
         ) {
-          exists(ParamNodeEx p, Ap innerReturnAp |
+          exists(Ap innerReturnAp |
             revFlowThrough(call, returnCtx, p, state, _, returnAp, ap, innerReturnAp) and
             flowThroughIntoCall(call, arg, p, _, ap, innerReturnAp)
           )
         }
 
         pragma[nomagic]
-        predicate callMayFlowThroughRev(DataFlowCall call) {
+        additional predicate callMayFlowThroughRev(DataFlowCall call, ParamNodeEx p) {
           exists(ArgNodeEx arg, FlowState state, ReturnCtx returnCtx, ApOption returnAp, Ap ap |
             revFlow(arg, state, returnCtx, returnAp, ap) and
-            revFlowThroughArg(call, arg, state, returnCtx, returnAp, ap)
+            revFlowThroughArg(call, arg, p, state, returnCtx, returnAp, ap)
           )
         }
 
+        predicate callMayFlowThroughRev(DataFlowCall call) { callMayFlowThroughRev(call, _) }
+
         predicate callEdgeArgParam(
           DataFlowCall call, DataFlowCallable c, ArgNodeEx arg, ParamNodeEx p,
           boolean allowsFieldFlow, Ap ap
@@ -3856,24 +3858,68 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
       private module StoreReadMatchingInput implements StoreReadMatchingInputSig {
         class NodeEx = NodeExAlias;
 
-        predicate nodeRange(NodeEx node, boolean fromArg) {
-          exists(PrevStage::Ap ap |
-            PrevStage::revFlowAp(node, ap) and
+        /**
+         * Gets a call context for `node` that is relevant for either improved virtual
+         * dispatch or for flow-through.
+         */
+        pragma[nomagic]
+        private DataFlowCallOption getACallCtx(NodeEx node, PrevStage::Ap ap, boolean fromArg) {
+          exists(PrevStage::Cc cc, ParamNodeOption summaryCtx, PrevStage::ApOption argAp |
+            PrevStage::fwdFlow(node, _, cc, summaryCtx, _, argAp, _, ap, _)
+          |
+            PrevStage::instanceofCcCall(cc) and
+            fromArg = true and
             (
-              ap = true
-              or
-              PrevStage::storeStepCand(node, ap, _, _, _, _)
+              // virtual dispatch may be improved
+              exists(DataFlowCall call, DataFlowCallable c |
+                PrevStage::callEdgeArgParam(call, c, _, _, _, _) and
+                cc = Stage2Param::getSpecificCallContextCall(call, c) and
+                c = node.getEnclosingCallable() and
+                result = TDataFlowCallSome(call)
+              )
               or
-              PrevStage::readStepCand(_, _, node)
+              not cc = Stage2Param::getSpecificCallContextCall(_, _) and
+              (
+                // flow-through not possible
+                summaryCtx instanceof TParamNodeNone and
+                result = TDataFlowCallNone()
+                or
+                exists(DataFlowCall call, ParamNodeEx p, PrevStage::Ap argApSome |
+                  summaryCtx = TParamNodeSome(p.asNode()) and
+                  argAp = PrevStage::apSome(argApSome)
+                |
+                  if
+                    PrevStage::parameterMayFlowThrough(p, argApSome) and
+                    PrevStage::callMayFlowThroughRev(call, p)
+                  then
+                    // flow-through possible
+                    result = TDataFlowCallSome(call)
+                  else (
+                    // flow-through not possible, but `node` can reach a sink without
+                    // flowing back out
+                    PrevStage::callEdgeArgParam(call, _, _, p, _, _) and
+                    result = TDataFlowCallNone()
+                  )
+                )
+              )
             )
+            or
+            PrevStage::instanceofCcNoCall(cc) and
+            fromArg = false and
+            result = TDataFlowCallNone()
+          )
+        }
+
+        predicate nodeRange(NodeEx node, boolean fromArg, DataFlowCallOption summaryCtx) {
+          exists(PrevStage::Ap ap |
+            PrevStage::revFlowAp(node, ap) and
+            summaryCtx = getACallCtx(node, ap, fromArg)
           |
-            exists(PrevStage::Cc cc | PrevStage::fwdFlow(node, _, cc, _, _, _, _, ap, _) |
-              PrevStage::instanceofCcCall(cc) and
-              fromArg = true
-              or
-              PrevStage::instanceofCcNoCall(cc) and
-              fromArg = false
-            )
+            ap = true
+            or
+            PrevStage::storeStepCand(node, ap, _, _, _, _)
+            or
+            PrevStage::readStepCand(_, _, node)
           )
         }
 
@@ -3899,12 +3945,51 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
           )
         }
 
-        predicate callEdgeArgParam(NodeEx arg, NodeEx param) {
-          PrevStage::callEdgeArgParam(_, _, arg, param, true, true)
+        pragma[nomagic]
+        private predicate callEdgeArgParam(
+          DataFlowCall call, DataFlowCallable c, NodeEx arg, NodeEx param,
+          DataFlowCallOption innerCallCtx
+        ) {
+          PrevStage::callEdgeArgParam(call, c, arg, param, true, true) and
+          innerCallCtx = getACallCtx(param, true, true) and
+          (
+            innerCallCtx = TDataFlowCallNone()
+            or
+            innerCallCtx = TDataFlowCallSome(call)
+          )
+        }
+
+        pragma[nomagic]
+        private predicate callEdgeArgParamCallContextReduced(
+          DataFlowCall call, NodeEx arg, NodeEx param, Stage2Param::CcCall outerCcCall,
+          DataFlowCallOption innerCallCtx
+        ) {
+          exists(DataFlowCallable c |
+            callEdgeArgParam(call, c, arg, param, innerCallCtx) and
+            Stage2Param::viableImplCallContextReduced(call, outerCcCall) = c
+          )
+        }
+
+        bindingset[outerCallCtx]
+        predicate callEdgeArgParam(
+          NodeEx arg, NodeEx param, DataFlowCallOption outerCallCtx, DataFlowCallOption innerCallCtx
+        ) {
+          exists(DataFlowCall call | callEdgeArgParam(call, _, arg, param, innerCallCtx) |
+            outerCallCtx = TDataFlowCallNone()
+            or
+            exists(DataFlowCall outerCall, Stage2Param::CcCall outerCcCall |
+              outerCallCtx = TDataFlowCallSome(outerCall) and
+              outerCcCall = Stage2Param::getCallContextCall(outerCall, call.getEnclosingCallable())
+            |
+              callEdgeArgParamCallContextReduced(call, arg, param, outerCcCall, innerCallCtx)
+              or
+              Stage2Param::viableImplNotCallContextReduced(call, outerCcCall)
+            )
+          )
         }
 
-        predicate callEdgeReturn(NodeEx ret, NodeEx out, boolean mayFlowThrough) {
-          PrevStage::callEdgeReturn(_, _, ret, _, out, true, true) and
+        predicate callEdgeReturn(DataFlowCall call, NodeEx ret, NodeEx out, boolean mayFlowThrough) {
+          PrevStage::callEdgeReturn(call, _, ret, _, out, true, true) and
           if flowThroughOutOfCall(ret, out) then mayFlowThrough = true else mayFlowThrough = false
         }