diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll b/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll index 58fc674f48d..c15f1751edb 100644 --- a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll +++ b/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll @@ -1729,24 +1729,34 @@ class CapturedVariablesArgumentNode extends CfgNode, ArgumentNode { } } -class ComprehensionCapturedVariablesArgumentNode extends CfgNode, ArgumentNode { +class ComprehensionCapturedVariablesArgumentNode extends Node, ArgumentNode { Comp comp; ComprehensionCapturedVariablesArgumentNode() { - node.getNode() = comp and + this = TSynthCompCapturedVariablesArgumentNode(comp) and exists(Function target | target = comp.getFunction() | target = any(VariableCapture::CapturedVariable v).getACapturingScope() ) } - override string toString() { result = "Capturing closure argument (comp)" } - override predicate argumentOf(DataFlowCall call, ArgumentPosition pos) { call.(ComprehensionCall).getComprehension() = comp and pos.isLambdaSelf() } } +class SynthCompCapturedVariablesArgumentNode extends Node, TSynthCompCapturedVariablesArgumentNode { + Comp comp; + + SynthCompCapturedVariablesArgumentNode() { this = TSynthCompCapturedVariablesArgumentNode(comp) } + + override string toString() { result = "Capturing closure argument (comp)" } + + override Scope getScope() { result = comp.getFunction() } + + override Location getLocation() { result = comp.getLocation() } +} + /** Gets a viable run-time target for the call `call`. */ DataFlowCallable viableCallable(DataFlowCall call) { call instanceof ExtractedDataFlowCall and diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPublic.qll b/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPublic.qll index fa9e844eda3..eca9da79791 100644 --- a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPublic.qll +++ b/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPublic.qll @@ -124,9 +124,11 @@ newtype TNode = /** 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)) } or + TSynthCompCapturedVariablesArgumentNode(Comp comp) { + comp.getFunction() = any(VariableCapture::CapturedVariable v).getACapturingScope() + } or /** An empty, unused node type that exists to prevent unwanted dependencies on data flow nodes. */ TForbiddenRecursionGuard() { none() and diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/VariableCapture.qll b/python/ql/lib/semmle/python/dataflow/new/internal/VariableCapture.qll index 86dcee2bfe3..6cb80881e2a 100644 --- a/python/ql/lib/semmle/python/dataflow/new/internal/VariableCapture.qll +++ b/python/ql/lib/semmle/python/dataflow/new/internal/VariableCapture.qll @@ -127,6 +127,11 @@ module Flow = Shared::Flow; private Flow::ClosureNode asClosureNode(Node n) { result = n.(SynthCaptureNode).getSynthesizedCaptureNode() or + exists(Comp comp | n = TSynthCompCapturedVariablesArgumentNode(comp) | + result.(Flow::ExprNode).getExpr().getNode() = comp + ) + or + // TODO: Should the `Comp`s above be excluded here? result.(Flow::ExprNode).getExpr() = n.(CfgNode).getNode() or result.(Flow::VariableWriteSourceNode).getVariableWrite() = n.(CfgNode).getNode()