Python: use synthetic node for comprehension capture argument

We used to use the CfgNode for the comprehension itself.
In cases where that is also an argument, say
```python
",".join([x for x in l])
```
that would be an argument to two different calls causing a dataflow consistency violation.

Author: yoff

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

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

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

@@ -127,6 +127,11 @@ module Flow = Shared::Flow<Location, CaptureInput>;
 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()