Variable capture: allow arbitrary data-flow nodes to be the source of a write

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

@@ -113,11 +113,6 @@ private module CaptureInput implements VariableCapture::InputSig {
     VariableWrite() { super.getDestVar() = v }
 
     CapturedVariable getVariable() { result = v }
-
-    Expr getSource() {
-      result = this.(VariableAssign).getSource() or
-      result = this.(AssignOp)
-    }
   }
 
   class VariableRead extends Expr instanceof RValue {
@@ -155,14 +150,27 @@ class CapturedParameter = CaptureInput::CapturedParameter;
 module CaptureFlow = VariableCapture::Flow<CaptureInput>;
 
 private CaptureFlow::ClosureNode asClosureNode(Node n) {
-  result = n.(CaptureNode).getSynthesizedCaptureNode() or
-  result.(CaptureFlow::ExprNode).getExpr() = n.asExpr() or
+  result = n.(CaptureNode).getSynthesizedCaptureNode()
+  or
+  result.(CaptureFlow::ExprNode).getExpr() = n.asExpr()
+  or
   result.(CaptureFlow::ExprPostUpdateNode).getExpr() =
-    n.(PostUpdateNode).getPreUpdateNode().asExpr() or
-  result.(CaptureFlow::ParameterNode).getParameter() = n.asParameter() or
-  result.(CaptureFlow::ThisParameterNode).getCallable() = n.(InstanceParameterNode).getCallable() or
+    n.(PostUpdateNode).getPreUpdateNode().asExpr()
+  or
+  result.(CaptureFlow::ParameterNode).getParameter() = n.asParameter()
+  or
+  result.(CaptureFlow::ThisParameterNode).getCallable() = n.(InstanceParameterNode).getCallable()
+  or
   exprNode(result.(CaptureFlow::MallocNode).getClosureExpr()).(PostUpdateNode).getPreUpdateNode() =
     n
+  or
+  exists(CaptureInput::VariableWrite write |
+    result.(CaptureFlow::VariableWriteSourceNode).getVariableWrite() = write
+  |
+    n.asExpr() = write.(VariableAssign).getSource()
+    or
+    n.asExpr() = write.(AssignOp)
+  )
 }
 
 private predicate captureStoreStep(Node node1, CapturedVariableContent c, Node node2) {

shared/dataflow/codeql/dataflow/VariableCapture.qll

@@ -95,9 +95,6 @@ signature module InputSig {
     /** Gets the variable that is the target of this write. */
     CapturedVariable getVariable();
 
-    /** Gets the expression that is the source of this write. */
-    Expr getSource();
-
     /** Gets the location of this write. */
     Location getLocation();
 
@@ -210,6 +207,22 @@ signature module OutputSig<InputSig I> {
     I::ClosureExpr getClosureExpr();
   }
 
+  /**
+   * A node representing the incoming value about to be written at the given assignment.
+   *
+   * The captured-variable library will generate flows out of this node, and assume that other
+   * parts of the language implementation produce the relevant data flows into this node.
+   *
+   * For ordinary assignments, this could be mapped to the right-hand side of the assignment.
+   *
+   * For more general cases, where an lvalue has no direct corresponding rvalue, this can be mapped
+   * to a data-flow node that wraps the lvalue, with language-specific incoming data flows.
+   */
+  class VariableWriteSourceNode extends ClosureNode {
+    /** Gets the variable write for which this node is the incoming value being written to the variable. */
+    I::VariableWrite getVariableWrite();
+  }
+
   /** Holds if `post` is a `PostUpdateNode` for `pre`. */
   predicate capturePostUpdateNode(SynthesizedCaptureNode post, SynthesizedCaptureNode pre);
 
@@ -239,7 +252,6 @@ module Flow<InputSig Input> implements OutputSig<Input> {
     private class RelevantExpr extends FinalExpr {
       RelevantExpr() {
         this instanceof VariableRead or
-        any(VariableWrite vw).getSource() = this or
         this instanceof ClosureExpr or
         any(ClosureExpr ce).hasAliasedAccess(this)
       }
@@ -353,14 +365,6 @@ module Flow<InputSig Input> implements OutputSig<Input> {
 
     query predicate uniqueWriteTarget(string msg) { uniqueWriteTarget(_, msg) }
 
-    private predicate uniqueWriteSource(VariableWrite vw, string msg) {
-      msg = "VariableWrite has no source expression" and not exists(vw.getSource())
-      or
-      msg = "VariableWrite has multiple source expressions" and 2 <= strictcount(vw.getSource())
-    }
-
-    query predicate uniqueWriteSource(string msg) { uniqueWriteSource(_, msg) }
-
     private predicate uniqueWriteCfgNode(VariableWrite vw, string msg) {
       msg = "VariableWrite has no cfg node" and not vw.hasCfgNode(_, _)
       or
@@ -370,17 +374,6 @@ module Flow<InputSig Input> implements OutputSig<Input> {
 
     query predicate uniqueWriteCfgNode(string msg) { uniqueWriteCfgNode(_, msg) }
 
-    private predicate localWriteStep(VariableWrite vw, string msg) {
-      exists(BasicBlock bb1, BasicBlock bb2 |
-        vw.hasCfgNode(bb1, _) and
-        vw.getSource().hasCfgNode(bb2, _) and
-        bb1.getEnclosingCallable() != bb2.getEnclosingCallable() and
-        msg = "VariableWrite is not a local step"
-      )
-    }
-
-    query predicate localWriteStep(string msg) { localWriteStep(_, msg) }
-
     query predicate uniqueReadVariable(VariableRead vr, string msg) {
       msg = "VariableRead has no source variable" and not exists(vr.getVariable())
       or
@@ -435,9 +428,7 @@ module Flow<InputSig Input> implements OutputSig<Input> {
       n = strictcount(Expr e | uniqueLocation(e, msg)) or
       n = strictcount(Expr e | uniqueCfgNode(e, msg)) or
       n = strictcount(VariableWrite vw | uniqueWriteTarget(vw, msg)) or
-      n = strictcount(VariableWrite vw | uniqueWriteSource(vw, msg)) or
       n = strictcount(VariableWrite vw | uniqueWriteCfgNode(vw, msg)) or
-      n = strictcount(VariableWrite vw | localWriteStep(vw, msg)) or
       n = strictcount(VariableRead vr | uniqueReadVariable(vr, msg)) or
       n = strictcount(ClosureExpr ce | closureMustHaveBody(ce, msg)) or
       n = strictcount(ClosureExpr ce, Expr access | closureAliasMustBeLocal(ce, access, msg)) or
@@ -689,13 +680,12 @@ module Flow<InputSig Input> implements OutputSig<Input> {
     TExprNode(Expr expr, boolean isPost) {
       expr instanceof VariableRead and isPost = [false, true]
       or
-      exists(VariableWrite vw | expr = vw.getSource() and isPost = false)
-      or
       synthRead(_, _, _, _, expr) and isPost = [false, true]
     } or
     TParamNode(CapturedParameter p) or
     TThisParamNode(Callable c) { captureAccess(_, c) } or
-    TMallocNode(ClosureExpr ce) { hasConstructorCapture(ce, _) }
+    TMallocNode(ClosureExpr ce) { hasConstructorCapture(ce, _) } or
+    TVariableWriteSourceNode(VariableWrite write)
 
   class ClosureNode extends TClosureNode {
     /** Gets a textual representation of this node. */
@@ -721,6 +711,11 @@ module Flow<InputSig Input> implements OutputSig<Input> {
       result = "this" and this = TThisParamNode(_)
       or
       result = "malloc" and this = TMallocNode(_)
+      or
+      exists(VariableWrite write |
+        this = TVariableWriteSourceNode(write) and
+        result = "Source of write to " + write.getVariable().toString()
+      )
     }
 
     /** Gets the location of this node. */
@@ -748,6 +743,10 @@ module Flow<InputSig Input> implements OutputSig<Input> {
       exists(Callable c | this = TThisParamNode(c) and result = c.getLocation())
       or
       exists(ClosureExpr ce | this = TMallocNode(ce) and result = ce.getLocation())
+      or
+      exists(VariableWrite write |
+        this = TVariableWriteSourceNode(write) and result = write.getLocation()
+      )
     }
   }
 
@@ -807,6 +806,10 @@ module Flow<InputSig Input> implements OutputSig<Input> {
     ClosureExpr getClosureExpr() { this = TMallocNode(result) }
   }
 
+  class VariableWriteSourceNode extends ClosureNode, TVariableWriteSourceNode {
+    VariableWrite getVariableWrite() { this = TVariableWriteSourceNode(result) }
+  }
+
   predicate capturePostUpdateNode(SynthesizedCaptureNode post, SynthesizedCaptureNode pre) {
     exists(CapturedVariable v, BasicBlock bb, int i |
       pre = TSynthRead(v, bb, i, false) and post = TSynthRead(v, bb, i, true)
@@ -851,7 +854,7 @@ module Flow<InputSig Input> implements OutputSig<Input> {
     or
     exists(VariableWrite vw, CapturedVariable v |
       captureWrite(v, bb, i, true, vw) and
-      n = TExprNode(vw.getSource(), false) and
+      n = TVariableWriteSourceNode(vw) and
       isPost = false and
       cc = TVariable(v)
     )
@@ -898,7 +901,7 @@ module Flow<InputSig Input> implements OutputSig<Input> {
     // write to v inside the closure body
     exists(BasicBlock bb, int i, VariableWrite vw |
       captureWrite(v, bb, i, false, vw) and
-      node1 = TExprNode(vw.getSource(), false) and
+      node1 = TVariableWriteSourceNode(vw) and
       node2 = TSynthThisQualifier(bb, i, true)
     )
   }