Data flow: Workaround for lambda + capture flow

Author: hvitved

shared/dataflow/codeql/dataflow/DataFlow.qll

@@ -352,6 +352,8 @@ signature module InputSig<LocationSig Location> {
   Content getLambdaArgumentContent(LambdaCallKind kind, ArgumentPosition pos);
 
   predicate isLambdaInstanceParameter(ParameterNode p);
+
+  predicate isVariableCaptureContentSet(ContentSet c);
 }
 
 module Configs<LocationSig Location, InputSig<Location> Lang> {

shared/dataflow/codeql/dataflow/VariableCapture.qll

@@ -1051,8 +1051,8 @@ module Flow<LocationSig Location, InputSig<Location> Input> implements OutputSig
      * since normal use-use flow for `fn` does not take the overwrite at (2) into account.
      */
 
-    storeStepClosure(_, v, node, true)
-    or
+    // storeStepClosure(_, v, node, true)
+    // or
     exists(BasicBlock bb, int i |
       captureWrite(v, bb, i, false, _) and
       node = TSynthThisQualifier(bb, i, false)

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

@@ -450,13 +450,13 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
     bindingset[c]
     private predicate expectsContentEx(NodeEx n, Content c) {
       exists(ContentSet cs |
-        expectsContentCached(n.asNode(), cs) and
+        expectsContentCached(n, cs) and
         pragma[only_bind_out](c) = pragma[only_bind_into](cs).getAReadContent()
       )
     }
 
     pragma[nomagic]
-    private predicate notExpectsContent(NodeEx n) { not expectsContentCached(n.asNode(), _) }
+    private predicate notExpectsContent(NodeEx n) { not expectsContentCached(n, _) }
 
     pragma[nomagic]
     private predicate storeExUnrestricted(
@@ -2618,7 +2618,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
               (
                 castNode(this.asNode()) or
                 clearsContentCached(this.asNode(), _) or
-                expectsContentCached(this.asNode(), _) or
+                expectsContentCached(this, _) or
                 neverSkipInPathGraph(this.asNode()) or
                 Config::neverSkip(this.asNode())
               )

shared/dataflow/codeql/dataflow/internal/DataFlowImplCommon.qll

@@ -893,6 +893,8 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
       or
       result = this.asLambdaArgsNode().toString() + " [LambdaArgs]"
       or
+      result = this.asLambdaCaptureNode().toString() + " [LambdaCapture]"
+      or
       result = this.asLambdaInstancePostUpdateNode().toString() + " [LambdaPostUpdate]"
       or
       exists(DataFlowCall synthCall, ArgumentPosition apos, boolean isPost |
@@ -918,6 +920,8 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
 
     Node asLambdaArgsNode() { this = TNodeLambdaArgs(result) }
 
+    Node asLambdaCaptureNode() { this = TNodeLambdaCapture(result) }
+
     predicate isLambdaArgNode(DataFlowCall synthCall, ArgumentPosition apos, boolean isPost) {
       this = TNodeLambdaArg(synthCall, apos, isPost)
     }
@@ -935,6 +939,10 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
       or
       this = TNodeLambdaArgs(result)
       or
+      this = TNodeLambdaCapture(result)
+      or
+      this = TNodeLambdaCapture(result)
+      or
       exists(DataFlowCall synthCall |
         this = TNodeLambdaArg(synthCall, _, _) and
         lambdaCreation(result, _, _, synthCall)
@@ -963,6 +971,8 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
       or
       nodeDataFlowType(this.asLambdaArgsNode(), result)
       or
+      nodeDataFlowType(this.asLambdaCaptureNode(), result)
+      or
       exists(
         DataFlowCall synthCall, ArgumentPosition apos, DataFlowCallable c, ParameterNode p,
         ParameterPosition ppos
@@ -1112,7 +1122,12 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
     predicate clearsContentCached(Node n, ContentSet c) { clearsContent(n, c) }
 
     cached
-    predicate expectsContentCached(Node n, ContentSet c) { expectsContent(n, c) }
+    predicate expectsContentCached(NodeEx n, ContentSet c) {
+      expectsContent(n.asNode(), c)
+      or
+      exists(n.asLambdaCaptureNode()) and
+      isVariableCaptureContentSet(c)
+    }
 
     cached
     predicate isUnreachableInCallCached(NodeRegion nr, DataFlowCall call) {
@@ -1133,6 +1148,7 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
       exists(n.asLambdaInstancePostUpdateNode()) or
       exists(n.asLambdaMallocNode()) or
       exists(n.asLambdaArgsNode()) or
+      exists(n.asLambdaCaptureNode()) or
       n.isLambdaArgNode(_, _, _) or
       hiddenNode(any(NodeEx p | n.asParamReturnNode() = p.asNode()))
     }
@@ -1554,7 +1570,7 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
               argumentValueFlowsThroughCand(arg, node, false)
             )
           ) and
-          not expectsContentCached(node, _)
+          not expectsContent(node, _)
         }
 
         pragma[nomagic]
@@ -2064,10 +2080,11 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
           lambdaCreation(_, _, c, synthCall) and
           isParameterNode(p, c, ppos) and
           parameterMatch(ppos, apos) and
-          not isLambdaInstanceParameter(p) and
+          // not isLambdaInstanceParameter(p) and
           exists(ispost)
         )
-      }
+      } or
+      TNodeLambdaCapture(Node receiver) { lambdaCall(_, _, receiver) }
 
     /*
      * foo(() => "taint"); // taint --store(ReturnValue)--> this (post-update) [ReturnValue]
@@ -2108,6 +2125,17 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
       or
       LambdaFlow::lambdaFlowsToPostUpdate(node2.asLambdaArgsNode(), node1.asNode()) and
       model = ""
+      or
+      // When data is stored in a captured variable content and reaches a lambda call,
+      // we need it to propagate back out to the lambda. We do this by adding flow
+      // from the lambda receiver to the post-update of the lambda receiver, but _only_
+      // for captured variable content. The latter restriction is enforced by going via
+      // an intermediate `expectsContent` node.
+      node1.asNode() = node2.asLambdaCaptureNode() and
+      model = ""
+      or
+      node2.asNode().(PostUpdateNode).getPreUpdateNode() = node1.asLambdaCaptureNode() and
+      model = ""
     }
 
     cached