python: instantiate module for variable capture

This provides variable capture in standard situations:
- nested functions
- lambdas
There are some deficiencies:
- we do not yet handle objects capturing variables.
- we do not handle variables captured via the `nonlocal` keyword.
  This should be solved at the AST level, though, and then it
  should "just work".

Author: yoff

python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll

@@ -43,6 +43,7 @@ private import semmle.python.dataflow.new.internal.TypeTracker::CallGraphConstru
 newtype TParameterPosition =
   /** Used for `self` in methods, and `cls` in classmethods. */
   TSelfParameterPosition() or
+  TLambdaSelfParameterPosition() or
   TPositionalParameterPosition(int index) {
     index = any(Parameter p).getPosition()
     or
@@ -79,6 +80,9 @@ class ParameterPosition extends TParameterPosition {
   /** Holds if this position represents a `self`/`cls` parameter. */
   predicate isSelf() { this = TSelfParameterPosition() }
 
+  /** Holds if this position represents a reference to a lambda itself. Only used for tracking flow through captured variables. */
+  predicate isLambdaSelf() { this = TLambdaSelfParameterPosition() }
+
   /** Holds if this position represents a positional parameter at (0-based) `index`. */
   predicate isPositional(int index) { this = TPositionalParameterPosition(index) }
 
@@ -110,6 +114,8 @@ class ParameterPosition extends TParameterPosition {
   string toString() {
     this.isSelf() and result = "self"
     or
+    this.isLambdaSelf() and result = "lambda self"
+    or
     exists(int index | this.isPositional(index) and result = "position " + index)
     or
     exists(string name | this.isKeyword(name) and result = "keyword " + name)
@@ -130,6 +136,7 @@ class ParameterPosition extends TParameterPosition {
 newtype TArgumentPosition =
   /** Used for `self` in methods, and `cls` in classmethods. */
   TSelfArgumentPosition() or
+  TLambdaSelfArgumentPosition() or
   TPositionalArgumentPosition(int index) {
     exists(any(CallNode c).getArg(index))
     or
@@ -154,6 +161,9 @@ class ArgumentPosition extends TArgumentPosition {
   /** Holds if this position represents a `self`/`cls` argument. */
   predicate isSelf() { this = TSelfArgumentPosition() }
 
+  /** Holds if this position represents a lambda `self` argument. Only used for tracking flow through captured variables. */
+  predicate isLambdaSelf() { this = TLambdaSelfArgumentPosition() }
+
   /** Holds if this position represents a positional argument at (0-based) `index`. */
   predicate isPositional(int index) { this = TPositionalArgumentPosition(index) }
 
@@ -184,6 +194,8 @@ class ArgumentPosition extends TArgumentPosition {
 predicate parameterMatch(ParameterPosition ppos, ArgumentPosition apos) {
   ppos.isSelf() and apos.isSelf()
   or
+  ppos.isLambdaSelf() and apos.isLambdaSelf()
+  or
   exists(int index | ppos.isPositional(index) and apos.isPositional(index))
   or
   exists(string name | ppos.isKeyword(name) and apos.isKeyword(name))
@@ -1507,6 +1519,36 @@ abstract class ParameterNodeImpl extends Node {
   }
 }
 
+/**
+ * The value of a lambda itself at function entry, viewed as a node in a data
+ * flow graph.
+ *
+ * This is used for tracking flow through captured variables, and we use a
+ * separate node and parameter/argument positions in order to distinguish
+ * "lambda self" from "normal self", as lambdas may also access outer `self`
+ * variables (through variable capture).
+ */
+class LambdaSelfReferenceNode extends ParameterNodeImpl, TLambdaSelfReferenceNode {
+  private Function callable;
+
+  LambdaSelfReferenceNode() { this = TLambdaSelfReferenceNode(callable) }
+
+  final Function getCallable() { result = callable }
+
+  override Parameter getParameter() { none() }
+
+  override predicate isParameterOf(DataFlowCallable c, ParameterPosition pos) {
+    c = TFunction(callable) and
+    pos.isLambdaSelf()
+  }
+
+  override Scope getScope() { result = callable }
+
+  override Location getLocation() { result = callable.getLocation() }
+
+  override string toString() { result = "lambda self in " + callable }
+}
+
 /** A parameter for a library callable with a flow summary. */
 class SummaryParameterNode extends ParameterNodeImpl, FlowSummaryNode {
   SummaryParameterNode() {
@@ -1578,6 +1620,28 @@ private class SummaryPostUpdateNode extends FlowSummaryNode, PostUpdateNodeImpl
   override Node getPreUpdateNode() { result = pre }
 }
 
+private class CapturePostUpdateNode extends PostUpdateNodeImpl, CaptureNode {
+  private CaptureNode pre;
+
+  CapturePostUpdateNode() {
+    VariableCapture::Flow::capturePostUpdateNode(this.getSynthesizedCaptureNode(),
+      pre.getSynthesizedCaptureNode())
+  }
+
+  override Node getPreUpdateNode() { result = pre }
+}
+
+private class CaptureArgumentNode extends CfgNode, ArgumentNode {
+  CallCfgNode callNode;
+
+  CaptureArgumentNode() { this = callNode.getFunction() }
+
+  override predicate argumentOf(DataFlowCall call, ArgumentPosition pos) {
+    callNode.asCfgNode() = call.getNode() and
+    pos.isLambdaSelf()
+  }
+}
+
 /** Gets a viable run-time target for the call `call`. */
 DataFlowCallable viableCallable(DataFlowCall call) {
   call instanceof ExtractedDataFlowCall and

python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPrivate.qll

@@ -378,6 +378,167 @@ module LocalFlow {
   }
 }
 
+/** Provides logic related to captured variables. */
+module VariableCapture {
+  private import codeql.dataflow.VariableCapture as Shared
+
+  class ExprCfgNode extends ControlFlowNode {
+    ExprCfgNode() { isExpressionNode(this) }
+  }
+
+  private predicate closureFlowStep(ExprCfgNode e1, ExprCfgNode e2) {
+    // simpleAstFlowStep(e1, e2)
+    // or
+    exists(SsaVariable def |
+      def.getAUse() = e2 and
+      def.getAnUltimateDefinition().getDefinition().(DefinitionNode).getValue() = e1
+    )
+  }
+
+  private module CaptureInput implements Shared::InputSig {
+    private import python as P
+
+    class Location = P::Location;
+
+    class BasicBlock extends P::BasicBlock {
+      Callable getEnclosingCallable() { result = this.getScope() }
+
+      // TODO: check that this gives useful results
+      Location getLocation() { result = super.getNode(0).getLocation() }
+    }
+
+    BasicBlock getImmediateBasicBlockDominator(BasicBlock bb) {
+      result = bb.getImmediateDominator()
+    }
+
+    BasicBlock getABasicBlockSuccessor(BasicBlock bb) { result = bb.getASuccessor() }
+
+    class CapturedVariable extends LocalVariable {
+      Function f;
+
+      CapturedVariable() {
+        //  exists(this.getScopeEntryDefinition())
+        this.getScope() = f and
+        (
+          this.getALoad().getScope() != f
+          or
+          this.getAStore().getScope() != f
+        )
+      }
+
+      Callable getCallable() { result = f }
+
+      Location getLocation() { result = f.getLocation() }
+    }
+
+    class CapturedParameter extends CapturedVariable {
+      CapturedParameter() { this.isParameter() }
+
+      ControlFlowNode getCfgNode() { result.getNode().(Parameter) = this.getAnAccess() }
+    }
+
+    class Expr extends ExprCfgNode {
+      predicate hasCfgNode(BasicBlock bb, int i) { this = bb.getNode(i) }
+    }
+
+    class VariableWrite extends ControlFlowNode {
+      CapturedVariable v;
+
+      VariableWrite() { this = v.getAStore().getAFlowNode() }
+
+      CapturedVariable getVariable() { result = v }
+
+      predicate hasCfgNode(BasicBlock bb, int i) { this = bb.getNode(i) }
+    }
+
+    class VariableRead extends Expr {
+      CapturedVariable v;
+
+      VariableRead() { this = v.getALoad().getAFlowNode() }
+
+      CapturedVariable getVariable() { result = v }
+    }
+
+    class ClosureExpr extends Expr {
+      ClosureExpr() { this.getNode() instanceof CallableExpr }
+
+      predicate hasBody(Callable body) { body = this.getNode().(CallableExpr).getInnerScope() }
+
+      predicate hasAliasedAccess(Expr f) { closureFlowStep+(this, f) and not closureFlowStep(f, _) }
+    }
+
+    // TODO: Some basic blocks will not have an enclosing callable
+    // as their scope are not `Function`s. This leads to inconsistency failures.
+    class Callable extends Scope {
+      predicate isConstructor() { none() }
+    }
+  }
+
+  class CapturedVariable = CaptureInput::CapturedVariable;
+
+  class ClosureExpr = CaptureInput::ClosureExpr;
+
+  module Flow = Shared::Flow<CaptureInput>;
+
+  private Flow::ClosureNode asClosureNode(Node n) {
+    result = n.(CaptureNode).getSynthesizedCaptureNode()
+    or
+    result.(Flow::ExprNode).getExpr() = n.(CfgNode).getNode()
+    or
+    result.(Flow::VariableWriteSourceNode).getVariableWrite() = n.(CfgNode).getNode()
+    or
+    result.(Flow::ExprPostUpdateNode).getExpr() =
+      n.(PostUpdateNode).getPreUpdateNode().(CfgNode).getNode()
+    or
+    result.(Flow::ParameterNode).getParameter().getCfgNode() = n.(CfgNode).getNode()
+    or
+    result.(Flow::ThisParameterNode).getCallable() = n.(LambdaSelfReferenceNode).getCallable()
+  }
+
+  predicate storeStep(Node nodeFrom, CapturedVariableContent c, Node nodeTo) {
+    Flow::storeStep(asClosureNode(nodeFrom), c.getVariable(), asClosureNode(nodeTo))
+  }
+
+  predicate readStep(Node nodeFrom, CapturedVariableContent c, Node nodeTo) {
+    Flow::readStep(asClosureNode(nodeFrom), c.getVariable(), asClosureNode(nodeTo))
+  }
+
+  predicate valueStep(Node nodeFrom, Node nodeTo) {
+    Flow::localFlowStep(asClosureNode(nodeFrom), asClosureNode(nodeTo))
+  }
+
+  // Note: Learn from JS, https://github.com/github/codeql/pull/14412
+  // - JS: Capture flow
+  // - JS: Disallow consecutive captured contents
+  private module Debug {
+    predicate flowStoreStep(
+      Node nodeFrom, Flow::ClosureNode closureNodeFrom, CapturedVariable v,
+      Flow::ClosureNode closureNodeTo, Node nodeTo
+    ) {
+      closureNodeFrom = asClosureNode(nodeFrom) and
+      closureNodeTo = asClosureNode(nodeTo) and
+      Flow::storeStep(closureNodeFrom, v, closureNodeTo)
+    }
+
+    predicate flowReadStep(
+      Node nodeFrom, Flow::ClosureNode closureNodeFrom, CapturedVariable v,
+      Flow::ClosureNode closureNodeTo, Node nodeTo
+    ) {
+      closureNodeFrom = asClosureNode(nodeFrom) and
+      closureNodeTo = asClosureNode(nodeTo) and
+      Flow::readStep(closureNodeFrom, v, closureNodeTo)
+    }
+
+    predicate flowValueStep(
+      Node nodeFrom, Flow::ClosureNode closureNodeFrom, Flow::ClosureNode closureNodeTo, Node nodeTo
+    ) {
+      closureNodeFrom = asClosureNode(nodeFrom) and
+      closureNodeTo = asClosureNode(nodeTo) and
+      Flow::localFlowStep(closureNodeFrom, closureNodeTo)
+    }
+  }
+}
+
 //--------
 // Local flow
 //--------
@@ -471,6 +632,8 @@ predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo) {
   simpleLocalFlowStepForTypetracking(nodeFrom, nodeTo)
   or
   summaryFlowSteps(nodeFrom, nodeTo)
+  or
+  variableCaptureFlowStep(nodeFrom, nodeTo)
 }
 
 /**
@@ -494,6 +657,11 @@ predicate summaryFlowSteps(Node nodeFrom, Node nodeTo) {
   IncludePostUpdateFlow<PhaseDependentFlow<summaryLocalStep/2>::step/2>::step(nodeFrom, nodeTo)
 }
 
+predicate variableCaptureFlowStep(Node nodeFrom, Node nodeTo) {
+  IncludePostUpdateFlow<PhaseDependentFlow<VariableCapture::valueStep/2>::step/2>::step(nodeFrom,
+    nodeTo)
+}
+
 /** `ModuleVariable`s are accessed via jump steps at runtime. */
 predicate runtimeJumpStep(Node nodeFrom, Node nodeTo) {
   // Module variable read
@@ -557,7 +725,7 @@ predicate compatibleTypes(DataFlowType t1, DataFlowType t2) { any() }
 
 predicate typeStrongerThan(DataFlowType t1, DataFlowType t2) { none() }
 
-predicate localMustFlowStep(Node node1, Node node2) { none() }
+predicate localMustFlowStep(Node nodeFrom, Node nodeTo) { none() }
 
 /**
  * Gets the type of `node`.
@@ -661,6 +829,8 @@ predicate storeStep(Node nodeFrom, ContentSet c, Node nodeTo) {
   synthStarArgsElementParameterNodeStoreStep(nodeFrom, c, nodeTo)
   or
   synthDictSplatArgumentNodeStoreStep(nodeFrom, c, nodeTo)
+  or
+  VariableCapture::storeStep(nodeFrom, c, nodeTo)
 }
 
 /**
@@ -864,6 +1034,8 @@ predicate readStep(Node nodeFrom, ContentSet c, Node nodeTo) {
     nodeTo.(FlowSummaryNode).getSummaryNode())
   or
   synthDictSplatParameterNodeReadStep(nodeFrom, c, nodeTo)
+  or
+  VariableCapture::readStep(nodeFrom, c, nodeTo)
 }
 
 /** Data flows from a sequence to a subscript of the sequence. */
@@ -993,6 +1165,8 @@ predicate nodeIsHidden(Node n) {
   n instanceof SynthDictSplatArgumentNode
   or
   n instanceof SynthDictSplatParameterNode
+  // or
+  // n instanceof CaptureNode
 }
 
 class LambdaCallKind = Unit;
@@ -1029,6 +1203,11 @@ predicate additionalLambdaFlowStep(Node nodeFrom, Node nodeTo, boolean preserves
  */
 predicate allowParameterReturnInSelf(ParameterNode p) {
   FlowSummaryImpl::Private::summaryAllowParameterReturnInSelf(p)
+  or
+  exists(Function f |
+    VariableCapture::Flow::heuristicAllowInstanceParameterReturnInSelf(f) and
+    p = TLambdaSelfReferenceNode(f)
+  )
 }
 
 /** An approximated `Content`. */