Ruby: Cleanup flow through self

Author: hvitved

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

@@ -24,21 +24,8 @@ class Expr extends Stmt, TExpr {
   }
 }
 
-/**
- * A reference to the current object. For example:
- * - `self == other`
- * - `self.method_name`
- * - `def self.method_name ... end`
- *
- * This also includes implicit references to the current object in method
- * calls.  For example, the method call `foo(123)` has an implicit `self`
- * receiver, and is equivalent to the explicit `self.foo(123)`.
- */
-class Self extends Expr, TSelf {
-  final override string getAPrimaryQlClass() { result = "Self" }
-
-  final override string toString() { result = "self" }
-}
+/** DEPRECATED: Use `SelfVariableAccess` instead. */
+deprecated class Self = SelfVariableAccess;
 
 /**
  * A sequence of expressions in the right-hand side of an assignment or

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

@@ -201,7 +201,16 @@ class ClassVariableWriteAccess extends ClassVariableAccess, VariableWriteAccess
 /** An access to a class variable where the value is read. */
 class ClassVariableReadAccess extends ClassVariableAccess, VariableReadAccess { }
 
-/** An access to the `self` variable */
+/**
+ * An access to the `self` variable. For example:
+ * - `self == other`
+ * - `self.method_name`
+ * - `def self.method_name ... end`
+ *
+ * This also includes implicit references to the current object in method
+ * calls.  For example, the method call `foo(123)` has an implicit `self`
+ * receiver, and is equivalent to the explicit `self.foo(123)`.
+ */
 class SelfVariableAccess extends LocalVariableAccess instanceof SelfVariableAccessImpl {
   final override string getAPrimaryQlClass() { result = "SelfVariableAccess" }
 }

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

@@ -73,7 +73,7 @@ private module Cached {
       m = resolveConstantReadAccess(c.getReceiver())
       or
       m = enclosingModule(c).getModule() and
-      c.getReceiver() instanceof Self
+      c.getReceiver() instanceof SelfVariableAccess
     ) and
     result = resolveConstantReadAccess(c.getAnArgument())
   }
@@ -437,7 +437,7 @@ private module ResolveImpl {
         encl = enclosingModule(this) and
         result = [qualifiedModuleNameNonRec(encl, _, _), qualifiedModuleNameRec(encl, _, _)]
       |
-        this.getReceiver() instanceof Self
+        this.getReceiver() instanceof SelfVariableAccess
         or
         not exists(this.getReceiver())
       )

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

@@ -366,7 +366,23 @@ private module Cached {
 
   cached
   predicate isCapturedAccess(LocalVariableAccess access) {
-    access.getVariable().getDeclaringScope() != access.getCfgScope()
+    exists(Scope scope1, Scope scope2 |
+      scope1 = access.getVariable().getDeclaringScope() and
+      scope2 = access.getCfgScope() and
+      scope1 != scope2
+    |
+      if access instanceof SelfVariableAccess
+      then
+        // ```
+        // class C
+        //   def self.m // not a captured access
+        //   end
+        // end
+        // ```
+        not scope2 instanceof Toplevel or
+        not access = any(SingletonMethod m).getObject()
+      else any()
+    )
   }
 
   cached

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

@@ -233,6 +233,8 @@ module Ssa {
       )
     }
 
+    override SelfVariable getSourceVariable() { result = v }
+
     final override string toString() { result = "self (" + v.getDeclaringScope() + ")" }
 
     final override Location getLocation() { result = this.getControlFlowNode().getLocation() }

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

@@ -203,7 +203,7 @@ private module Cached {
           result = lookupMethod(tp, method) and
           if result.(Method).isPrivate()
           then
-            exists(Self self |
+            exists(SelfVariableAccess self |
               self = call.getReceiver().getExpr() and
               pragma[only_bind_out](self.getEnclosingModule().getModule().getSuperClass*()) =
                 pragma[only_bind_out](result.getEnclosingModule().getModule())
@@ -300,28 +300,14 @@ private DataFlow::LocalSourceNode trackInstance(Module tp, TypeTracker t) {
     )
     or
     // `self` in method
-    exists(Self self, Method enclosing |
-      self = result.asExpr().getExpr() and
-      enclosing = self.getEnclosingMethod() and
-      tp = enclosing.getEnclosingModule().getModule() and
-      not self.getEnclosingModule().getEnclosingMethod() = enclosing
-    )
+    tp = result.(SsaSelfDefinitionNode).getSelfScope().(Method).getEnclosingModule().getModule()
     or
     // `self` in singleton method
-    exists(Self self, MethodBase enclosing |
-      self = result.asExpr().getExpr() and
-      flowsToSingletonMethodObject(trackInstance(tp), enclosing) and
-      enclosing = self.getEnclosingMethod() and
-      not self.getEnclosingModule().getEnclosingMethod() = enclosing
-    )
+    flowsToSingletonMethodObject(trackInstance(tp), result.(SsaSelfDefinitionNode).getSelfScope())
     or
     // `self` in top-level
-    exists(Self self, Toplevel enclosing |
-      self = result.asExpr().getExpr() and
-      enclosing = self.getEnclosingModule() and
-      tp = TResolved("Object") and
-      not self.getEnclosingMethod().getEnclosingModule() = enclosing
-    )
+    result.(SsaSelfDefinitionNode).getSelfScope() instanceof Toplevel and
+    tp = TResolved("Object")
     or
     // a module or class
     exists(Module m |
@@ -371,7 +357,7 @@ private predicate singletonMethod(MethodBase method, Expr object) {
 
 pragma[nomagic]
 private predicate flowsToSingletonMethodObject(DataFlow::LocalSourceNode nodeFrom, MethodBase method) {
-  exists(DataFlow::LocalSourceNode nodeTo |
+  exists(DataFlow::Node nodeTo |
     nodeFrom.flowsTo(nodeTo) and
     singletonMethod(method, nodeTo.asExpr().getExpr())
   )
@@ -409,13 +395,8 @@ private DataFlow::LocalSourceNode trackSingletonMethod(MethodBase m, string name
   name = m.getName()
 }
 
-private DataFlow::Node selfInModule(Module tp) {
-  exists(Self self, ModuleBase enclosing |
-    self = result.asExpr().getExpr() and
-    enclosing = self.getEnclosingModule() and
-    tp = enclosing.getModule() and
-    not self.getEnclosingMethod().getEnclosingModule() = enclosing
-  )
+private SsaSelfDefinitionNode selfInModule(Module tp) {
+  tp = result.getSelfScope().(ModuleBase).getModule()
 }
 
 private DataFlow::LocalSourceNode trackModule(Module tp, TypeTracker t) {

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

@@ -70,6 +70,20 @@ module LocalFlow {
     )
   }
 
+  /** Gets the SSA definition node corresponding to the implicit `self` parameter for `m`. */
+  private SsaDefinitionNode getSelfParameterDefNode(MethodBase m) {
+    result.getDefinition().(Ssa::SelfDefinition).getSourceVariable().getDeclaringScope() = m
+  }
+
+  /**
+   * Holds if `nodeFrom` is a parameter node, and `nodeTo` is a corresponding SSA node.
+   */
+  predicate localFlowSsaParamInput(Node nodeFrom, Node nodeTo) {
+    nodeTo = getParameterDefNode(nodeFrom.(ParameterNode).getParameter())
+    or
+    nodeTo = getSelfParameterDefNode(nodeFrom.(SelfParameterNode).getMethod())
+  }
+
   /**
    * Holds if there is a local use-use flow step from `nodeFrom` to `nodeTo`
    * involving SSA definition `def`.
@@ -115,9 +129,6 @@ module LocalFlow {
   predicate localFlowStepCommon(Node nodeFrom, Node nodeTo) {
     localSsaFlowStep(nodeFrom, nodeTo)
     or
-    nodeFrom.(SelfParameterNode).getMethod() = nodeTo.asExpr().getExpr().getEnclosingCallable() and
-    nodeTo.asExpr().getExpr() instanceof Self
-    or
     nodeFrom.asExpr() = nodeTo.asExpr().(CfgNodes::ExprNodes::AssignExprCfgNode).getRhs()
     or
     nodeFrom.asExpr() = nodeTo.asExpr().(CfgNodes::ExprNodes::BlockArgumentCfgNode).getValue()
@@ -236,7 +247,7 @@ private module Cached {
     or
     defaultValueFlow(nodeTo.(ParameterNode).getParameter(), nodeFrom)
     or
-    nodeTo = LocalFlow::getParameterDefNode(nodeFrom.(ParameterNode).getParameter())
+    LocalFlow::localFlowSsaParamInput(nodeFrom, nodeTo)
     or
     nodeTo.(SynthReturnNode).getAnInput() = nodeFrom
     or
@@ -253,7 +264,7 @@ private module Cached {
     or
     defaultValueFlow(nodeTo.(ParameterNode).getParameter(), nodeFrom)
     or
-    nodeTo = LocalFlow::getParameterDefNode(nodeFrom.(ParameterNode).getParameter())
+    LocalFlow::localFlowSsaParamInput(nodeFrom, nodeTo)
     or
     LocalFlow::localSsaFlowStepUseUse(_, nodeFrom, nodeTo)
     or
@@ -275,27 +286,34 @@ private module Cached {
     LocalFlow::localSsaFlowStepUseUse(_, nodeFrom, nodeTo)
   }
 
+  private predicate entrySsaDefinition(SsaDefinitionNode n) {
+    n = LocalFlow::getParameterDefNode(_)
+    or
+    exists(Ssa::Definition def | def = n.getDefinition() |
+      def instanceof Ssa::SelfDefinition
+      or
+      def instanceof Ssa::CapturedEntryDefinition
+    )
+  }
+
   cached
   predicate isLocalSourceNode(Node n) {
     n instanceof ParameterNode
     or
-    // This case should not be needed once we have proper use-use flow
-    // for `self`. At that point, the `self`s returned by `trackInstance`
-    // in `DataFlowDispatch.qll` should refer to the post-update node,
-    // and we can remove this case.
-    n.asExpr().getExpr() instanceof Self
+    n instanceof PostUpdateNodes::ExprPostUpdateNode
     or
-    // Nodes that can't be reached from another parameter or expression.
-    not localFlowStepTypeTracker+(any(Node e |
-        e instanceof ExprNode
+    // Expressions that can't be reached from another entry definition or expression.
+    not localFlowStepTypeTracker+(any(Node n0 |
+        n0 instanceof ExprNode
         or
-        e instanceof ParameterNode
-      ), n)
+        entrySsaDefinition(n0)
+      ), n.(ExprNode))
     or
-    // Ensure all parameter SSA nodes are local sources -- this is needed by type tracking.
-    // Note that when the parameter has a default value, it will be reachable from an
-    // expression (the default value) and therefore won't be caught by the rule above.
-    n = LocalFlow::getParameterDefNode(_)
+    // Ensure all entry SSA definitions are local sources -- for parameters, this
+    // is needed by type tracking. Note that when the parameter has a default value,
+    // it will be reachable from an expression (the default value) and therefore
+    // won't be caught by the rule above.
+    entrySsaDefinition(n)
   }
 
   cached
@@ -358,6 +376,16 @@ class SsaDefinitionNode extends NodeImpl, TSsaDefinitionNode {
   override string toStringImpl() { result = def.toString() }
 }
 
+/** An SSA definition for a `self` variable. */
+class SsaSelfDefinitionNode extends LocalSourceNode, SsaDefinitionNode {
+  private SelfVariable self;
+
+  SsaSelfDefinitionNode() { self = def.getSourceVariable() }
+
+  /** Gets the scope in which the `self` variable is declared. */
+  Scope getSelfScope() { result = self.getDeclaringScope() }
+}
+
 /**
  * A value returning statement, viewed as a node in a data flow graph.
  *
@@ -745,13 +773,6 @@ predicate jumpStep(Node pred, Node succ) {
   SsaImpl::captureFlowOut(pred.(SsaDefinitionNode).getDefinition(),
     succ.(SsaDefinitionNode).getDefinition())
   or
-  exists(Self s, Method m |
-    s = succ.asExpr().getExpr() and
-    pred.(SelfParameterNode).getMethod() = m and
-    m = s.getEnclosingMethod() and
-    m != s.getEnclosingCallable()
-  )
-  or
   succ.asExpr().getExpr().(ConstantReadAccess).getValue() = pred.asExpr().getExpr()
 }
 

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

@@ -1,4 +1,5 @@
 private import SsaImplCommon
+private import SsaImplSpecific as SsaImplSpecific
 private import codeql.ruby.AST
 private import codeql.ruby.CFG
 private import codeql.ruby.ast.Variable
@@ -57,14 +58,16 @@ private predicate hasCapturedRead(Variable v, CfgScope scope) {
   ).isCapturedAccess()
 }
 
+pragma[noinline]
+private predicate variableWriteInOuterScope(BasicBlock bb, LocalVariable v, CfgScope scope) {
+  SsaImplSpecific::variableWrite(bb, _, v, _) and
+  scope.getOuterCfgScope() = bb.getScope()
+}
+
 pragma[noinline]
 private predicate hasVariableWriteWithCapturedRead(BasicBlock bb, LocalVariable v, CfgScope scope) {
   hasCapturedRead(v, scope) and
-  exists(VariableWriteAccess write |
-    write = bb.getANode().getNode() and
-    write.getVariable() = v and
-    bb.getScope() = scope.getOuterCfgScope()
-  )
+  variableWriteInOuterScope(bb, v, scope)
 }
 
 /**

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

@@ -105,7 +105,7 @@ private class ActionControllerContextCall extends MethodCall {
   private ActionControllerControllerClass controllerClass;
 
   ActionControllerContextCall() {
-    this.getReceiver() instanceof Self and
+    this.getReceiver() instanceof SelfVariableAccess and
     this.getEnclosingModule() = controllerClass
   }
 

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

@@ -61,7 +61,7 @@ private class ActionViewHtmlEscapeCall extends HtmlEscapeCall {
 // A call in a context where some commonly used `ActionView` methods are available.
 private class ActionViewContextCall extends MethodCall {
   ActionViewContextCall() {
-    this.getReceiver() instanceof Self and
+    this.getReceiver() instanceof SelfVariableAccess and
     inActionViewContext(this)
   }