Refactor Ruby callable body identity for methods and lambdas

Agent-Logs-Url: https://github.com/github/codeql/sessions/c8fcf73a-5bb2-4182-b8fb-a251eec43ef4

Co-authored-by: aschackmull <28296824+aschackmull@users.noreply.github.com>

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

@@ -8,7 +8,7 @@ private import internal.TreeSitter
 private import internal.Method
 
 /** A callable. */
-class Callable extends StmtSequence, Expr, Scope, TCallable {
+class Callable extends Expr, Scope, TCallable {
   /** Gets the number of parameters of this callable. */
   final int getNumberOfParameters() { result = count(this.getAParameter()) }
 
@@ -26,7 +26,10 @@ class Callable extends StmtSequence, Expr, Scope, TCallable {
 }
 
 /** A method. */
-class MethodBase extends Callable, BodyStmt, Scope, TMethodBase {
+class MethodBase extends Callable, Scope, TMethodBase {
+  /** Gets the body of this method. */
+  BodyStmt getBody() { none() }
+
   /** Gets the name of this method. */
   string getName() { none() }
 
@@ -36,7 +39,7 @@ class MethodBase extends Callable, BodyStmt, Scope, TMethodBase {
   override AstNode getAChild(string pred) {
     result = Callable.super.getAChild(pred)
     or
-    result = BodyStmt.super.getAChild(pred)
+    pred = "getBody" and result = this.getBody()
   }
 
   /**
@@ -218,6 +221,8 @@ class Method extends MethodBase, TMethod {
     toGenerated(result) = g.getParameters().getChild(n)
   }
 
+  final override BodyStmt getBody() { synthChild(this, this.getNumberOfParameters(), result) }
+
   final override string toString() { result = this.getName() }
 
   overlay[global]
@@ -280,6 +285,8 @@ class SingletonMethod extends MethodBase, TSingletonMethod {
     toGenerated(result) = g.getParameters().getChild(n)
   }
 
+  final override BodyStmt getBody() { synthChild(this, this.getNumberOfParameters() + 1, result) }
+
   final override string toString() { result = this.getName() }
 
   final override AstNode getAChild(string pred) {
@@ -321,7 +328,7 @@ class SingletonMethod extends MethodBase, TSingletonMethod {
  * -> (x) { x + 1 }
  * ```
  */
-class Lambda extends Callable, BodyStmt, TLambda {
+class Lambda extends Callable, TLambda {
   private Ruby::Lambda g;
 
   Lambda() { this = TLambda(g) }
@@ -332,12 +339,14 @@ class Lambda extends Callable, BodyStmt, TLambda {
     toGenerated(result) = g.getParameters().getChild(n)
   }
 
+  final BodyStmt getBody() { synthChild(this, this.getNumberOfParameters(), result) }
+
   final override string toString() { result = "-> { ... }" }
 
   final override AstNode getAChild(string pred) {
     result = Callable.super.getAChild(pred)
     or
-    result = BodyStmt.super.getAChild(pred)
+    pred = "getBody" and result = this.getBody()
   }
 }
 

ruby/ql/lib/codeql/ruby/ast/internal/AST.qll

@@ -198,6 +198,7 @@ private module Cached {
     TLShiftExprSynth(Ast::AstNode parent, int i) { mkSynthChild(LShiftExprKind(), parent, i) } or
     TLTExpr(Ruby::Binary g) { g instanceof @ruby_binary_langle } or
     TLambda(Ruby::Lambda g) or
+    TLambdaBodyStmt(Ast::Lambda parent, int i) { i = parent.getNumberOfParameters() } or
     TLine(Ruby::Line g) or
     TLeftAssignmentList(Ruby::LeftAssignmentList g) or
     TLocalVariableAccessReal(Ruby::Identifier g, TLocalVariableReal v) {
@@ -220,6 +221,7 @@ private module Cached {
     TLogicalOrExprSynth(Ast::AstNode parent, int i) { mkSynthChild(LogicalOrExprKind(), parent, i) } or
     TMatchPattern(Ruby::MatchPattern g) or
     TMethod(Ruby::Method g) or
+    TMethodBodyStmt(Ast::Method parent, int i) { i = parent.getNumberOfParameters() } or
     TMethodCallSynth(Ast::AstNode parent, int i, string name, boolean setter, int arity) {
       mkSynthChild(MethodCallKind(name, setter, arity), parent, i)
     } or
@@ -285,6 +287,9 @@ private module Cached {
     } or
     TSingletonClass(Ruby::SingletonClass g) or
     TSingletonMethod(Ruby::SingletonMethod g) or
+    TSingletonMethodBodyStmt(Ast::SingletonMethod parent, int i) {
+      i = parent.getNumberOfParameters() + 1
+    } or
     TSpaceshipExpr(Ruby::Binary g) { g instanceof @ruby_binary_langleequalrangle } or
     TSplatExprReal(Ruby::SplatArgument g) or
     TSplatExprSynth(Ast::AstNode parent, int i) { mkSynthChild(SplatExprKind(), parent, i) } or
@@ -651,6 +656,12 @@ private module Cached {
     result = TPairSynth(parent, i)
     or
     result = TSimpleSymbolLiteralSynth(parent, i, _)
+    or
+    result = TMethodBodyStmt(parent, i)
+    or
+    result = TLambdaBodyStmt(parent, i)
+    or
+    result = TSingletonMethodBodyStmt(parent, i)
   }
 
   /**
@@ -759,7 +770,9 @@ class TStmtSequence =
   TBeginBlock or TEndBlock or TThen or TElse or TDo or TEnsure or TStringInterpolationComponent or
       TBlock or TBodyStmt or TParenthesizedExpr or TStmtSequenceSynth;
 
-class TBodyStmt = TBeginExpr or TModuleBase or TMethod or TLambda or TDoBlock or TSingletonMethod;
+class TBodyStmt =
+  TBeginExpr or TModuleBase or TDoBlock or TMethodBodyStmt or TLambdaBodyStmt or
+      TSingletonMethodBodyStmt;
 
 class TNilLiteral = TNilLiteralReal or TNilLiteralSynth;
 

ruby/ql/lib/codeql/ruby/ast/internal/Expr.qll

@@ -64,21 +64,30 @@ class Ensure extends StmtSequence, TEnsure {
 
 // Not defined by dispatch, as it should not be exposed
 Ruby::AstNode getBodyStmtChild(TBodyStmt b, int i) {
-  exists(Ruby::Method g, Ruby::AstNode body | b = TMethod(g) and body = g.getBody() |
-    result = body.(Ruby::BodyStatement).getChild(i)
-    or
-    i = 0 and result = body and not body instanceof Ruby::BodyStatement
-  )
-  or
-  exists(Ruby::SingletonMethod g, Ruby::AstNode body |
-    b = TSingletonMethod(g) and body = g.getBody()
+  exists(Method m, Ruby::Method g, Ruby::AstNode body |
+    b = TMethodBodyStmt(m, _) and
+    m = TMethod(g) and
+    body = g.getBody()
   |
     result = body.(Ruby::BodyStatement).getChild(i)
     or
     i = 0 and result = body and not body instanceof Ruby::BodyStatement
   )
   or
-  exists(Ruby::Lambda g | b = TLambda(g) |
+  exists(SingletonMethod m, Ruby::SingletonMethod g, Ruby::AstNode body |
+    b = TSingletonMethodBodyStmt(m, _) and
+    m = TSingletonMethod(g) and
+    body = g.getBody()
+  |
+    result = body.(Ruby::BodyStatement).getChild(i)
+    or
+    i = 0 and result = body and not body instanceof Ruby::BodyStatement
+  )
+  or
+  exists(Lambda l, Ruby::Lambda g |
+    b = TLambdaBodyStmt(l, _) and
+    l = TLambda(g)
+  |
     result = g.getBody().(Ruby::DoBlock).getBody().getChild(i) or
     result = g.getBody().(Ruby::Block).getBody().getChild(i)
   )

ruby/ql/lib/codeql/ruby/controlflow/internal/ControlFlowGraphImpl.qll

@@ -100,7 +100,7 @@ private class EndBlockScope extends CfgScopeImpl, EndBlock {
   }
 }
 
-private class BodyStmtCallableScope extends CfgScopeImpl, AstInternal::TBodyStmt, Callable {
+private class BodyStmtCallableScope extends CfgScopeImpl, TDoBlock, Callable {
   final override predicate entry(AstNode first) { this.(Trees::BodyStmtTree).firstInner(first) }
 
   final override predicate exit(AstNode last, Completion c) {
@@ -108,6 +108,22 @@ private class BodyStmtCallableScope extends CfgScopeImpl, AstInternal::TBodyStmt
   }
 }
 
+private class MethodBodyScope extends CfgScopeImpl, TMethodBase {
+  final override predicate entry(AstNode first) { this.getBody().(Trees::BodyStmtTree).firstInner(first) }
+
+  final override predicate exit(AstNode last, Completion c) {
+    this.getBody().(Trees::BodyStmtTree).lastInner(last, c)
+  }
+}
+
+private class LambdaBodyScope extends CfgScopeImpl, TLambda {
+  final override predicate entry(AstNode first) { this.getBody().(Trees::BodyStmtTree).firstInner(first) }
+
+  final override predicate exit(AstNode last, Completion c) {
+    this.getBody().(Trees::BodyStmtTree).lastInner(last, c)
+  }
+}
+
 private class BraceBlockScope extends CfgScopeImpl, BraceBlock {
   final override predicate entry(AstNode first) {
     first(this.(Trees::BraceBlockTree).getBodyChild(0, _), first)
@@ -1073,14 +1089,18 @@ module Trees {
     final override AstNode getAccessNode() { result = super.getDefiningAccess() }
   }
 
-  private class LambdaTree extends BodyStmtTree instanceof Lambda {
+  private class LambdaTree extends LeafTree instanceof Lambda { }
+
+  private class LambdaBodyStmtTree extends BodyStmtTree instanceof AstInternal::TLambdaBodyStmt {
     final override predicate propagatesAbnormal(AstNode child) { none() }
 
+    private Lambda getLambda() { result = this.getParent() }
+
     /** Gets the `i`th child in the body of this block. */
     final override AstNode getBodyChild(int i, boolean rescuable) {
-      result = super.getParameter(i) and rescuable = false
+      result = this.getLambda().getParameter(i) and rescuable = false
       or
-      result = BodyStmtTree.super.getBodyChild(i - super.getNumberOfParameters(), rescuable)
+      result = BodyStmtTree.super.getBodyChild(i - this.getLambda().getNumberOfParameters(), rescuable)
     }
   }
 
@@ -1151,14 +1171,18 @@ module Trees {
   private class MethodNameTree extends LeafTree instanceof MethodName, AstInternal::TTokenMethodName
   { }
 
-  private class MethodTree extends BodyStmtTree instanceof Method {
+  private class MethodTree extends LeafTree instanceof Method { }
+
+  private class MethodBodyStmtTree extends BodyStmtTree instanceof AstInternal::TMethodBodyStmt {
     final override predicate propagatesAbnormal(AstNode child) { none() }
 
+    private Method getMethod() { result = this.getParent() }
+
     /** Gets the `i`th child in the body of this block. */
     final override AstNode getBodyChild(int i, boolean rescuable) {
-      result = super.getParameter(i) and rescuable = false
+      result = this.getMethod().getParameter(i) and rescuable = false
       or
-      result = BodyStmtTree.super.getBodyChild(i - super.getNumberOfParameters(), rescuable)
+      result = BodyStmtTree.super.getBodyChild(i - this.getMethod().getNumberOfParameters(), rescuable)
     }
   }
 
@@ -1351,24 +1375,23 @@ module Trees {
     }
   }
 
-  private class SingletonMethodTree extends BodyStmtTree instanceof SingletonMethod {
+  private class SingletonMethodTree extends StandardPostOrderTree instanceof SingletonMethod {
+    final override ControlFlowTree getChildNode(int i) { result = super.getObject() and i = 0 }
+  }
+
+  private class SingletonMethodBodyStmtTree extends BodyStmtTree instanceof
+    AstInternal::TSingletonMethodBodyStmt
+  {
     final override predicate propagatesAbnormal(AstNode child) { none() }
 
+    private SingletonMethod getMethod() { result = this.getParent() }
+
     /** Gets the `i`th child in the body of this block. */
     final override AstNode getBodyChild(int i, boolean rescuable) {
-      result = super.getParameter(i) and rescuable = false
+      result = this.getMethod().getParameter(i) and rescuable = false
       or
-      result = BodyStmtTree.super.getBodyChild(i - super.getNumberOfParameters(), rescuable)
-    }
-
-    override predicate first(AstNode first) { first(super.getObject(), first) }
-
-    override predicate succ(AstNode pred, AstNode succ, Completion c) {
-      BodyStmtTree.super.succ(pred, succ, c)
-      or
-      last(super.getObject(), pred, c) and
-      succ = this and
-      c instanceof NormalCompletion
+      result =
+        BodyStmtTree.super.getBodyChild(i - this.getMethod().getNumberOfParameters(), rescuable)
     }
   }
 

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

@@ -1656,13 +1656,21 @@ private module ReturnNodes {
     result = implicitReturn(c, n).getParent()
   }
 
+  private Stmt getACallableBodyStmt(Callable c) {
+    result = c.(MethodBase).getBody().getAStmt()
+    or
+    result = c.(Lambda).getBody().getAStmt()
+    or
+    result = c.(StmtSequence).getAStmt()
+  }
+
   /**
    * A data-flow node that represents an expression implicitly returned by
    * a callable. An implicit return happens when an expression can be the
    * last thing that is evaluated in the body of the callable.
    */
   class ExprReturnNode extends SourceReturnNode, ExprNode {
-    ExprReturnNode() { exists(Callable c | implicitReturn(c, this) = c.getAStmt()) }
+    ExprReturnNode() { exists(Callable c | implicitReturn(c, this) = getACallableBodyStmt(c)) }
 
     override ReturnKind getKindSource() {
       exists(CfgScope scope | scope = this.(NodeImpl).getCfgScope() |

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

@@ -1392,7 +1392,7 @@ class StmtSequenceNode extends ExprNode {
 /**
  * A data flow node corresponding to a method, block, or lambda expression.
  */
-class CallableNode extends StmtSequenceNode {
+class CallableNode extends ExprNode {
   private Callable callable;
 
   CallableNode() { this.asExpr().getExpr() = callable }
@@ -1400,6 +1400,17 @@ class CallableNode extends StmtSequenceNode {
   /** Gets the underlying AST node as a `Callable`. */
   Callable asCallableAstNode() { result = callable }
 
+  private StmtSequence getBodySequence() {
+    result = callable.(MethodBase).getBody()
+    or
+    result = callable.(Lambda).getBody()
+    or
+    result = callable.(StmtSequence)
+  }
+
+  /** Gets the last statement in this callable body, if any. */
+  ExprNode getLastStmt() { result.asExpr().getExpr() = this.getBodySequence().getLastStmt() }
+
   private ParameterPosition getParameterPosition(ParameterNodeImpl node) {
     result = getSourceParameterPosition(node, callable)
   }

ruby/ql/lib/codeql/ruby/frameworks/Slim.qll

@@ -18,7 +18,7 @@ module Slim {
 
     override DataFlow::Node getTemplate() {
       result.asExpr().getExpr() =
-        this.getBlock().(DataFlow::BlockNode).asCallableAstNode().getAStmt()
+        this.getBlock().(DataFlow::BlockNode).asCallableAstNode().getLastStmt()
     }
   }
 

ruby/ql/lib/codeql/ruby/security/ImproperMemoizationQuery.qll

@@ -70,7 +70,7 @@ private predicate memoReturnedFromMethod(Method m, MemoStmt s) {
   or
   // If we don't have flow (e.g. due to the dataflow library not supporting instance variable flow yet),
   // fall back to a syntactic heuristic: does the last statement in the method mention the memoization variable?
-  m.getLastStmt().getAChild*().(InstanceVariableReadAccess).getVariable() =
+  m.getBody().getLastStmt().getAChild*().(InstanceVariableReadAccess).getVariable() =
     s.getVariableAccess().getVariable()
 }