Java: Add proper support for variable capture flow.

java/ql/lib/semmle/code/java/dataflow/internal/DataFlowDispatch.qll

@@ -156,7 +156,7 @@ private module DispatchImpl {
 
   private module Unification = MkUnification<unificationTargetLeft/1, unificationTargetRight/1>;
 
-  private int parameterPosition() { result in [-1, any(Parameter p).getPosition()] }
+  private int parameterPosition() { result in [-2, -1, any(Parameter p).getPosition()] }
 
   /** A parameter position represented by an integer. */
   class ParameterPosition extends int {

java/ql/lib/semmle/code/java/dataflow/internal/DataFlowNodes.qll

@@ -55,7 +55,9 @@ private module Cached {
       )
     } or
     TFlowSummaryNode(FlowSummaryImpl::Private::SummaryNode sn) or
-    TFieldValueNode(Field f)
+    TFieldValueNode(Field f) or
+    TParameterPostUpdate(CapturedParameter p) or
+    TClosureNode(CaptureFlow::ClosureNode cn)
 
   cached
   newtype TContent =
@@ -64,6 +66,7 @@ private module Cached {
     TCollectionContent() or
     TMapKeyContent() or
     TMapValueContent() or
+    TClosureContent(CapturedVariable v) or
     TSyntheticFieldContent(SyntheticField s)
 
   cached
@@ -73,6 +76,7 @@ private module Cached {
     TCollectionContentApprox() or
     TMapKeyContentApprox() or
     TMapValueContentApprox() or
+    TClosureContentApprox(CapturedVariable v) or
     TSyntheticFieldApproxContent()
 }
 
@@ -127,6 +131,8 @@ module Public {
       or
       result = this.(ImplicitPostUpdateNode).getPreUpdateNode().getType()
       or
+      result = this.(ClosureNode).getTypeImpl()
+      or
       result = this.(FieldValueNode).getField().getType()
     }
 
@@ -359,6 +365,10 @@ private class ImplicitExprPostUpdate extends ImplicitPostUpdateNode, TImplicitEx
   }
 }
 
+private class ParameterPostUpdate extends ImplicitPostUpdateNode, TParameterPostUpdate {
+  override Node getPreUpdateNode() { this = TParameterPostUpdate(result.asParameter()) }
+}
+
 module Private {
   private import DataFlowDispatch
 
@@ -372,6 +382,7 @@ module Private {
     result.asCallable() = n.(MallocNode).getClassInstanceExpr().getEnclosingCallable() or
     result = nodeGetEnclosingCallable(n.(ImplicitPostUpdateNode).getPreUpdateNode()) or
     result.asSummarizedCallable() = n.(FlowSummaryNode).getSummarizedCallable() or
+    result = n.(ClosureNode).getCaptureFlowNode().getEnclosingCallable() or
     result.asFieldScope() = n.(FieldValueNode).getField()
   }
 
@@ -400,6 +411,8 @@ module Private {
       this = getInstanceArgument(_)
       or
       this.(FlowSummaryNode).isArgumentOf(_, _)
+      or
+      this.(ClosureNode).isArgumentOf(_, _)
     }
 
     /**
@@ -417,6 +430,8 @@ module Private {
       pos = -1 and this = getInstanceArgument(call.asCall())
       or
       this.(FlowSummaryNode).isArgumentOf(call, pos)
+      or
+      this.(ClosureNode).isArgumentOf(call, pos)
     }
 
     /** Gets the call in which this node is an argument. */
@@ -491,6 +506,34 @@ module Private {
       c.asSummarizedCallable() = this.getSummarizedCallable() and pos = this.getPosition()
     }
   }
+
+  /**
+   * A synthesized data flow node representing a closure object that tracks
+   * captured variables.
+   */
+  class ClosureNode extends Node, TClosureNode {
+    CaptureFlow::ClosureNode getCaptureFlowNode() { this = TClosureNode(result) }
+
+    override Location getLocation() { result = this.getCaptureFlowNode().getLocation() }
+
+    override string toString() { result = this.getCaptureFlowNode().toString() }
+
+    predicate isParameter(DataFlowCallable c) { this.getCaptureFlowNode().isParameter(c) }
+
+    predicate isArgumentOf(DataFlowCall call, int pos) {
+      this.getCaptureFlowNode().isArgument(call) and pos = -2
+    }
+
+    Type getTypeImpl() { result instanceof TypeObject }
+  }
+
+  class ClosureParameterNode extends ParameterNode, ClosureNode {
+    ClosureParameterNode() { this.isParameter(_) }
+
+    override predicate isParameterOf(DataFlowCallable c, int pos) {
+      this.isParameter(c) and pos = -2
+    }
+  }
 }
 
 private import Private
@@ -520,3 +563,13 @@ private class SummaryPostUpdateNode extends FlowSummaryNode, PostUpdateNode {
 
   override Node getPreUpdateNode() { result = pre }
 }
+
+private class ClosurePostUpdateNode extends PostUpdateNode, ClosureNode {
+  private ClosureNode pre;
+
+  ClosurePostUpdateNode() {
+    CaptureFlow::closurePostUpdateNode(this.getCaptureFlowNode(), pre.getCaptureFlowNode())
+  }
+
+  override Node getPreUpdateNode() { result = pre }
+}

java/ql/lib/semmle/code/java/dataflow/internal/DataFlowPrivate.qll

@@ -10,6 +10,7 @@ private import semmle.code.java.dataflow.FlowSummary
 private import FlowSummaryImpl as FlowSummaryImpl
 private import DataFlowImplConsistency
 private import DataFlowNodes
+private import codeql.dataflow.VariableCapture as VariableCapture
 import DataFlowNodes::Private
 
 private newtype TReturnKind = TNormalReturnKind()
@@ -51,24 +52,110 @@ private predicate fieldStep(Node node1, Node node2) {
   )
 }
 
-/**
- * Holds if data can flow from `node1` to `node2` through variable capture.
- */
-private predicate variableCaptureStep(Node node1, ExprNode node2) {
-  exists(SsaImplicitInit closure, SsaVariable captured |
-    closure.captures(captured) and
-    node2.getExpr() = closure.getAFirstUse()
+private module CaptureInput implements VariableCapture::InputSig {
+  private import java as J
+
+  class Location = J::Location;
+
+  class BasicBlock instanceof J::BasicBlock {
+    string toString() { result = super.toString() }
+
+    DataFlowCallable getEnclosingCallable() { result.asCallable() = super.getEnclosingCallable() }
+  }
+
+  BasicBlock getImmediateBasicBlockDominator(BasicBlock bb) { bbIDominates(result, bb) }
+
+  BasicBlock getABasicBlockSuccessor(BasicBlock bb) {
+    result = bb.(J::BasicBlock).getABBSuccessor()
+  }
+
+  //TODO: support capture of `this` in lambdas
+  class CapturedVariable instanceof LocalScopeVariable {
+    CapturedVariable() {
+      2 <=
+        strictcount(J::Callable c |
+          c = this.getCallable() or c = this.getAnAccess().getEnclosingCallable()
+        )
+    }
+
+    string toString() { result = super.toString() }
+
+    DataFlowCallable getCallable() { result.asCallable() = super.getCallable() }
+  }
+
+  class CapturedParameter extends CapturedVariable instanceof Parameter { }
+
+  additional predicate capturedVarUpdate(
+    J::BasicBlock bb, int i, CapturedVariable v, VariableUpdate upd
+  ) {
+    upd.getDestVar() = v and bb.getNode(i) = upd
+  }
+
+  additional predicate capturedVarRead(J::BasicBlock bb, int i, CapturedVariable v, RValue rv) {
+    v.(LocalScopeVariable).getAnAccess() = rv and bb.getNode(i) = rv
+  }
+
+  predicate variableWrite(BasicBlock bb, int i, CapturedVariable v, Location loc) {
+    exists(VariableUpdate upd | capturedVarUpdate(bb, i, v, upd) and loc = upd.getLocation())
+  }
+
+  predicate variableRead(BasicBlock bb, int i, CapturedVariable v, Location loc) {
+    exists(RValue rv | capturedVarRead(bb, i, v, rv) and loc = rv.getLocation())
+  }
+
+  class Callable = DataFlowCallable;
+
+  class Call instanceof DataFlowCall {
+    string toString() { result = super.toString() }
+
+    Location getLocation() { result = super.getLocation() }
+
+    DataFlowCallable getEnclosingCallable() { result = super.getEnclosingCallable() }
+
+    predicate hasCfgNode(BasicBlock bb, int i) { super.asCall() = bb.(J::BasicBlock).getNode(i) }
+  }
+}
+
+class CapturedVariable = CaptureInput::CapturedVariable;
+
+class CapturedParameter = CaptureInput::CapturedParameter;
+
+module CaptureFlow = VariableCapture::Flow<CaptureInput>;
+
+private predicate captureStoreStep(Node node1, ClosureContent c, Node node2) {
+  exists(BasicBlock bb, int i, CaptureInput::CapturedVariable v, VariableUpdate upd |
+    upd.(VariableAssign).getSource() = node1.asExpr() or
+    upd.(AssignOp) = node1.asExpr()
   |
-    node1.asExpr() = captured.getAUse()
-    or
-    not exists(captured.getAUse()) and
-    exists(SsaVariable capturedDef | capturedDef = captured.getAnUltimateDefinition() |
-      capturedDef.(SsaImplicitInit).isParameterDefinition(node1.asParameter()) or
-      capturedDef.(SsaExplicitUpdate).getDefiningExpr().(VariableAssign).getSource() =
-        node1.asExpr() or
-      capturedDef.(SsaExplicitUpdate).getDefiningExpr().(AssignOp) = node1.asExpr()
-    )
+    CaptureInput::capturedVarUpdate(bb, i, v, upd) and
+    c.getVariable() = v and
+    CaptureFlow::storeStep(bb, i, v, node2.(ClosureNode).getCaptureFlowNode())
   )
+  or
+  exists(Parameter p |
+    node1.asParameter() = p and
+    c.getVariable() = p and
+    CaptureFlow::parameterStoreStep(p, node2.(ClosureNode).getCaptureFlowNode())
+  )
+}
+
+private predicate captureReadStep(Node node1, ClosureContent c, Node node2) {
+  exists(BasicBlock bb, int i, CaptureInput::CapturedVariable v |
+    CaptureFlow::readStep(node1.(ClosureNode).getCaptureFlowNode(), bb, i, v) and
+    c.getVariable() = v and
+    CaptureInput::capturedVarRead(bb, i, v, node2.asExpr())
+  )
+  or
+  exists(Parameter p |
+    CaptureFlow::parameterReadStep(node1.(ClosureNode).getCaptureFlowNode(), p) and
+    c.getVariable() = p and
+    node2.(PostUpdateNode).getPreUpdateNode().asParameter() = p
+  )
+}
+
+predicate captureValueStep(Node node1, Node node2) {
+  CaptureFlow::localFlowStep(node1.(ClosureNode).getCaptureFlowNode(),
+    node2.(ClosureNode).getCaptureFlowNode())
 }
 
 /**
@@ -78,10 +165,6 @@ private predicate variableCaptureStep(Node node1, ExprNode node2) {
 predicate jumpStep(Node node1, Node node2) {
   fieldStep(node1, node2)
   or
-  variableCaptureStep(node1, node2)
-  or
-  variableCaptureStep(node1.(PostUpdateNode).getPreUpdateNode(), node2)
-  or
   any(AdditionalValueStep a).step(node1, node2) and
   node1.getEnclosingCallable() != node2.getEnclosingCallable()
   or
@@ -117,6 +200,8 @@ predicate storeStep(Node node1, ContentSet f, Node node2) {
   or
   FlowSummaryImpl::Private::Steps::summaryStoreStep(node1.(FlowSummaryNode).getSummaryNode(), f,
     node2.(FlowSummaryNode).getSummaryNode())
+  or
+  captureStoreStep(node1, f, node2)
 }
 
 /**
@@ -149,6 +234,8 @@ predicate readStep(Node node1, ContentSet f, Node node2) {
   or
   FlowSummaryImpl::Private::Steps::summaryReadStep(node1.(FlowSummaryNode).getSummaryNode(), f,
     node2.(FlowSummaryNode).getSummaryNode())
+  or
+  captureReadStep(node1, f, node2)
 }
 
 /**
@@ -447,6 +534,8 @@ ContentApprox getContentApprox(Content c) {
   or
   c instanceof MapValueContent and result = TMapValueContentApprox()
   or
+  exists(CapturedVariable v | c = TClosureContent(v) and result = TClosureContentApprox(v))
+  or
   c instanceof SyntheticFieldContent and result = TSyntheticFieldApproxContent()
 }
 

java/ql/lib/semmle/code/java/dataflow/internal/DataFlowUtil.qll

@@ -135,6 +135,10 @@ private module Cached {
 
 import Cached
 
+private predicate capturedVariableRead(Node n) {
+  n.asExpr().(RValue).getVariable() instanceof CapturedVariable
+}
+
 private predicate simpleLocalFlowStep0(Node node1, Node node2) {
   TaintTrackingUtil::forceCachingInSameStage() and
   // Variable flow steps through adjacent def-use and use-use pairs.
@@ -142,23 +146,27 @@ private predicate simpleLocalFlowStep0(Node node1, Node node2) {
     upd.getDefiningExpr().(VariableAssign).getSource() = node1.asExpr() or
     upd.getDefiningExpr().(AssignOp) = node1.asExpr()
   |
-    node2.asExpr() = upd.getAFirstUse()
+    node2.asExpr() = upd.getAFirstUse() and
+    not capturedVariableRead(node2)
   )
   or
   exists(SsaImplicitInit init |
     init.isParameterDefinition(node1.asParameter()) and
-    node2.asExpr() = init.getAFirstUse()
+    node2.asExpr() = init.getAFirstUse() and
+    not capturedVariableRead(node2)
   )
   or
   adjacentUseUse(node1.asExpr(), node2.asExpr()) and
   not exists(FieldRead fr |
     hasNonlocalValue(fr) and fr.getField().isStatic() and fr = node1.asExpr()
   ) and
-  not FlowSummaryImpl::Private::Steps::prohibitsUseUseFlow(node1, _)
+  not FlowSummaryImpl::Private::Steps::prohibitsUseUseFlow(node1, _) and
+  not capturedVariableRead(node2)
   or
   ThisFlow::adjacentThisRefs(node1, node2)
   or
-  adjacentUseUse(node1.(PostUpdateNode).getPreUpdateNode().asExpr(), node2.asExpr())
+  adjacentUseUse(node1.(PostUpdateNode).getPreUpdateNode().asExpr(), node2.asExpr()) and
+  not capturedVariableRead(node2)
   or
   ThisFlow::adjacentThisRefs(node1.(PostUpdateNode).getPreUpdateNode(), node2)
   or
@@ -185,6 +193,8 @@ private predicate simpleLocalFlowStep0(Node node1, Node node2) {
   or
   FlowSummaryImpl::Private::Steps::summaryLocalStep(node1.(FlowSummaryNode).getSummaryNode(),
     node2.(FlowSummaryNode).getSummaryNode(), true)
+  or
+  captureValueStep(node1, node2)
 }
 
 /**
@@ -256,6 +266,19 @@ class MapValueContent extends Content, TMapValueContent {
   override string toString() { result = "<map.value>" }
 }
 
+/** A captured variable. */
+class ClosureContent extends Content, TClosureContent {
+  CapturedVariable v;
+
+  ClosureContent() { this = TClosureContent(v) }
+
+  CapturedVariable getVariable() { result = v }
+
+  override DataFlowType getType() { result = getErasedRepr(v.(Variable).getType()) }
+
+  override string toString() { result = v.toString() }
+}
+
 /** A reference through a synthetic instance field. */
 class SyntheticFieldContent extends Content, TSyntheticFieldContent {
   SyntheticField s;