Resolve yield calls to blocks

ql/src/codeql_ruby/controlflow/CfgNodes.qll

@@ -231,9 +231,23 @@ module ExprNodes {
     final ExprCfgNode getRightOperand() { e.hasCfgChild(e.getRightOperand(), this, result) }
   }
 
+  private class BlockArgumentChildMapping extends ExprChildMapping, BlockArgument {
+    override predicate relevantChild(Expr e) { e = this.getValue() }
+  }
+
+  /** A control-flow node that wraps a `BlockArgument` AST expression. */
+  class BlockArgumentCfgNode extends ExprCfgNode {
+    override BlockArgumentChildMapping e;
+
+    final override BlockArgument getExpr() { result = ExprCfgNode.super.getExpr() }
+
+    /** Gets the value of this block argument. */
+    final ExprCfgNode getValue() { e.hasCfgChild(e.getValue(), this, result) }
+  }
+
   private class CallExprChildMapping extends ExprChildMapping, Call {
     override predicate relevantChild(Expr e) {
-      e = [this.getAnArgument(), this.(MethodCall).getReceiver()]
+      e = [this.getAnArgument(), this.(MethodCall).getReceiver(), this.(MethodCall).getBlock()]
     }
   }
 
@@ -248,6 +262,9 @@ module ExprNodes {
 
     /** Gets the receiver of this call. */
     final ExprCfgNode getReceiver() { e.hasCfgChild(e.(MethodCall).getReceiver(), this, result) }
+
+    /** Gets the block of this call. */
+    final ExprCfgNode getBlock() { e.hasCfgChild(e.(MethodCall).getBlock(), this, result) }
   }
 
   private class CaseExprChildMapping extends ExprChildMapping, CaseExpr {

ql/src/codeql_ruby/dataflow/internal/DataFlowDispatch.qll

@@ -53,6 +53,14 @@ class DataFlowCall extends CfgNodes::ExprNodes::CallCfgNode {
     )
   }
 
+  private Block yieldCall() {
+    this.getExpr() instanceof YieldCall and
+    exists(BlockParameterNode node |
+      node = trackBlock(result) and
+      node.getMethod() = this.getExpr().getEnclosingMethod()
+    )
+  }
+
   pragma[nomagic]
   private predicate superCall(Module superClass, string method) {
     this.getExpr() instanceof SuperCall and
@@ -89,6 +97,8 @@ class DataFlowCall extends CfgNodes::ExprNodes::CallCfgNode {
       this.superCall(superClass, method) and
       result = lookupMethod(superClass, method)
     )
+    or
+    result = this.yieldCall()
   }
 }
 
@@ -163,6 +173,16 @@ private DataFlow::LocalSourceNode trackInstance(Module tp) {
   result = trackInstance(tp, TypeTracker::end())
 }
 
+private DataFlow::LocalSourceNode trackBlock(Block block, TypeTracker t) {
+  t.start() and result.asExpr().getExpr() = block
+  or
+  exists(TypeTracker t2 | result = trackBlock(block, t2).track(t2, t))
+}
+
+private DataFlow::LocalSourceNode trackBlock(Block block) {
+  result = trackBlock(block, TypeTracker::end())
+}
+
 private predicate singletonMethod(MethodBase method, Expr object) {
   object = method.(SingletonMethod).getObject()
   or

ql/src/codeql_ruby/dataflow/internal/DataFlowPrivate.qll

@@ -115,9 +115,12 @@ private module Cached {
     TExprNode(CfgNodes::ExprCfgNode n) or
     TReturningNode(CfgNodes::ReturningCfgNode n) or
     TSsaDefinitionNode(Ssa::Definition def) or
-    TParameterNode(Parameter p) or
+    TNormalParameterNode(Parameter p) { not p instanceof BlockParameter } or
+    TBlockParameterNode(MethodBase m) or
     TExprPostUpdateNode(CfgNodes::ExprCfgNode n) { n.getNode() instanceof Argument }
 
+  class TParameterNode = TNormalParameterNode or TBlockParameterNode;
+
   /**
    * This is the local flow predicate that is used as a building block in global
    * data flow. It excludes SSA flow through instance fields, as flow through fields
@@ -136,6 +139,8 @@ private module Cached {
     or
     nodeFrom.asExpr() = nodeTo.asExpr().(CfgNodes::ExprNodes::AssignExprCfgNode).getRhs()
     or
+    nodeFrom.asExpr() = nodeTo.asExpr().(CfgNodes::ExprNodes::BlockArgumentCfgNode).getValue()
+    or
     nodeFrom.asExpr() = nodeTo.asExpr().(CfgNodes::ExprNodes::StmtSequenceCfgNode).getLastStmt()
     or
     nodeFrom.asExpr() = nodeTo.asExpr().(CfgNodes::ExprNodes::ConditionalExprCfgNode).getBranch(_)
@@ -221,10 +226,10 @@ private module ParameterNodes {
    * The value of an explicit parameter at function entry, viewed as a node in a data
    * flow graph.
    */
-  class ExplicitParameterNode extends ParameterNodeImpl, TParameterNode {
+  class ExplicitParameterNode extends ParameterNodeImpl, TNormalParameterNode {
     private Parameter parameter;
 
-    ExplicitParameterNode() { this = TParameterNode(parameter) }
+    ExplicitParameterNode() { this = TNormalParameterNode(parameter) }
 
     override Parameter getParameter() { result = parameter }
 
@@ -236,6 +241,38 @@ private module ParameterNodes {
 
     override string toStringImpl() { result = parameter.toString() }
   }
+
+  /**
+   * The value of a block parameter at function entry, viewed as a node in a data
+   * flow graph.
+   */
+  class BlockParameterNode extends ParameterNodeImpl, TBlockParameterNode {
+    private MethodBase method;
+
+    BlockParameterNode() { this = TBlockParameterNode(method) }
+
+    final MethodBase getMethod() { result = method }
+
+    override Parameter getParameter() {
+      result = method.getAParameter() and result instanceof BlockParameter
+    }
+
+    override predicate isParameterOf(Callable c, int i) { c = method and i = -2 }
+
+    override CfgScope getCfgScope() { result = method }
+
+    override Location getLocationImpl() {
+      result = getParameter().getLocation()
+      or
+      not exists(getParameter()) and result = method.getLocation()
+    }
+
+    override string toStringImpl() {
+      result = getParameter().toString()
+      or
+      not exists(getParameter()) and result = "&block"
+    }
+  }
 }
 
 import ParameterNodes
@@ -253,7 +290,10 @@ abstract class ArgumentNode extends Node {
 private module ArgumentNodes {
   /** A data-flow node that represents an explicit call argument. */
   class ExplicitArgumentNode extends ArgumentNode {
-    ExplicitArgumentNode() { this.asExpr().getExpr() instanceof Argument }
+    ExplicitArgumentNode() {
+      this.asExpr().getExpr() instanceof Argument and
+      not this.asExpr().getExpr() instanceof BlockArgument
+    }
 
     override predicate argumentOf(DataFlowCall call, int pos) {
       this.asExpr() = call.getReceiver() and
@@ -262,6 +302,27 @@ private module ArgumentNodes {
       this.asExpr() = call.getArgument(pos)
     }
   }
+
+  /** A data-flow node that represents a block argument. */
+  class BlockArgumentNode extends ArgumentNode {
+    BlockArgumentNode() {
+      this.asExpr().getExpr() instanceof BlockArgument or
+      exists(CfgNodes::ExprNodes::CallCfgNode c | c.getBlock() = this.asExpr())
+    }
+
+    override predicate argumentOf(DataFlowCall call, int pos) {
+      pos = -2 and
+      (
+        this.asExpr() = call.getBlock()
+        or
+        exists(CfgNodes::ExprCfgNode arg, int n |
+          arg = call.getArgument(n) and
+          this.asExpr() = arg and
+          arg.getExpr() instanceof BlockArgument
+        )
+      )
+    }
+  }
 }
 
 import ArgumentNodes

ql/src/codeql_ruby/dataflow/internal/DataFlowPublic.qll

@@ -60,18 +60,14 @@ class ExprNode extends Node, TExprNode {
  * flow graph.
  */
 class ParameterNode extends Node, TParameterNode {
-  private Parameter p;
-
-  ParameterNode() { this = TParameterNode(p) }
-
   /** Gets the parameter corresponding to this node, if any. */
-  Parameter getParameter() { result = p }
+  Parameter getParameter() { none() }
 
   /**
    * Holds if this node is the parameter of callable `c` at the specified
    * (zero-based) position.
    */
-  predicate isParameterOf(Callable c, int i) { p = c.getParameter(i) }
+  predicate isParameterOf(Callable c, int i) { none() }
 }
 
 /**

ql/src/codeql_ruby/typetracking/TypeTrackerSpecific.qll

@@ -31,9 +31,7 @@ predicate jumpStep(Node nodeFrom, Node nodeTo) {
 string getPossibleContentName() { result = getSetterCallAttributeName(_) }
 
 /** Holds if `nodeFrom` steps to `nodeTo` by being passed as a parameter in a call. */
-predicate callStep(
-  DataFlowPrivate::ArgumentNode nodeFrom, DataFlowPrivate::ExplicitParameterNode nodeTo
-) {
+predicate callStep(DataFlowPrivate::ArgumentNode nodeFrom, DataFlowPublic::ParameterNode nodeTo) {
   exists(DataFlowDispatch::DataFlowCall call, DataFlowDispatch::DataFlowCallable callable, int i |
     call.getTarget() = callable and
     nodeFrom.argumentOf(call, i) and

ql/test/library-tests/modules/callgraph.expected

@@ -21,27 +21,37 @@ getTarget
 | calls.rb:60:1:60:12 | call to instance_m | calls.rb:16:5:16:23 | instance_m |
 | calls.rb:63:5:63:16 | call to bit_length | calls.rb:78:5:78:23 | bit_length |
 | calls.rb:64:5:64:16 | call to bit_length | calls.rb:78:5:78:23 | bit_length |
+| calls.rb:68:5:68:11 | yield ... | calls.rb:74:16:74:29 | { ... } |
+| calls.rb:68:5:68:11 | yield ... | calls.rb:140:10:140:28 | { ... } |
 | calls.rb:72:11:72:18 | call to new | calls.rb:98:5:98:16 | new |
 | calls.rb:73:5:73:10 | ...[...] | calls.rb:102:5:102:15 | [] |
 | calls.rb:74:5:74:29 | call to call_block | calls.rb:67:1:69:3 | call_block |
 | calls.rb:74:22:74:27 | ...[...] | calls.rb:102:5:102:15 | [] |
 | calls.rb:97:5:97:18 | call to include | calls.rb:92:5:92:20 | include |
 | calls.rb:111:15:111:25 | call to length | calls.rb:107:3:107:17 | length |
+| calls.rb:112:9:112:24 | yield ... | calls.rb:128:23:128:62 | { ... } |
+| calls.rb:112:9:112:24 | yield ... | calls.rb:130:17:130:35 | { ... } |
+| calls.rb:112:9:112:24 | yield ... | calls.rb:132:17:132:40 | { ... } |
+| calls.rb:112:9:112:24 | yield ... | calls.rb:134:18:134:37 | { ... } |
 | calls.rb:112:18:112:24 | ...[...] | calls.rb:106:3:106:13 | [] |
+| calls.rb:119:5:119:20 | yield ... | calls.rb:122:7:122:30 | { ... } |
 | calls.rb:122:1:122:30 | call to funny | calls.rb:118:1:120:3 | funny |
 | calls.rb:122:13:122:29 | call to puts | calls.rb:87:5:87:17 | puts |
+| calls.rb:122:18:122:29 | call to capitalize | calls.rb:83:5:83:23 | capitalize |
 | calls.rb:124:1:124:14 | call to capitalize | calls.rb:83:5:83:23 | capitalize |
 | calls.rb:125:1:125:12 | call to bit_length | calls.rb:78:5:78:23 | bit_length |
 | calls.rb:126:1:126:5 | call to abs | calls.rb:79:5:79:16 | abs |
 | calls.rb:128:1:128:62 | call to foreach | calls.rb:109:3:115:5 | foreach |
 | calls.rb:128:32:128:61 | call to puts | calls.rb:87:5:87:17 | puts |
 | calls.rb:130:1:130:35 | call to foreach | calls.rb:109:3:115:5 | foreach |
+| calls.rb:130:23:130:34 | call to bit_length | calls.rb:78:5:78:23 | bit_length |
 | calls.rb:132:1:132:40 | call to foreach | calls.rb:109:3:115:5 | foreach |
 | calls.rb:132:23:132:39 | call to puts | calls.rb:87:5:87:17 | puts |
 | calls.rb:134:1:134:37 | call to foreach | calls.rb:109:3:115:5 | foreach |
 | calls.rb:134:27:134:36 | call to puts | calls.rb:87:5:87:17 | puts |
 | calls.rb:137:5:137:17 | call to call_block | calls.rb:67:1:69:3 | call_block |
 | calls.rb:140:1:140:28 | call to indirect | calls.rb:136:1:138:3 | indirect |
+| calls.rb:140:16:140:27 | call to bit_length | calls.rb:78:5:78:23 | bit_length |
 | calls.rb:159:1:159:5 | call to new | calls.rb:98:5:98:16 | new |
 | calls.rb:159:1:159:14 | call to s_method | calls.rb:144:5:146:7 | s_method |
 | calls.rb:160:1:160:5 | call to new | calls.rb:98:5:98:16 | new |
@@ -74,13 +84,7 @@ unresolvedCall
 | calls.rb:42:9:42:24 | call to singleton_m |
 | calls.rb:48:1:48:13 | call to singleton_m |
 | calls.rb:59:1:59:13 | call to singleton_m |
-| calls.rb:68:5:68:11 | yield ... |
-| calls.rb:112:9:112:24 | yield ... |
-| calls.rb:119:5:119:20 | yield ... |
-| calls.rb:122:18:122:29 | call to capitalize |
 | calls.rb:128:48:128:59 | call to capitalize |
-| calls.rb:130:23:130:34 | call to bit_length |
 | calls.rb:132:28:132:39 | call to capitalize |
 | calls.rb:134:32:134:36 | call to abs |
-| calls.rb:140:16:140:27 | call to bit_length |
 | calls.rb:145:9:145:17 | call to to_s |