Ruby: Reimplement flow through captured variables using field flow

ruby/ql/consistency-queries/DataFlowConsistency.ql

@@ -11,8 +11,6 @@ private module Input implements InputSig<RubyDataFlow> {
   predicate postWithInFlowExclude(Node n) { n instanceof FlowSummaryNode }
 
   predicate argHasPostUpdateExclude(ArgumentNode n) {
-    n instanceof BlockArgumentNode
-    or
     n instanceof FlowSummaryNode
     or
     n instanceof SynthHashSplatArgumentNode

ruby/ql/consistency-queries/VariablesConsistency.ql

@@ -1,4 +1,5 @@
 import codeql.ruby.ast.Variable
+import codeql.ruby.dataflow.internal.DataFlowPrivate::VariableCapture::Flow::ConsistencyChecks
 
 query predicate ambiguousVariable(VariableAccess access, Variable variable) {
   access.getVariable() = variable and

ruby/ql/lib/change-notes/2023-08-23-variable-capture-flow.md

@@ -0,0 +1,4 @@
+---
+category: majorAnalysis
+---
+* Improved support for flow through captured variables that properly adheres to inter-procedural control flow.

ruby/ql/lib/codeql/ruby/ast/Call.qll

@@ -117,6 +117,17 @@ class MethodCall extends Call instanceof MethodCallImpl {
    */
   final Block getBlock() { result = super.getBlockImpl() }
 
+  /**
+   * Gets the block argument of this method call, if any.
+   * ```rb
+   * foo(&block)
+   * ```
+   */
+  final BlockArgument getBlockArgument() { result = this.getAnArgument() }
+
+  /** Holds if this method call has a block or block argument. */
+  final predicate hasBlock() { exists(this.getBlock()) or exists(this.getBlockArgument()) }
+
   /**
    * Holds if the safe navigation operator (`&.`) is used in this call.
    * ```rb

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

@@ -201,8 +201,15 @@ module ExprNodes {
 
     override LhsExpr getExpr() { result = super.getExpr() }
 
-    /** Gets a variable used in (or introduced by) this LHS. */
-    Variable getAVariable() { result = e.(VariableAccess).getVariable() }
+    /**
+     * DEPRECATED: use `getVariable` instead.
+     *
+     * Gets a variable used in (or introduced by) this LHS.
+     */
+    deprecated Variable getAVariable() { result = e.(VariableAccess).getVariable() }
+
+    /** Gets the variable used in (or introduced by) this LHS. */
+    Variable getVariable() { result = e.(VariableAccess).getVariable() }
   }
 
   private class AssignExprChildMapping extends ExprChildMapping, AssignExpr {

ruby/ql/lib/codeql/ruby/dataflow/FlowSummary.qll

@@ -1,6 +1,8 @@
 /** Provides classes and predicates for defining flow summaries. */
 
 import codeql.ruby.AST
+private import codeql.ruby.CFG
+private import codeql.ruby.typetracking.TypeTracker
 import codeql.ruby.DataFlow
 private import internal.FlowSummaryImpl as Impl
 private import internal.DataFlowDispatch
@@ -158,3 +160,57 @@ abstract class SimpleSummarizedCallable extends SummarizedCallable {
 }
 
 class RequiredSummaryComponentStack = Impl::Public::RequiredSummaryComponentStack;
+
+/**
+ * Provides a set of special flow summaries to ensure that callbacks passed into
+ * library methods will be passed as `lambda-self` arguments into themselves. That is,
+ * we are assuming that callbacks passed into library methods will be called, which is
+ * needed for flow through captured variables.
+ */
+private module LibraryCallbackSummaries {
+  private predicate libraryCall(CfgNodes::ExprNodes::CallCfgNode call) {
+    not exists(getTarget(call))
+  }
+
+  private DataFlow::LocalSourceNode trackLambdaCreation(TypeTracker t) {
+    t.start() and
+    lambdaCreation(result, TLambdaCallKind(), _)
+    or
+    exists(TypeTracker t2 | result = trackLambdaCreation(t2).track(t2, t)) and
+    not result instanceof DataFlow::SelfParameterNode
+  }
+
+  private predicate libraryCallHasLambdaArg(CfgNodes::ExprNodes::CallCfgNode call, int i) {
+    exists(CfgNodes::ExprCfgNode arg |
+      arg = call.getArgument(i) and
+      arg = trackLambdaCreation(TypeTracker::end()).getALocalUse().asExpr() and
+      libraryCall(call) and
+      not arg instanceof CfgNodes::ExprNodes::BlockArgumentCfgNode
+    )
+  }
+
+  private class LibraryLambdaMethod extends SummarizedCallable {
+    LibraryLambdaMethod() { this = "<library method accepting a callback>" }
+
+    final override MethodCall getACall() {
+      libraryCall(result.getAControlFlowNode()) and
+      result.hasBlock()
+      or
+      libraryCallHasLambdaArg(result.getAControlFlowNode(), _)
+    }
+
+    override predicate propagatesFlowExt(string input, string output, boolean preservesValue) {
+      (
+        input = "Argument[block]" and
+        output = "Argument[block].Parameter[lambda-self]"
+        or
+        exists(int i |
+          i in [0 .. 10] and
+          input = "Argument[" + i + "]" and
+          output = "Argument[" + i + "].Parameter[lambda-self]"
+        )
+      ) and
+      preservesValue = true
+    }
+  }
+}

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

@@ -21,8 +21,7 @@ private class SelfLocalSourceNode extends DataFlow::LocalSourceNode {
   SelfLocalSourceNode() {
     self = this.(SelfParameterNodeImpl).getSelfVariable()
     or
-    self = this.(SsaSelfDefinitionNode).getVariable() and
-    not LocalFlow::localFlowSsaParamInput(_, this)
+    self = this.(SsaSelfDefinitionNode).getVariable()
   }
 
   /** Gets the `self` variable. */
@@ -424,6 +423,7 @@ private module Cached {
   cached
   newtype TArgumentPosition =
     TSelfArgumentPosition() or
+    TLambdaSelfArgumentPosition() or
     TBlockArgumentPosition() or
     TPositionalArgumentPosition(int pos) {
       exists(Call c | exists(c.getArgument(pos)))
@@ -446,6 +446,7 @@ private module Cached {
   cached
   newtype TParameterPosition =
     TSelfParameterPosition() or
+    TLambdaSelfParameterPosition() or
     TBlockParameterPosition() or
     TPositionalParameterPosition(int pos) {
       pos = any(Parameter p).getPosition()
@@ -941,20 +942,24 @@ private module TrackSingletonMethodOnInstanceInput implements CallGraphConstruct
   private predicate paramReturnFlow(
     DataFlow::Node nodeFrom, DataFlow::PostUpdateNode nodeTo, StepSummary summary
   ) {
-    exists(RelevantCall call, DataFlow::Node arg, DataFlow::ParameterNode p, Expr nodeFromPreExpr |
+    exists(
+      RelevantCall call, DataFlow::Node arg, DataFlow::ParameterNode p,
+      CfgNodes::ExprCfgNode nodeFromPreExpr
+    |
       TypeTrackerSpecific::callStep(call, arg, p) and
       nodeTo.getPreUpdateNode() = arg and
       summary.toString() = "return" and
       (
-        nodeFromPreExpr = nodeFrom.(DataFlow::PostUpdateNode).getPreUpdateNode().asExpr().getExpr()
+        nodeFromPreExpr = nodeFrom.(DataFlow::PostUpdateNode).getPreUpdateNode().asExpr()
         or
-        nodeFromPreExpr = nodeFrom.asExpr().getExpr() and
-        singletonMethodOnInstance(_, _, nodeFromPreExpr)
+        nodeFromPreExpr = nodeFrom.asExpr() and
+        singletonMethodOnInstance(_, _, nodeFromPreExpr.getExpr())
       )
     |
-      nodeFromPreExpr = p.getParameter().(NamedParameter).getVariable().getAnAccess()
+      nodeFromPreExpr =
+        LocalFlow::getParameterDefNode(p.getParameter()).getDefinitionExt().getARead()
       or
-      nodeFromPreExpr = p.(SelfParameterNodeImpl).getSelfVariable().getAnAccess()
+      nodeFromPreExpr = p.(SelfParameterNodeImpl).getSelfDefinition().getARead()
     )
   }
 
@@ -1276,6 +1281,9 @@ class ParameterPosition extends TParameterPosition {
   /** Holds if this position represents a `self` 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 block parameter. */
   predicate isBlock() { this = TBlockParameterPosition() }
 
@@ -1313,6 +1321,8 @@ class ParameterPosition extends TParameterPosition {
   string toString() {
     this.isSelf() and result = "self"
     or
+    this.isLambdaSelf() and result = "lambda self"
+    or
     this.isBlock() and result = "block"
     or
     exists(int pos | this.isPositional(pos) and result = "position " + pos)
@@ -1342,6 +1352,9 @@ class ArgumentPosition extends TArgumentPosition {
   /** Holds if this position represents a `self` 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 block argument. */
   predicate isBlock() { this = TBlockArgumentPosition() }
 
@@ -1374,6 +1387,8 @@ class ArgumentPosition extends TArgumentPosition {
   string toString() {
     this.isSelf() and result = "self"
     or
+    this.isLambdaSelf() and result = "lambda self"
+    or
     this.isBlock() and result = "block"
     or
     exists(int pos | this.isPositional(pos) and result = "position " + pos)
@@ -1393,16 +1408,24 @@ class ArgumentPosition extends TArgumentPosition {
 }
 
 pragma[nomagic]
-private predicate parameterPositionIsNotSelf(ParameterPosition ppos) { not ppos.isSelf() }
+private predicate parameterPositionIsNotSelf(ParameterPosition ppos) {
+  not ppos.isSelf() and
+  not ppos.isLambdaSelf()
+}
 
 pragma[nomagic]
-private predicate argumentPositionIsNotSelf(ArgumentPosition apos) { not apos.isSelf() }
+private predicate argumentPositionIsNotSelf(ArgumentPosition apos) {
+  not apos.isSelf() and
+  not apos.isLambdaSelf()
+}
 
 /** Holds if arguments at position `apos` match parameters at position `ppos`. */
 pragma[nomagic]
 predicate parameterMatch(ParameterPosition ppos, ArgumentPosition apos) {
   ppos.isSelf() and apos.isSelf()
   or
+  ppos.isLambdaSelf() and apos.isLambdaSelf()
+  or
   ppos.isBlock() and apos.isBlock()
   or
   exists(int pos | ppos.isPositional(pos) and apos.isPositional(pos))
@@ -1441,8 +1464,6 @@ predicate parameterMatch(ParameterPosition ppos, ArgumentPosition apos) {
  * This is a temporary hook to support technical debt in the Go language; do not use.
  */
 pragma[inline]
-predicate golangSpecificParamArgFilter(
-  DataFlowCall call, DataFlow::ParameterNode p, ArgumentNode arg
-) {
+predicate golangSpecificParamArgFilter(DataFlowCall call, ParameterNodeImpl p, ArgumentNode arg) {
   any()
 }

ruby/ql/lib/codeql/ruby/dataflow/internal/DataFlowImplSpecific.qll

@@ -17,6 +17,9 @@ module RubyDataFlow implements InputSig {
   import Private
   import Public
 
+  // includes `LambdaSelfReferenceNode`, which is not part of the public API
+  class ParameterNode = Private::ParameterNodeImpl;
+
   predicate isParameterNode(ParameterNode p, DataFlowCallable c, ParameterPosition pos) {
     Private::isParameterNode(p, c, pos)
   }