Unified: captured variables

shared/dataflow/codeql/dataflow/VariableCapture.qll

@@ -483,6 +483,7 @@ module Flow<
   }
 
   /** Gets the enclosing callable of `ce`. */
+  pragma[nomagic]
   private Callable closureExprGetEnclosingCallable(ClosureExpr ce) {
     exists(BasicBlock bb | ce.hasCfgNode(bb, _) and result = bb.getEnclosingCallable())
   }
@@ -496,6 +497,13 @@ module Flow<
     )
   }
 
+  /** Holds if `outer` contains or equals `inner` */
+  bindingset[outer, inner]
+  pragma[inline_late]
+  private predicate isEnclosingCallable(Callable outer, Callable inner) {
+    outer = callableGetEnclosingCallable*(inner)
+  }
+
   /**
    * Gets a callable that contains `ce`, or a reference to `ce` into which `ce` could be inlined without
    * bringing any variables out of scope.
@@ -512,7 +520,7 @@ module Flow<
       expr.hasCfgNode(bb, _) and
       result = bb.getEnclosingCallable() and
       // The reference to `ce` is allowed to occur in a more deeply nested context
-      closureExprGetEnclosingCallable(ce) = callableGetEnclosingCallable*(result)
+      isEnclosingCallable(closureExprGetEnclosingCallable(ce), result)
     )
   }
 
@@ -812,6 +820,16 @@ module Flow<
       or
       exists(SsaFlow::SsaNode n | this = TSynthSsa(n) and n.getSourceVariable() = TThis(_))
     }
+
+    predicate hasCfgNode(BasicBlock bb, int i) {
+      this = TSynthRead(_, bb, i, _)
+      or
+      this = TSynthThisQualifier(bb, i, _)
+      or
+      exists(SsaFlow::SsaNode n |
+        this = TSynthSsa(n) and n.getBasicBlock() = bb and n.getIndex() = i
+      )
+    }
   }
 
   class ExprNode extends ClosureNode, TExprNode {

unified/extractor/src/languages/swift/swift.rs

@@ -823,10 +823,8 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
                 return_type: {ret}
                 body: (block stmt: {body}))
         ),
-        // A closure capture (`[weak self]`, `[x]`, `[y = expr]`). The optional
-        // ownership specifier (`weak`/`unowned`) becomes a modifier; the
-        // captured name becomes the bound `name_node`; an explicit capture
-        // initializer (`[y = expr]`) becomes the bound value.
+        // A closure capture (`[weak self]`, `[x]`, `[y = expr]`).
+        // Normalize so there is always an explicit initializer, expanding `[x]` into `[x = x]`.
         rule!(
             (closureCapture
                 specifier: (closureCaptureSpecifier specifier: @@spec)?
@@ -836,7 +834,13 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
             (variable_declaration
                 modifier: (modifier #{spec})?
                 pattern: (identifier #{name})
-                value: {val})
+                value: {
+                    // Expand [x] into [x = x]
+                    match val {
+                        Some(val) => val,
+                        None => tree!((identifier #{name})),
+                    }
+                })
         ),
         // A closure parameter clause (`(x: Int, y)`) unwraps to its parameters.
         rule!((closureParameterClause parameters: _* @params) => parameter* { params }),

unified/extractor/tests/corpus/swift/closures/closure-with-capture-list.output

@@ -70,9 +70,10 @@ top_level source="⟨body⟩"
           value:
             function_expr source="{ [⟨capture_declaration⟩] in ⟨body⟩ }"
               capture_declaration:
-                variable_declaration source="⟨modifier⟩ ⟨pattern⟩"
+                variable_declaration source="⟨modifier⟩ ⟨pattern⟩⟨value⟩"
                   modifier: modifier "weak" source="weak"
                   pattern: identifier "self" source="self"
+                  value: identifier "self" source="self"
               body:
                 block source="⟨stmt⟩"
                   stmt:

unified/ql/consistency-queries/CaptureSsaConsistency.ql

@@ -0,0 +1,3 @@
+private import unified
+private import codeql.unified.internal.dataflow.CaptureSsa
+import CaptureSsaOutput::ConsistencyChecks

unified/ql/lib/codeql/unified/internal/ControlFlowGraph.qll

@@ -43,6 +43,8 @@ module Ast implements AstSig<Location> {
     not skipControlFlow(result)
     or
     n.(FunctionExpr).getCaptureDeclaration(index) = result
+    or
+    n.(FunctionDeclaration).getNameNode() = result and index = 0
   }
 
   Callable getEnclosingCallable(AstNode node) { result = node.getEnclosingCallable() }

unified/ql/lib/codeql/unified/internal/FacadeAst.qll

@@ -54,6 +54,12 @@ module Unified {
         this = func.getACaptureDeclaration() and
         result = func.getParent()
       )
+      or
+      exists(FunctionDeclaration func |
+        // The name of a function declaration is a variable belonging to the enclosing callable.
+        this = func.getNameNode() and
+        result = func.getParent()
+      )
     }
 
     /**

unified/ql/lib/codeql/unified/internal/LocalNameBinding.qll

@@ -72,35 +72,10 @@ private module LocalNameBindingInput implements LocalNameBindingInputSig<Locatio
   }
 
   private AstNode getChild1(AstNode n, int index) {
-    result = n.(Block).getStmt(index) and
-    not n instanceof BlockWithGuardStmts
-    or
     result = n.(BlockWithGuardStmts).getTranslatedChild(index)
     or
     result = n.(LogicalAndRoot).getNthLeaf(index)
     or
-    exists(PatternGuardExpr guard | n = guard |
-      index = 0 and result = guard.getPattern()
-      or
-      index = 1 and result = guard.getValue()
-    )
-    or
-    exists(IfExpr expr | n = expr |
-      index = 0 and result = expr.getCondition()
-      or
-      index = 1 and result = expr.getThen()
-      or
-      index = 2 and result = expr.getElse()
-    )
-    or
-    exists(VariableDeclaration decl | n = decl |
-      index = 0 and result = decl.getPattern()
-      or
-      index = 1 and result = decl.getType()
-      or
-      index = 2 and result = decl.getValue()
-    )
-    or
     index = 0 and
     relocatedClassMember(n, result)
   }
@@ -123,7 +98,7 @@ private module LocalNameBindingInput implements LocalNameBindingInputSig<Locatio
     not n instanceof LogicalAndExpr and // also ignore intermediate nodes within a 'logical and' tree
     not n instanceof GuardIfStmt and
     not relocatedClassMember(_, result) and
-    index = 0 and
+    index = result.getParentIndex() and
     result = n.getAFieldOrChild()
   }
 
@@ -176,11 +151,6 @@ private module LocalNameBindingInput implements LocalNameBindingInputSig<Locatio
 
   private class LocalVariableDeclarationSiblingShadowingDecl extends SiblingShadowingDecl instanceof LocalVariableDeclaration
   {
-    LocalVariableDeclarationSiblingShadowingDecl() {
-      // Capture-declarations act as local variables, but are not sibling-shadowing
-      not this = any(FunctionExpr e).getACaptureDeclaration()
-    }
-
     override Expr getPattern() { result = LocalVariableDeclaration.super.getPattern() }
 
     override AstNode getRhs() { result = LocalVariableDeclaration.super.getValue() }

unified/ql/lib/codeql/unified/internal/dataflow/AllDataFlow.qll

@@ -1,6 +1,7 @@
 /** Re-exports all the files in the internal dataflow folder (except DataFlowPublic). */
 
 import CallGraph
+import CaptureSsa
 import Content
 import ConstructorPatterns
 import DataFlowCall
@@ -14,3 +15,4 @@ import ParameterPositions
 import Step
 import TaintTrackingInstantiation
 import VariableRefKind
+import LocalVariableRefNode

unified/ql/lib/codeql/unified/internal/dataflow/CaptureSsa.qll

@@ -0,0 +1,180 @@
+/**
+ * SSA for captured variables.
+ */
+
+private import unified
+private import unified as U
+private import AllDataFlow
+private import codeql.dataflow.VariableCapture
+private import codeql.unified.internal.ControlFlowGraph
+private import codeql.unified.internal.ExprPositions
+
+module CaptureSsaInput implements InputSig<Location, BasicBlock> {
+  final private class FinalCallable = U::Callable;
+
+  class Callable extends FinalCallable {
+    Callable() { this.fromSource() }
+
+    predicate isConstructor() { none() } // TODO
+  }
+
+  Callable basicBlockGetEnclosingCallable(BasicBlock bb) { result = bb.getEnclosingCallable() }
+
+  class CapturedVariable extends LocalVariable {
+    CapturedVariable() { this.isCaptured() }
+
+    Callable getCallable() { result = super.getDeclaringCallable() }
+  }
+
+  class CapturedParameter extends CapturedVariable {
+    // This class should be empty, as the capture library never gets to work with a node that corresponds to a `DataFlowInput::ParameterNode`.
+    //
+    // When a variable is mentioned in a pattern in the parameter list, the pattern itself (TValueNode) becomes the ParameterNode
+    // and the assignment to the variable is modeled as an ordinary variable-assignment at the point where the variable mention appears in the CFG.
+    // Likewise, the receiver is represented by TImplicitParameterNode and its assignment to 'self/this' is modeled as an ordinary assignment.
+    CapturedParameter() { none() }
+  }
+
+  private class TExpr = TCallableNode or TLocalVariableRefNode;
+
+  private predicate isTopLevelCallable(BuilderNode callable) {
+    callable.isCallable(any(TopLevel t))
+  }
+
+  class Expr extends TExpr {
+    Expr() { not isTopLevelCallable(this) }
+
+    string toString() { result = this.(BuilderNode).toString() }
+
+    Location getLocation() { result = this.(BuilderNode).getLocation() }
+
+    predicate hasCfgNode(BasicBlock bb, int i) {
+      this.(LocalVariableRefNode).hasCfgNode(bb, i)
+      or
+      exists(DataFlowCallable callable |
+        this = TCallableNode(callable) and
+        bb.getNode(i).injects(callable.asSourceCallable())
+      )
+    }
+  }
+
+  final private class FinalLocalVariableRefNode = LocalVariableRefNode;
+
+  class VariableWrite extends FinalLocalVariableRefNode {
+    VariableWrite() {
+      this.getRefKind().isWrite() and super.getVariable() instanceof CapturedVariable
+    }
+
+    CapturedVariable getVariable() { result = super.getVariable() }
+  }
+
+  class VariableRead extends Expr instanceof LocalVariableRefNode {
+    VariableRead() {
+      super.getVariable() instanceof CapturedVariable and
+      super.getRefKind().isRead()
+    }
+
+    CapturedVariable getVariable() { result = super.getVariable() }
+  }
+
+  /**
+   * Holds if `node` is a possible alias for `callable`.
+   */
+  private predicate callableHasLocalAlias(Callable callable, Stage2Node node) {
+    node.(BuilderNode).isCallable(callable)
+    or
+    exists(Stage2Node prev | callableHasLocalAlias(callable, prev) |
+      step(prev, any(Step s | s.value()), node)
+      or
+      localSsaStep(prev, node, _)
+    )
+    or
+    exists(CapturedVariable var |
+      callableHasLocalAliasVar(callable, var) and
+      node.(BuilderNode).isLocalVariableRead(_, var)
+    )
+  }
+
+  pragma[nomagic]
+  private predicate callableHasLocalAliasVar(Callable callable, CapturedVariable var) {
+    exists(BuilderNode ref |
+      callableHasLocalAlias(callable, ref) and
+      ref.isLocalVariableWrite(_, var)
+    )
+  }
+
+  class ClosureExpr extends Expr instanceof TCallableNode {
+    // TODO: In order to support classes with captures, this needs to cover constructor calls to such classes as well.
+    private Callable callable;
+
+    ClosureExpr() { this.(BuilderNode).isCallable(callable) }
+
+    predicate hasBody(Callable body) { callable = body }
+
+    predicate hasAliasedAccess(Expr f) {
+      // TODO: Restrict to places where the function is called or escapes into a call where it might get called.
+      callableHasLocalAlias(callable, f) and not f.(LocalVariableRefNode).getRefKind().isWrite()
+    }
+  }
+}
+
+module CaptureSsaOutput = Flow<Location, Cfg, CaptureSsaInput>;
+
+Node getNodeFromCaptureSsaNode(CaptureSsaOutput::ClosureNode n) {
+  result = TCaptureSsaNode(n)
+  or
+  result = n.(CaptureSsaOutput::ExprNode).getExpr()
+  or
+  result = n.(CaptureSsaOutput::ExprPostUpdateNode).getExpr().(BuilderNode).getPostUpdateNode()
+  or
+  result = n.(CaptureSsaOutput::VariableWriteSourceNode).getVariableWrite()
+  or
+  // NOTE: This only supports lambdas at the moment. Local classes in Swift cannot capture variables.
+  result = n.(CaptureSsaOutput::MallocNode).getClosureExpr()
+  or
+  exists(CaptureSsaOutput::ThisParameterNode thisParam, Callable callable |
+    n = thisParam and
+    callable = thisParam.getCallable() and
+    result
+        .(BuilderNode)
+        .isImplicitParameter(toDataFlowCallable(callable), any(ParameterPosition p | p.isCallee()))
+  )
+}
+
+CaptureSsaOutput::ClosureNode getCaptureSsaNodeFromNode(Node n) {
+  n = getNodeFromCaptureSsaNode(result)
+}
+
+predicate captureSsaLocalFlowStep(Node node1, Node node2) {
+  CaptureSsaOutput::localFlowStep(getCaptureSsaNodeFromNode(node1), getCaptureSsaNodeFromNode(node2))
+}
+
+predicate captureSsaStoreStep(Node node1, ContentSet contents, Node node2) {
+  CaptureSsaOutput::storeStep(getCaptureSsaNodeFromNode(node1),
+    contents.asSingleton().asCapturedVariable(), getCaptureSsaNodeFromNode(node2))
+}
+
+predicate captureSsaReadStep(Node node1, ContentSet contents, Node node2) {
+  CaptureSsaOutput::readStep(getCaptureSsaNodeFromNode(node1),
+    contents.asSingleton().asCapturedVariable(), getCaptureSsaNodeFromNode(node2))
+}
+
+predicate captureSsaClearsContent(Node node, ContentSet contents) {
+  CaptureSsaOutput::clearsContent(getCaptureSsaNodeFromNode(node),
+    contents.asSingleton().asCapturedVariable())
+}
+
+Node getCaptureSsaPostUpdate(Node pre) {
+  CaptureSsaOutput::capturePostUpdateNode(getCaptureSsaNodeFromNode(result),
+    getCaptureSsaNodeFromNode(pre))
+}
+
+predicate captureSsaAllowParameterReturnInSelf(Node param) {
+  exists(Callable callable |
+    CaptureSsaOutput::heuristicAllowInstanceParameterReturnInSelf(callable) and
+    param =
+      getNodeFromCaptureSsaNode(any(CaptureSsaOutput::ThisParameterNode n |
+          n.getCallable() = callable
+        ))
+  )
+}

unified/ql/lib/codeql/unified/internal/dataflow/Content.qll

@@ -11,7 +11,8 @@ private newtype TContent =
     name = [0 .. 20].toString()
     or
     name = getEnumCaseParameterFieldFromArgument(_, _)
-  }
+  } or
+  TCapturedVariable(CaptureSsaInput::CapturedVariable v)
 
 class Content extends TContent {
   string asNamedMember() { this = TNamedMember(result) }
@@ -22,9 +23,13 @@ class Content extends TContent {
     result = this.asNamedMember()
     or
     this.isArrayElement() and result = "ArrayElement"
+    or
+    result = this.asCapturedVariable().toString()
   }
 
-  Location getLocation() { none() }
+  LocalVariable asCapturedVariable() { this = TCapturedVariable(result) }
+
+  Location getLocation() { result = this.asCapturedVariable().getLocation() }
 }
 
 private newtype TContentSet = TSingleton(Content content)
@@ -45,4 +50,6 @@ module ContentSet {
   ContentSet namedMember(string name) { result.asSingleton().asNamedMember() = name }
 
   ContentSet arrayElement() { result.asSingleton().isArrayElement() }
+
+  ContentSet capturedVariable(LocalVariable v) { result.asSingleton().asCapturedVariable() = v }
 }