Swift: Fix 'parameter' -> 'argument' flow into closures.

swift/ql/lib/codeql/swift/controlflow/CfgNodes.qll

@@ -173,7 +173,7 @@ class ApplyExprCfgNode extends ExprCfgNode {
 
   Callable getStaticTarget() { result = e.getStaticTarget() }
 
-  Expr getFunction() { result = e.getFunction() }
+  CfgNode getFunction() { result.getAst() = e.getFunction() }
 }
 
 class CallExprCfgNode extends ApplyExprCfgNode {

swift/ql/lib/codeql/swift/dataflow/internal/DataFlowDispatch.qll

@@ -293,12 +293,14 @@ private module Cached {
   newtype TArgumentPosition =
     TThisArgument() or
     // we rely on default exprs generated in the caller for ordering
-    TPositionalArgument(int n) { n = any(Argument arg).getIndex() }
+    TPositionalArgument(int n) { n = any(Argument arg).getIndex() } or
+    TClosureSelfArgument()
 
   cached
   newtype TParameterPosition =
     TThisParameter() or
-    TPositionalParameter(int n) { n = any(Argument arg).getIndex() }
+    TPositionalParameter(int n) { n = any(Argument arg).getIndex() } or
+    TClosureSelfParameter()
 }
 
 import Cached
@@ -331,6 +333,10 @@ class ThisParameterPosition extends ParameterPosition, TThisParameter {
   override string toString() { result = "this" }
 }
 
+class ClosureSelfParameter extends ParameterPosition, TClosureSelfParameter {
+  override string toString() { result = "{ ... }" }
+}
+
 /** An argument position. */
 class ArgumentPosition extends TArgumentPosition {
   /** Gets a textual representation of this position. */
@@ -347,6 +353,10 @@ class ThisArgumentPosition extends ArgumentPosition, TThisArgument {
   override string toString() { result = "this" }
 }
 
+class ClosureSelfArgument extends ArgumentPosition, TClosureSelfArgument {
+  override string toString() { result = "{ ... }" }
+}
+
 /** Holds if arguments at position `apos` match parameters at position `ppos`. */
 pragma[inline]
 predicate parameterMatch(ParameterPosition ppos, ArgumentPosition apos) {
@@ -354,4 +364,7 @@ predicate parameterMatch(ParameterPosition ppos, ArgumentPosition apos) {
   apos instanceof TThisArgument
   or
   ppos.(PositionalParameterPosition).getIndex() = apos.(PositionalArgumentPosition).getIndex()
+  or
+  ppos instanceof TClosureSelfParameter and
+  apos instanceof TClosureSelfArgument
 }

swift/ql/lib/codeql/swift/dataflow/internal/DataFlowPrivate.qll

@@ -382,6 +382,12 @@ private class DictionarySubscriptNode extends NodeImpl, TDictionarySubscriptNode
   SubscriptExpr getExpr() { result = expr }
 }
 
+private class ClosureSelfReferenceNode extends ExprNodeImpl {
+  override ClosureExpr expr;
+
+  ClosureExpr getClosure() { result = expr }
+}
+
 private module ParameterNodes {
   abstract class ParameterNodeImpl extends NodeImpl {
     predicate isParameterOf(DataFlowCallable c, ParameterPosition pos) { none() }
@@ -412,6 +418,13 @@ private module ParameterNodes {
     override ParamDecl getParameter() { result = param }
   }
 
+  class ClosureSelfParameterNode extends ParameterNodeImpl, ClosureSelfReferenceNode {
+    override predicate isParameterOf(DataFlowCallable c, ParameterPosition pos) {
+      c.asSourceCallable() = this.getClosure() and
+      pos instanceof ClosureSelfParameter
+    }
+  }
+
   class SummaryParameterNode extends ParameterNodeImpl, FlowSummaryNode {
     SummaryParameterNode() {
       FlowSummaryImpl::Private::summaryParameterNode(this.getSummaryNode(), _)
@@ -626,6 +639,18 @@ private module ArgumentNodes {
       pos = TPositionalArgument(0)
     }
   }
+
+  class SelfClosureArgumentNode extends ExprNode, ArgumentNode {
+    ApplyExprCfgNode apply;
+
+    SelfClosureArgumentNode() { n = apply.getFunction() }
+
+    override predicate argumentOf(DataFlowCall call, ArgumentPosition pos) {
+      apply = call.asCall() and
+      not exists(apply.getStaticTarget()) and
+      pos instanceof ClosureSelfArgument
+    }
+  }
 }
 
 import ArgumentNodes
@@ -878,8 +903,8 @@ private module CaptureInput implements VariableCapture::InputSig {
         source = a.getSource()
       )
       or
-      exists(S::PatternBindingDecl pbd, S::NamedPattern np | this = pbd and pbd.getAPattern() = np |
-        np.getVarDecl() = variable and
+      exists(S::NamedPattern np | this = np |
+        variable = np.getVarDecl() and
         source = np.getMatchingExpr()
       )
       // TODO: support multiple variables in LHS of =, in both of above cases.
@@ -918,13 +943,23 @@ 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.getParameter() or
-  result.(CaptureFlow::ThisParameterNode).getCallable().getSelfParam() = n.getParameter() or
+    n.(PostUpdateNode).getPreUpdateNode().asExpr()
+  or
+  result.(CaptureFlow::ParameterNode).getParameter() = n.getParameter()
+  or
+  result.(CaptureFlow::ThisParameterNode).getCallable() = n.(ClosureSelfParameterNode).getClosure()
+  or
   result.(CaptureFlow::MallocNode).getClosureExpr() = n.getCfgNode().getNode().asAstNode() // TODO: figure out why the java version had PostUpdateNode logic here
+  or
+  exists(CaptureInput::VariableWrite write |
+    result.(CaptureFlow::VariableWriteSourceNode).getVariableWrite() = write and
+    n.asExpr() = write.getSource()
+  )
 }
 
 private predicate captureStoreStep(Node node1, Content::CapturedVariableContent c, Node node2) {