Ruby: Reimplement flow through captured variables using field flow

Author: hvitved

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

@@ -727,16 +727,30 @@ module ExprNodes {
 
     override VariableAccess e;
 
-    final override VariableAccess getExpr() { result = ExprCfgNode.super.getExpr() }
+    override VariableAccess getExpr() { result = ExprCfgNode.super.getExpr() }
+
+    /** Gets the variable that is being accessed. */
+    Variable getVariable() { result = this.getExpr().getVariable() }
   }
 
   /** A control-flow node that wraps a `VariableReadAccess` AST expression. */
-  class VariableReadAccessCfgNode extends ExprCfgNode {
+  class VariableReadAccessCfgNode extends VariableAccessCfgNode {
     override string getAPrimaryQlClass() { result = "VariableReadAccessCfgNode" }
 
     override VariableReadAccess e;
 
-    final override VariableReadAccess getExpr() { result = ExprCfgNode.super.getExpr() }
+    override VariableReadAccess getExpr() { result = VariableAccessCfgNode.super.getExpr() }
+  }
+
+  /** A control-flow node that wraps a `LocalVariableReadAccess` AST expression. */
+  class LocalVariableReadAccessCfgNode extends VariableReadAccessCfgNode {
+    override string getAPrimaryQlClass() { result = "LocalVariableReadAccessCfgNode" }
+
+    override LocalVariableReadAccess e;
+
+    final override LocalVariableReadAccess getExpr() { result = super.getExpr() }
+
+    final override LocalVariable getVariable() { result = super.getVariable() }
   }
 
   private class InstanceVariableAccessMapping extends ExprChildMapping, InstanceVariableAccess {
@@ -762,21 +776,32 @@ module ExprNodes {
   }
 
   /** A control-flow node that wraps a `SelfVariableAccess` AST expression. */
-  class SelfVariableAccessCfgNode extends ExprCfgNode {
+  class SelfVariableAccessCfgNode extends VariableAccessCfgNode {
     final override string getAPrimaryQlClass() { result = "SelfVariableAccessCfgNode" }
 
     override SelfVariableAccessMapping e;
 
-    override SelfVariableAccess getExpr() { result = ExprCfgNode.super.getExpr() }
+    override SelfVariableAccess getExpr() { result = VariableAccessCfgNode.super.getExpr() }
   }
 
   /** A control-flow node that wraps a `VariableWriteAccess` AST expression. */
-  class VariableWriteAccessCfgNode extends ExprCfgNode {
+  class VariableWriteAccessCfgNode extends VariableAccessCfgNode {
     override string getAPrimaryQlClass() { result = "VariableWriteAccessCfgNode" }
 
     override VariableWriteAccess e;
 
-    final override VariableWriteAccess getExpr() { result = ExprCfgNode.super.getExpr() }
+    override VariableWriteAccess getExpr() { result = VariableAccessCfgNode.super.getExpr() }
+  }
+
+  /** A control-flow node that wraps a `LocalVariableWriteAccess` AST expression. */
+  class LocalVariableWriteAccessCfgNode extends VariableWriteAccessCfgNode {
+    override string getAPrimaryQlClass() { result = "LocalVariableWriteAccessCfgNode" }
+
+    override LocalVariableWriteAccess e;
+
+    final override LocalVariableWriteAccess getExpr() { result = super.getExpr() }
+
+    final override LocalVariable getVariable() { result = super.getVariable() }
   }
 
   /** A control-flow node that wraps a `ConstantReadAccess` AST expression. */

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

@@ -178,6 +178,9 @@ module Ssa {
 
     /** Gets the location of this SSA definition. */
     Location getLocation() { result = this.getControlFlowNode().getLocation() }
+
+    /** Gets the scope of this SSA definition. */
+    CfgScope getScope() { result = this.getBasicBlock().getScope() }
   }
 
   /**
@@ -289,7 +292,7 @@ module Ssa {
       )
     }
 
-    final override string toString() { result = "<captured> " + this.getSourceVariable() }
+    final override string toString() { result = "<captured entry> " + this.getSourceVariable() }
 
     override Location getLocation() { result = this.getBasicBlock().getLocation() }
   }
@@ -324,7 +327,7 @@ module Ssa {
      */
     final Definition getPriorDefinition() { result = SsaImpl::uncertainWriteDefinitionInput(this) }
 
-    override string toString() { result = this.getControlFlowNode().toString() }
+    override string toString() { result = "<captured exit> " + this.getSourceVariable() }
   }
 
   /**

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

@@ -0,0 +1,154 @@
+private import codeql.ssa.Ssa as SsaImplCommon
+private import codeql.ruby.AST
+private import codeql.ruby.CFG as Cfg
+private import codeql.ruby.controlflow.internal.ControlFlowGraphImplShared as ControlFlowGraphImplShared
+private import codeql.ruby.dataflow.SSA
+private import codeql.ruby.ast.Variable
+private import Cfg::CfgNodes::ExprNodes
+
+class CapturedVariableAccess extends LocalVariableAccess {
+  CapturedVariableAccess() { this.isCapturedAccess() }
+}
+
+class CapturedVariable extends LocalVariable {
+  CapturedVariable() { this.isCaptured() }
+}
+
+class CapturingCallable extends Callable {
+  CapturingCallable() { any(CapturedVariableAccess access).getCfgScope() = this }
+
+  CapturedVariable getACapturedVariable() { result.getAnAccess().getCfgScope() = this }
+}
+
+private predicate entryWrite(Cfg::EntryBasicBlock bb, int i, CapturingCallable v) {
+  i = -1 and
+  (
+    bb.getScope() = v
+    or
+    bb.getScope() = v.getACapturedVariable().getDeclaringScope()
+  )
+}
+
+private predicate captureRead1(
+  Cfg::BasicBlock bb, int i, CapturedVariable var, VariableAccessCfgNode access, CapturingCallable v
+) {
+  bb.getNode(i) = access and
+  access.getNode() = var.getAnAccess() and
+  var = v.getACapturedVariable() and
+  (bb.getScope() = v or bb.getScope() = var.getDefiningAccess().getCfgScope())
+}
+
+private predicate captureRead2(Cfg::BasicBlock bb, int i, CapturingCallable v) {
+  bb.getNode(i).getNode() = v
+}
+
+newtype TCaptureAccess =
+  TEntryDef(Cfg::EntryBasicBlock bb, int i, CapturingCallable v) { entryWrite(bb, i, v) } or
+  TVariableAccess(VariableAccessCfgNode access, CapturingCallable v) {
+    captureRead1(_, _, _, access, v)
+  } or
+  TLambdaAccess(Cfg::BasicBlock bb, int i, CapturingCallable v) { captureRead2(bb, i, v) }
+
+class CaptureAccess extends TCaptureAccess {
+  string toString() { result = "<captured variable>" }
+
+  Location getLocation() {
+    exists(Cfg::EntryBasicBlock bb, int i, CapturingCallable v |
+      this = TEntryDef(bb, i, v) and
+      result = bb.getLocation()
+    )
+    or
+    exists(VariableAccessCfgNode access | this = TVariableAccess(access, _) |
+      result = access.getLocation()
+    )
+    or
+    exists(CapturingCallable c |
+      this = TLambdaAccess(_, _, c) and
+      result = c.getLocation()
+    )
+  }
+
+  Cfg::CfgScope getCfgScope() {
+    exists(Cfg::EntryBasicBlock bb, int i, CapturingCallable v |
+      this = TEntryDef(bb, i, v) and
+      result = bb.getScope()
+    )
+    or
+    exists(VariableAccessCfgNode access | this = TVariableAccess(access, _) |
+      result = access.getScope()
+    )
+    or
+    exists(CapturingCallable c |
+      this = TLambdaAccess(_, _, c) and
+      result = c.getCfgScope()
+    )
+  }
+
+  predicate isLambdaAccess(Cfg::BasicBlock bb, int i, CapturingCallable v) {
+    captureRead2(bb, i, v) and
+    this = TLambdaAccess(bb, i, v)
+  }
+
+  predicate isRead(Cfg::BasicBlock bb, int i, CapturingCallable v) {
+    exists(VariableAccessCfgNode acc |
+      captureRead1(bb, i, _, acc, v) and
+      this = TVariableAccess(acc, v)
+    )
+    or
+    this.isLambdaAccess(bb, i, v)
+  }
+
+  predicate isWrite(Cfg::BasicBlock bb, int i, CapturingCallable v) {
+    entryWrite(bb, i, v) and
+    this = TEntryDef(bb, i, v)
+  }
+
+  predicate isReadOrWrite(Cfg::BasicBlock bb, int i, CapturingCallable v) {
+    this.isRead(bb, i, v) or
+    this.isWrite(bb, i, v)
+  }
+}
+
+private module SsaInput implements SsaImplCommon::InputSig {
+  private import codeql.ruby.controlflow.BasicBlocks as BasicBlocks
+
+  class BasicBlock = BasicBlocks::BasicBlock;
+
+  BasicBlock getImmediateBasicBlockDominator(BasicBlock bb) { result = bb.getImmediateDominator() }
+
+  BasicBlock getABasicBlockSuccessor(BasicBlock bb) { result = bb.getASuccessor() }
+
+  class ExitBasicBlock = BasicBlocks::ExitBasicBlock;
+
+  class SourceVariable = CapturingCallable;
+
+  /**
+   * Holds if the statement at index `i` of basic block `bb` contains a write to variable `v`.
+   * `certain` is true if the write definitely occurs.
+   */
+  predicate variableWrite(BasicBlock bb, int i, SourceVariable v, boolean certain) {
+    any(CaptureAccess access).isWrite(bb, i, v) and
+    certain = true
+  }
+
+  predicate variableRead(BasicBlock bb, int i, SourceVariable v, boolean certain) {
+    any(CaptureAccess access).isRead(bb, i, v) and
+    certain = true
+  }
+}
+
+private import SsaImplCommon::Make<SsaInput> as Impl
+
+class Definition = Impl::Definition;
+
+predicate adjacentStep(CaptureAccess access1, CaptureAccess access2) {
+  exists(
+    Definition def, SsaInput::BasicBlock bb1, int i1, SsaInput::BasicBlock bb2, int i2,
+    SsaInput::SourceVariable v
+  |
+    Impl::adjacentDefRead(def, bb1, i1, bb2, i2) and
+    v = def.getSourceVariable() and
+    access1.isReadOrWrite(bb1, i1, v) and
+    access2.isReadOrWrite(bb2, i2, v)
+  )
+}