Python: Basic implementation of variable capture

python/ql/consistency-queries/DataFlowConsistency.ql

@@ -14,6 +14,8 @@ private module Input implements InputSig<PythonDataFlow> {
   private import Private
   private import Public
 
+  predicate postWithInFlowExclude(Node n) { n instanceof FlowSummaryNode }
+
   predicate argHasPostUpdateExclude(ArgumentNode n) {
     // TODO: Implement post-updates for *args, see tests added in https://github.com/github/codeql/pull/14936
     exists(ArgumentPosition apos | n.argumentOf(_, apos) and apos.isStarArgs(_))
@@ -44,6 +46,13 @@ private module Input implements InputSig<PythonDataFlow> {
     )
   }
 
+  predicate uniqueEnclosingCallableExclude(Node n) {
+    // We only have a selection of valid callables.
+    // For instance, we do not have classes as `DataFlowCallable`s.
+    not n.(SynthCaptureNode).getSynthesizedCaptureNode().getEnclosingCallable() instanceof Function and
+    not n.(SynthCaptureNode).getSynthesizedCaptureNode().getEnclosingCallable() instanceof Module
+  }
+
   predicate uniqueCallEnclosingCallableExclude(DataFlowCall call) {
     not exists(call.getLocation().getFile().getRelativePath())
   }
@@ -53,7 +62,7 @@ private module Input implements InputSig<PythonDataFlow> {
   }
 
   predicate multipleArgumentCallExclude(ArgumentNode arg, DataFlowCall call) {
-    // since we can have multiple DataFlowCall for a CallNode (for example if can
+    // since we can have multiple DataFlowCall for a CallNode (for example if it can
     // resolve to multiple functions), but we only make _one_ ArgumentNode for each
     // argument in the CallNode, we end up violating this consistency check in those
     // cases. (see `getCallArg` in DataFlowDispatch.qll)

python/ql/lib/change-notes/2023-12-18-support-variable-capture.md

@@ -0,0 +1,4 @@
+---
+category: majorAnalysis
+---
+* Added support for global data-flow through captured variables.

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

@@ -42,6 +42,13 @@ private import semmle.python.dataflow.new.internal.TypeTrackingImpl::CallGraphCo
 newtype TParameterPosition =
   /** Used for `self` in methods, and `cls` in classmethods. */
   TSelfParameterPosition() or
+  /**
+   * This is used for tracking flow through captured variables, and
+   * we use separate parameter/argument positions in order to distinguish
+   * "lambda self" from "normal self", as lambdas may also access outer `self`
+   * variables (through variable capture).
+   */
+  TLambdaSelfParameterPosition() or
   TPositionalParameterPosition(int index) {
     index = any(Parameter p).getPosition()
     or
@@ -78,6 +85,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) }
 
@@ -109,6 +119,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)
@@ -129,6 +141,13 @@ class ParameterPosition extends TParameterPosition {
 newtype TArgumentPosition =
   /** Used for `self` in methods, and `cls` in classmethods. */
   TSelfArgumentPosition() or
+  /**
+   * This is used for tracking flow through captured variables, and
+   * we use separate parameter/argument positions in order to distinguish
+   * "lambda self" from "normal self", as lambdas may also access outer `self`
+   * variables (through variable capture).
+   */
+  TLambdaSelfArgumentPosition() or
   TPositionalArgumentPosition(int index) {
     exists(any(CallNode c).getArg(index))
     or
@@ -153,6 +172,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) }
 
@@ -169,6 +191,8 @@ class ArgumentPosition extends TArgumentPosition {
   string toString() {
     this.isSelf() and result = "self"
     or
+    this.isLambdaSelf() and result = "lambda self"
+    or
     exists(int pos | this.isPositional(pos) and result = "position " + pos)
     or
     exists(string name | this.isKeyword(name) and result = "keyword " + name)
@@ -183,6 +207,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))
@@ -1506,6 +1532,37 @@ abstract class ParameterNodeImpl extends Node {
   }
 }
 
+/**
+ * A synthetic parameter representing the values of the variables captured
+ * by the callable being called. This parameter represents a single object
+ * where all the values are stored as attributes.
+ * This is also known as the environment part of a closure.
+ *
+ * This is used for tracking flow through captured variables.
+ */
+class SynthCapturedVariablesParameterNode extends ParameterNodeImpl,
+  TSynthCapturedVariablesParameterNode
+{
+  private Function callable;
+
+  SynthCapturedVariablesParameterNode() { this = TSynthCapturedVariablesParameterNode(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() {
@@ -1580,6 +1637,39 @@ private class SummaryPostUpdateNode extends FlowSummaryNode, PostUpdateNodeImpl
   override Node getPreUpdateNode() { result = pre }
 }
 
+/**
+ * A synthetic argument representing the values of the variables captured
+ * by the callable being called. This argument represents a single object
+ * where all the values are stored as attributes.
+ * This is also known as the environment part of a closure.
+ *
+ * This is used for tracking flow through captured variables.
+ *
+ * TODO:
+ * We might want a synthetic node here, but currently that incurs problems
+ * with non-monotonic recursion, because of the use of `resolveCall` in the
+ * char pred. This may be solvable by using
+ * `CallGraphConstruction::Make` in stead of
+ * `CallGraphConstruction::Simple::Make` appropriately.
+ */
+class CapturedVariablesArgumentNode extends CfgNode, ArgumentNode {
+  CallNode callNode;
+
+  CapturedVariablesArgumentNode() {
+    node = callNode.getFunction() and
+    exists(Function target | resolveCall(callNode, target, _) |
+      target = any(VariableCapture::CapturedVariable v).getACapturingScope()
+    )
+  }
+
+  override string toString() { result = "Capturing closure argument" }
+
+  override predicate argumentOf(DataFlowCall call, ArgumentPosition pos) {
+    callNode = 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

@@ -17,6 +17,7 @@ private import semmle.python.Frameworks
 import MatchUnpacking
 import IterableUnpacking
 import DataFlowDispatch
+import VariableCapture as VariableCapture
 
 /** Gets the callable in which this node occurs. */
 DataFlowCallable nodeGetEnclosingCallable(Node n) { result = n.getEnclosingCallable() }
@@ -474,6 +475,8 @@ predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo) {
   simpleLocalFlowStepForTypetracking(nodeFrom, nodeTo)
   or
   summaryFlowSteps(nodeFrom, nodeTo)
+  or
+  variableCaptureLocalFlowStep(nodeFrom, nodeTo)
 }
 
 /**
@@ -496,6 +499,16 @@ predicate summaryFlowSteps(Node nodeFrom, Node nodeTo) {
   IncludePostUpdateFlow<PhaseDependentFlow<summaryLocalStep/2>::step/2>::step(nodeFrom, nodeTo)
 }
 
+predicate variableCaptureLocalFlowStep(Node nodeFrom, Node nodeTo) {
+  // Blindly applying use-use flow can result in a node that steps to itself, for
+  // example in while-loops. To uphold dataflow consistency checks, we don't want
+  // that. However, we do want to allow `[post] n` to `n` (to handle while loops), so
+  // we should only do the filtering after `IncludePostUpdateFlow` has ben applied.
+  IncludePostUpdateFlow<PhaseDependentFlow<VariableCapture::valueStep/2>::step/2>::step(nodeFrom,
+    nodeTo) and
+  nodeFrom != nodeTo
+}
+
 /** `ModuleVariable`s are accessed via jump steps at runtime. */
 predicate runtimeJumpStep(Node nodeFrom, Node nodeTo) {
   // Module variable read
@@ -559,7 +572,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`.
@@ -663,6 +676,38 @@ predicate storeStep(Node nodeFrom, ContentSet c, Node nodeTo) {
   synthStarArgsElementParameterNodeStoreStep(nodeFrom, c, nodeTo)
   or
   synthDictSplatArgumentNodeStoreStep(nodeFrom, c, nodeTo)
+  or
+  VariableCapture::storeStep(nodeFrom, c, nodeTo)
+}
+
+/**
+ * A synthesized data flow node representing a closure object that tracks
+ * captured variables.
+ */
+class SynthCaptureNode extends Node, TSynthCaptureNode {
+  private VariableCapture::Flow::SynthesizedCaptureNode cn;
+
+  SynthCaptureNode() { this = TSynthCaptureNode(cn) }
+
+  /** Gets the `SynthesizedCaptureNode` that this node represents. */
+  VariableCapture::Flow::SynthesizedCaptureNode getSynthesizedCaptureNode() { result = cn }
+
+  override Scope getScope() { result = cn.getEnclosingCallable() }
+
+  override Location getLocation() { result = cn.getLocation() }
+
+  override string toString() { result = cn.toString() }
+}
+
+private class SynthCapturePostUpdateNode extends PostUpdateNodeImpl, SynthCaptureNode {
+  private SynthCaptureNode pre;
+
+  SynthCapturePostUpdateNode() {
+    VariableCapture::Flow::capturePostUpdateNode(this.getSynthesizedCaptureNode(),
+      pre.getSynthesizedCaptureNode())
+  }
+
+  override Node getPreUpdateNode() { result = pre }
 }
 
 /**
@@ -866,6 +911,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. */
@@ -995,6 +1042,10 @@ predicate nodeIsHidden(Node n) {
   n instanceof SynthDictSplatArgumentNode
   or
   n instanceof SynthDictSplatParameterNode
+  or
+  n instanceof SynthCaptureNode
+  or
+  n instanceof SynthCapturedVariablesParameterNode
 }
 
 class LambdaCallKind = Unit;
@@ -1034,6 +1085,11 @@ predicate allowParameterReturnInSelf(ParameterNode p) {
     p.(ParameterNodeImpl).isParameterOf(c, pos) and
     FlowSummaryImpl::Private::summaryAllowParameterReturnInSelf(c.asLibraryCallable(), pos)
   )
+  or
+  exists(Function f |
+    VariableCapture::Flow::heuristicAllowInstanceParameterReturnInSelf(f) and
+    p = TSynthCapturedVariablesParameterNode(f)
+  )
 }
 
 /** An approximated `Content`. */

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

@@ -117,6 +117,14 @@ newtype TNode =
   /** A synthetic node to allow flow to keyword parameters from a `**kwargs` argument. */
   TSynthDictSplatParameterNode(DataFlowCallable callable) {
     exists(ParameterPosition ppos | ppos.isKeyword(_) | exists(callable.getParameter(ppos)))
+  } or
+  /** A synthetic node representing a captured variable. */
+  TSynthCaptureNode(VariableCapture::Flow::SynthesizedCaptureNode cn) or
+  /** A synthetic node representing the heap of a function. Used for variable capture. */
+  TSynthCapturedVariablesParameterNode(Function f) {
+    f = any(VariableCapture::CapturedVariable v).getACapturingScope() and
+    // TODO: Remove this restriction when adding proper support for captured variables in the body of the function we generate for comprehensions
+    exists(TFunction(f))
   }
 
 private import semmle.python.internal.CachedStages
@@ -627,7 +635,9 @@ newtype TContent =
     exists(string input, string output | ModelOutput::relevantSummaryModel(_, _, input, output, _) |
       attr = [input, output].regexpFind("(?<=(^|\\.)Attribute\\[)[^\\]]+(?=\\])", _, _).trim()
     )
-  }
+  } or
+  /** A captured variable. */
+  TCapturedVariableContent(VariableCapture::CapturedVariable v)
 
 /**
  * A data-flow value can have associated content.
@@ -690,6 +700,18 @@ class AttributeContent extends TAttributeContent, Content {
   override string toString() { result = "Attribute " + attr }
 }
 
+/** A captured variable. */
+class CapturedVariableContent extends Content, TCapturedVariableContent {
+  private VariableCapture::CapturedVariable v;
+
+  CapturedVariableContent() { this = TCapturedVariableContent(v) }
+
+  /** Gets the captured variable. */
+  VariableCapture::CapturedVariable getVariable() { result = v }
+
+  override string toString() { result = "captured " + v }
+}
+
 /**
  * An entity that represents a set of `Content`s.
  *

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

@@ -12,13 +12,16 @@ private import DataFlowImplSpecific::Public
 module Input implements InputSig<DataFlowImplSpecific::PythonDataFlow> {
   class SummarizedCallableBase = string;
 
-  ArgumentPosition callbackSelfParameterPosition() { none() }
+  ArgumentPosition callbackSelfParameterPosition() { result.isLambdaSelf() }
 
   ReturnKind getStandardReturnValueKind() { any() }
 
   string encodeParameterPosition(ParameterPosition pos) {
     pos.isSelf() and result = "self"
     or
+    pos.isLambdaSelf() and
+    result = "lambda-self"
+    or
     exists(int i |
       pos.isPositional(i) and
       result = i.toString()
@@ -33,6 +36,9 @@ module Input implements InputSig<DataFlowImplSpecific::PythonDataFlow> {
   string encodeArgumentPosition(ArgumentPosition pos) {
     pos.isSelf() and result = "self"
     or
+    pos.isLambdaSelf() and
+    result = "lambda-self"
+    or
     exists(int i |
       pos.isPositional(i) and
       result = i.toString()