Ruby: Handle captured variables in BarrierGuard::getAGuardedNode()

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

@@ -767,10 +767,10 @@ private module OutNodes {
 import OutNodes
 
 predicate jumpStep(Node pred, Node succ) {
-  SsaImpl::captureFlowIn(pred.(SsaDefinitionNode).getDefinition(),
+  SsaImpl::captureFlowIn(_, pred.(SsaDefinitionNode).getDefinition(),
     succ.(SsaDefinitionNode).getDefinition())
   or
-  SsaImpl::captureFlowOut(pred.(SsaDefinitionNode).getDefinition(),
+  SsaImpl::captureFlowOut(_, pred.(SsaDefinitionNode).getDefinition(),
     succ.(SsaDefinitionNode).getDefinition())
   or
   succ.asExpr().getExpr().(ConstantReadAccess).getValue() = pred.asExpr().getExpr()

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

@@ -5,6 +5,7 @@ private import codeql.ruby.CFG
 private import codeql.ruby.typetracking.TypeTracker
 private import codeql.ruby.dataflow.SSA
 private import FlowSummaryImpl as FlowSummaryImpl
+private import SsaImpl as SsaImpl
 
 /**
  * An element, viewed as a node in a data flow graph. Either an expression
@@ -256,6 +257,26 @@ abstract class BarrierGuard extends CfgNodes::ExprCfgNode {
    */
   abstract predicate checks(CfgNode expr, boolean branch);
 
+  /**
+   * Gets an implicit entry definition for a captured variable that
+   * may be guarded, because a call to the capturing callable is guarded.
+   *
+   * This is restricted to calls where the variable is captured inside a
+   * block.
+   */
+  private Ssa::Definition getAMaybeGuardedCapturedDef() {
+    exists(
+      boolean branch, CfgNodes::ExprCfgNode testedNode, Ssa::Definition def,
+      CfgNodes::ExprNodes::CallCfgNode call
+    |
+      def.getARead() = testedNode and
+      this.checks(testedNode, branch) and
+      SsaImpl::captureFlowIn(call, def, result) and
+      this.controlsBlock(call.getBasicBlock(), branch) and
+      result.getBasicBlock().getScope() = call.getExpr().(MethodCall).getBlock()
+    )
+  }
+
   final Node getAGuardedNode() {
     exists(boolean branch, CfgNodes::ExprCfgNode testedNode, Ssa::Definition def |
       def.getARead() = testedNode and
@@ -263,5 +284,7 @@ abstract class BarrierGuard extends CfgNodes::ExprCfgNode {
       this.checks(testedNode, branch) and
       this.controlsBlock(result.asExpr().getBasicBlock(), branch)
     )
+    or
+    result.asExpr() = this.getAMaybeGuardedCapturedDef().getARead()
   }
 }

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

@@ -72,16 +72,18 @@ private predicate hasVariableWriteWithCapturedRead(BasicBlock bb, LocalVariable
  * Holds if the call `call` at index `i` in basic block `bb` may reach
  * a callable that reads captured variable `v`.
  */
-private predicate capturedCallRead(Call call, BasicBlock bb, int i, LocalVariable v) {
+private predicate capturedCallRead(
+  CfgNodes::ExprNodes::CallCfgNode call, BasicBlock bb, int i, LocalVariable v
+) {
   exists(CfgScope scope |
     hasVariableWriteWithCapturedRead(bb.getAPredecessor*(), v, scope) and
-    call = bb.getNode(i).getNode()
+    call = bb.getNode(i)
   |
     // If the read happens inside a block, we restrict to the call that
     // contains the block
     not scope instanceof Block
     or
-    scope = call.(MethodCall).getBlock()
+    scope = call.getExpr().(MethodCall).getBlock()
   )
 }
 
@@ -148,16 +150,18 @@ private module Cached {
    * that writes captured variable `v`.
    */
   cached
-  predicate capturedCallWrite(Call call, BasicBlock bb, int i, LocalVariable v) {
+  predicate capturedCallWrite(
+    CfgNodes::ExprNodes::CallCfgNode call, BasicBlock bb, int i, LocalVariable v
+  ) {
     exists(CfgScope scope |
       hasVariableReadWithCapturedWrite(bb.getASuccessor*(), v, scope) and
-      call = bb.getNode(i).getNode()
+      call = bb.getNode(i)
     |
       // If the write happens inside a block, we restrict to the call that
       // contains the block
       not scope instanceof Block
       or
-      scope = call.(MethodCall).getBlock()
+      scope = call.getExpr().(MethodCall).getBlock()
     )
   }
 
@@ -189,7 +193,7 @@ private module Cached {
 
   pragma[noinline]
   private predicate defReachesCallReadInOuterScope(
-    Definition def, Call call, LocalVariable v, CfgScope scope
+    Definition def, CfgNodes::ExprNodes::CallCfgNode call, LocalVariable v, CfgScope scope
   ) {
     exists(BasicBlock bb, int i |
       ssaDefReachesRead(v, def, bb, i) and
@@ -217,8 +221,8 @@ private module Cached {
    * ```
    */
   cached
-  predicate captureFlowIn(Definition def, Definition entry) {
-    exists(Call call, LocalVariable v, CfgScope scope |
+  predicate captureFlowIn(CfgNodes::ExprNodes::CallCfgNode call, Definition def, Definition entry) {
+    exists(LocalVariable v, CfgScope scope |
       defReachesCallReadInOuterScope(def, call, v, scope) and
       hasCapturedEntryWrite(entry, v, scope)
     |
@@ -226,7 +230,7 @@ private module Cached {
       // contains the block
       not scope instanceof Block
       or
-      scope = call.(MethodCall).getBlock()
+      scope = call.getExpr().(MethodCall).getBlock()
     )
   }
 
@@ -242,7 +246,9 @@ private module Cached {
   }
 
   pragma[noinline]
-  private predicate hasCapturedExitRead(Definition exit, Call call, LocalVariable v, CfgScope scope) {
+  private predicate hasCapturedExitRead(
+    Definition exit, CfgNodes::ExprNodes::CallCfgNode call, LocalVariable v, CfgScope scope
+  ) {
     exists(BasicBlock bb, int i |
       capturedCallWrite(call, bb, i, v) and
       exit.definesAt(v, bb, i) and
@@ -261,8 +267,8 @@ private module Cached {
    * ```
    */
   cached
-  predicate captureFlowOut(Definition def, Definition exit) {
-    exists(Call call, LocalVariable v, CfgScope scope |
+  predicate captureFlowOut(CfgNodes::ExprNodes::CallCfgNode call, Definition def, Definition exit) {
+    exists(LocalVariable v, CfgScope scope |
       defReachesExitReadInInnerScope(def, v, scope) and
       hasCapturedExitRead(exit, call, v, _)
     |
@@ -270,7 +276,7 @@ private module Cached {
       // contains the block
       not scope instanceof Block
       or
-      scope = call.(MethodCall).getBlock()
+      scope = call.getExpr().(MethodCall).getBlock()
     )
   }
 

ruby/ql/test/library-tests/dataflow/barrier-guards/barrier-guards.expected

@@ -4,3 +4,4 @@
 | barrier-guards.rb:21:8:21:19 | ... == ... | barrier-guards.rb:24:5:24:7 | foo | barrier-guards.rb:21:8:21:10 | foo | true |
 | barrier-guards.rb:27:8:27:19 | ... != ... | barrier-guards.rb:28:5:28:7 | foo | barrier-guards.rb:27:8:27:10 | foo | false |
 | barrier-guards.rb:37:4:37:20 | call to include? | barrier-guards.rb:38:5:38:7 | foo | barrier-guards.rb:37:17:37:19 | foo | true |
+| barrier-guards.rb:43:4:43:15 | ... == ... | barrier-guards.rb:45:9:45:11 | foo | barrier-guards.rb:43:4:43:6 | foo | true |