Data flow: Add a synthetic return node

Author: hvitved

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

@@ -436,7 +436,7 @@ DataFlowCallable viableImplInCallContext(DataFlowCall call, DataFlowCall ctx) {
  * Holds if `e` is an `ExprNode` that may be returned by a call to `c`.
  */
 predicate exprNodeReturnedFrom(DataFlow::ExprNode e, Callable c) {
-  exists(ReturnNode r |
+  exists(ReturningNode r |
     r.getEnclosingCallable().asCallable() = c and
     (
       r.(ExplicitReturnNode).getReturningNode().getReturnedValueNode() = e.asExpr() or

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

@@ -48,7 +48,7 @@ module LocalFlow {
 
   /**
    * Holds if there is a local flow step from `nodeFrom` to `nodeTo` involving
-   * SSA definition `def.
+   * SSA definition `def`.
    */
   predicate localSsaFlowStep(Ssa::Definition def, Node nodeFrom, Node nodeTo) {
     // Flow from parameter into SSA definition
@@ -114,6 +114,12 @@ private module Cached {
   newtype TNode =
     TExprNode(CfgNodes::ExprCfgNode n) or
     TReturningNode(CfgNodes::ReturningCfgNode n) or
+    TSynthReturnNode(CfgScope scope, ReturnKind kind) {
+      exists(ReturningNode ret |
+        ret.(NodeImpl).getCfgScope() = scope and
+        ret.getKind() = kind
+      )
+    } or
     TSsaDefinitionNode(Ssa::Definition def) or
     TNormalParameterNode(Parameter p) { not p instanceof BlockParameter } or
     TSelfParameterNode(MethodBase m) or
@@ -140,14 +146,12 @@ private module Cached {
     TNormalParameterNode or TBlockParameterNode or TSelfParameterNode or TSummaryParameterNode;
 
   /**
-   * 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
-   * is handled by the global data-flow library, but includes various other steps
-   * that are only relevant for global flow.
+   * This is the local flow predicate that is shared between local data flow
+   * and global data flow.
    */
   cached
-  predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo) {
-    exists(Ssa::Definition def | LocalFlow::localSsaFlowStep(def, nodeFrom, nodeTo))
+  predicate simpleLocalFlowStepCommon(Node nodeFrom, Node nodeTo) {
+    LocalFlow::localSsaFlowStep(_, nodeFrom, nodeTo)
     or
     nodeTo.(ParameterNode).getParameter().(OptionalParameter).getDefaultValue() =
       nodeFrom.asExpr().getExpr()
@@ -186,12 +190,39 @@ private module Cached {
         ) and
         nodeFrom.asExpr() = for.getValue()
       )
+  }
+
+  /**
+   * 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
+   * is handled by the global data-flow library, but includes various other steps
+   * that are only relevant for global flow.
+   */
+  cached
+  predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo) {
+    simpleLocalFlowStepCommon(nodeFrom, nodeTo)
+    or
+    nodeTo.(SynthReturnNode).getAnInput() = nodeFrom
     or
     FlowSummaryImpl::Private::Steps::summaryLocalStep(nodeFrom, nodeTo, true)
   }
 
   cached
-  predicate isLocalSourceNode(Node n) { not simpleLocalFlowStep+(any(ExprNode e), n) }
+  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 instanceof SelfArgumentNode
+    or
+    not simpleLocalFlowStepCommon+(any(Node e |
+        e instanceof ExprNode
+        or
+        e instanceof ParameterNode
+      ), n)
+  }
 
   cached
   newtype TContent = TTodoContent() // stub
@@ -208,6 +239,8 @@ predicate nodeIsHidden(Node n) {
   n instanceof SummaryNode
   or
   n instanceof SummaryParameterNode
+  or
+  n instanceof SynthReturnNode
 }
 
 /** An SSA definition, viewed as a node in a data flow graph. */
@@ -234,7 +267,7 @@ class SsaDefinitionNode extends NodeImpl, TSsaDefinitionNode {
  * `ControlFlow::Node`s.
  */
 class ReturningStatementNode extends NodeImpl, TReturningNode {
-  private CfgNodes::ReturningCfgNode n;
+  CfgNodes::ReturningCfgNode n;
 
   ReturningStatementNode() { this = TReturningNode(n) }
 
@@ -436,6 +469,12 @@ private module ArgumentNodes {
 
 import ArgumentNodes
 
+/** A data-flow node that represents a value syntactically returned by a callable. */
+abstract class ReturningNode extends Node {
+  /** Gets the kind of this return node. */
+  abstract ReturnKind getKind();
+}
+
 /** A data-flow node that represents a value returned by a callable. */
 abstract class ReturnNode extends Node {
   /** Gets the kind of this return node. */
@@ -463,11 +502,9 @@ private module ReturnNodes {
    * A data-flow node that represents an expression returned by a callable,
    * either using an explict `return` statement or as the expression of a method body.
    */
-  class ExplicitReturnNode extends ReturnNode, ReturningStatementNode {
-    private CfgNodes::ReturningCfgNode n;
-
+  class ExplicitReturnNode extends ReturningNode, ReturningStatementNode {
     ExplicitReturnNode() {
-      isValid(this.getReturningNode()) and
+      isValid(n) and
       n.getASuccessor().(CfgNodes::AnnotatedExitNode).isNormal() and
       n.getScope() instanceof Callable
     }
@@ -479,7 +516,7 @@ private module ReturnNodes {
     }
   }
 
-  class ExprReturnNode extends ReturnNode, ExprNode {
+  class ExprReturnNode extends ReturningNode, ExprNode {
     ExprReturnNode() {
       this.getExprNode().getASuccessor().(CfgNodes::AnnotatedExitNode).isNormal() and
       this.(NodeImpl).getCfgScope() instanceof Callable
@@ -488,6 +525,34 @@ private module ReturnNodes {
     override ReturnKind getKind() { result instanceof NormalReturnKind }
   }
 
+  /**
+   * A synthetic data-flow node for joining flow from different syntactic
+   * returns into a single node.
+   *
+   * This node only exists to avoid computing the product of a large fan-in
+   * with a large fan-out.
+   */
+  class SynthReturnNode extends NodeImpl, ReturnNode, TSynthReturnNode {
+    private CfgScope scope;
+    private ReturnKind kind;
+
+    SynthReturnNode() { this = TSynthReturnNode(scope, kind) }
+
+    /** Get a syntactic return node that flows into this synthetic node. */
+    ReturningNode getAnInput() {
+      result.(NodeImpl).getCfgScope() = scope and
+      result.getKind() = kind
+    }
+
+    override ReturnKind getKind() { result = kind }
+
+    override CfgScope getCfgScope() { result = scope }
+
+    override Location getLocationImpl() { result = scope.getLocation() }
+
+    override string toStringImpl() { result = "return " + kind + " in " + scope }
+  }
+
   private class SummaryReturnNode extends SummaryNode, ReturnNode {
     private ReturnKind rk;
 
@@ -631,7 +696,7 @@ private import PostUpdateNodes
 
 /** A node that performs a type cast. */
 class CastNode extends Node {
-  CastNode() { none() }
+  CastNode() { this instanceof ReturningNode }
 }
 
 class DataFlowExpr = CfgNodes::ExprCfgNode;

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

@@ -120,7 +120,7 @@ predicate hasLocalSource(Node sink, Node source) {
   or
   exists(Node mid |
     hasLocalSource(mid, source) and
-    simpleLocalFlowStep(mid, sink)
+    simpleLocalFlowStepCommon(mid, sink)
   )
 }
 
@@ -137,7 +137,7 @@ ParameterNode parameterNode(Parameter p) { result.getParameter() = p }
  * (intra-procedural) step.
  */
 predicate localFlowStep(Node nodeFrom, Node nodeTo) {
-  simpleLocalFlowStep(nodeFrom, nodeTo)
+  simpleLocalFlowStepCommon(nodeFrom, nodeTo)
   or
   // Simple flow through library code is included in the exposed local
   // step relation, even though flow is technically inter-procedural

ql/lib/codeql/ruby/typetracking/TypeTrackerSpecific.qll

@@ -45,11 +45,18 @@ predicate callStep(DataFlowPrivate::ArgumentNode nodeFrom, DataFlowPublic::Param
  * recursion (or, at best, terrible performance), since identifying calls to library
  * methods is done using API graphs (which uses type tracking).
  */
-predicate returnStep(DataFlowPrivate::ReturnNode nodeFrom, Node nodeTo) {
+predicate returnStep(Node nodeFrom, Node nodeTo) {
   exists(ExprNodes::CallCfgNode call |
+    nodeFrom instanceof DataFlowPrivate::ReturnNode and
     nodeFrom.(DataFlowPrivate::NodeImpl).getCfgScope() = DataFlowDispatch::getTarget(call) and
     nodeTo.asExpr().getNode() = call.getNode()
   )
+  or
+  // In normal data-flow, this will be a local flow step. But for type tracking
+  // we model it as a returning flow step, in order to avoid computing a potential
+  // self-cross product of all calls to a function that returns one of its parameters
+  // (only to later filter that flow out using `TypeTracker::append`).
+  nodeTo.(DataFlowPrivate::SynthReturnNode).getAnInput() = nodeFrom
 }
 
 /**

ql/test/library-tests/dataflow/local/ReturnNodes.ql

@@ -1,4 +1,4 @@
 import ruby
 import codeql.ruby.dataflow.internal.DataFlowPrivate
 
-select any(ReturnNode node)
+select any(ReturningNode node)

ql/test/query-tests/security/cwe-601/UrlRedirect.expected

@@ -3,9 +3,7 @@ edges
 | UrlRedirect.rb:14:17:14:22 | call to params :  | UrlRedirect.rb:14:17:14:43 | call to fetch |
 | UrlRedirect.rb:19:17:19:22 | call to params :  | UrlRedirect.rb:19:17:19:37 | call to to_unsafe_hash |
 | UrlRedirect.rb:24:31:24:36 | call to params :  | UrlRedirect.rb:24:17:24:37 | call to filter_params |
-| UrlRedirect.rb:24:31:24:36 | call to params :  | UrlRedirect.rb:56:21:56:32 | input_params :  |
 | UrlRedirect.rb:34:20:34:25 | call to params :  | UrlRedirect.rb:34:17:34:37 | "#{...}/foo" |
-| UrlRedirect.rb:56:21:56:32 | input_params :  | UrlRedirect.rb:57:5:57:29 | call to permit :  |
 nodes
 | UrlRedirect.rb:4:17:4:22 | call to params | semmle.label | call to params |
 | UrlRedirect.rb:9:17:9:22 | call to params :  | semmle.label | call to params :  |
@@ -18,10 +16,7 @@ nodes
 | UrlRedirect.rb:24:31:24:36 | call to params :  | semmle.label | call to params :  |
 | UrlRedirect.rb:34:17:34:37 | "#{...}/foo" | semmle.label | "#{...}/foo" |
 | UrlRedirect.rb:34:20:34:25 | call to params :  | semmle.label | call to params :  |
-| UrlRedirect.rb:56:21:56:32 | input_params :  | semmle.label | input_params :  |
-| UrlRedirect.rb:57:5:57:29 | call to permit :  | semmle.label | call to permit :  |
 subpaths
-| UrlRedirect.rb:24:31:24:36 | call to params :  | UrlRedirect.rb:56:21:56:32 | input_params :  | UrlRedirect.rb:57:5:57:29 | call to permit :  | UrlRedirect.rb:24:17:24:37 | call to filter_params :  |
 #select
 | UrlRedirect.rb:4:17:4:22 | call to params | UrlRedirect.rb:4:17:4:22 | call to params | UrlRedirect.rb:4:17:4:22 | call to params | Untrusted URL redirection due to $@. | UrlRedirect.rb:4:17:4:22 | call to params | a user-provided value |
 | UrlRedirect.rb:9:17:9:28 | ...[...] | UrlRedirect.rb:9:17:9:22 | call to params :  | UrlRedirect.rb:9:17:9:28 | ...[...] | Untrusted URL redirection due to $@. | UrlRedirect.rb:9:17:9:22 | call to params | a user-provided value |