Merge pull request #22854 from asgerf/unified/shorthand-returns

Unified: Support shorthand returns

Author: asgerf

unified/ql/lib/codeql/unified/internal/ExprPositions.qll

@@ -79,3 +79,13 @@ predicate hasResultValue(Expr e) {
   not isInAssignmentContext(e, any(AssignExpr n)) and // non-compound assignment target
   not e instanceof IdentifierLabel
 }
+
+/** Holds if `e` is returned from `callable` either because it is an operand of `return` or is implicitly returned due to language-specific semantics. */
+predicate isReturnedExpr(Callable callable, Expr e) {
+  exists(ReturnExpr ret |
+    callable = ret.getEnclosingCallable() and
+    e = ret.getValue()
+  )
+  or
+  any(ExprPositionsPlugin p).isReturnedExpr(callable, e)
+}

unified/ql/lib/codeql/unified/internal/ExprPositionsPlugin.qll

@@ -7,4 +7,6 @@ private module Plugins {
 
 class ExprPositionsPlugin extends Unit {
   predicate isInTypeContext(Expr e) { none() }
+
+  predicate isReturnedExpr(Callable callable, Expr e) { none() }
 }

unified/ql/lib/codeql/unified/internal/ExprPositionsPluginSwift.qll

@@ -3,4 +3,24 @@ private import ExprPositionsPlugin
 
 private class ExprPositionsPluginSwift extends ExprPositionsPlugin {
   override predicate isInTypeContext(Expr e) { e = any(GenericTypeExpr g).getBase() }
+
+  override predicate isReturnedExpr(Callable callable, Expr e) {
+    // Short-hand return from a function body, e.g. `func foo() -> String { "42" }`
+    // This happens when the body contains a single expression and this expression resolves to a value.
+    // We don't detect precisely whether the expression resolves to a value, but we restrict to cases
+    // where a spurious return would be harmless in practice.
+    (
+      callable instanceof FunctionDeclaration and
+      exists(Expr ret | ret = callable.getReturnType() | not ret.(Identifier).getValue() = "Void")
+      or
+      callable instanceof GetterDeclaration
+      or
+      callable instanceof FunctionExpr // We can't tell whether closures have a return value
+    ) and
+    exists(Block body |
+      body = callable.getBody() and
+      not exists(body.getStmt(1)) and // Check that only one statement exists
+      e = body.getStmt(0)
+    )
+  }
 }

unified/ql/lib/codeql/unified/internal/FacadeAst.qll

@@ -437,4 +437,40 @@ module Unified {
     /** Gets the number of elements in this array literal. */
     int getNumberOfElements() { result = count(this.getAnElement()) }
   }
+
+  /**
+   * A declaration of a getter associated with a variable or field.
+   *
+   * For example, in Swift:
+   * ```swift
+   * var x: String {
+   *   get { "42" } // getter declaration
+   * }
+   * var x: String { "42" } // getter declaration
+   * ```
+   */
+  final class GetterDeclaration extends AccessorDeclaration {
+    GetterDeclaration() { this.getAccessorKind().getValue() = "get" }
+  }
+
+  /**
+   * A declaration of a setter associated with a variable or field.
+   *
+   * For example, in Swift:
+   * ```swift
+   * var x: String {
+   *   get { "42" }
+   *   set { print(newValue) } // setter declaration
+   * }
+   * ```
+   */
+  final class SetterDeclaration extends AccessorDeclaration {
+    SetterDeclaration() { this.getAccessorKind().getValue() = "set" }
+  }
+
+  /** An `if` statement or conditional expression. */
+  class IfExpr extends G::IfExpr {
+    /** Gets then `then` or `else` branch. */
+    Expr getABranch() { result = [this.getThen(), this.getElse()] }
+  }
 }

unified/ql/lib/codeql/unified/internal/dataflow/DataFlowGraph.qll

@@ -1,6 +1,7 @@
 private import unified
 private import AllDataFlow
 private import codeql.unified.internal.LocalNameBinding
+private import codeql.unified.internal.ExprPositions
 
 predicate step(BuilderNode node1, Step step, BuilderNode node2) {
   any(DataFlowPlugin p).step(node1, step, node2)
@@ -180,6 +181,19 @@ predicate step(BuilderNode node1, Step step, BuilderNode node2) {
     node2.isIncomingValue(fun.getNameNode())
   )
   or
+  exists(Callable callable, Expr expr |
+    isReturnedExpr(callable, expr) and
+    node1.isResultValue(expr) and
+    step.value() and
+    node2.isReturnNode(callable)
+  )
+  or
+  exists(IfExpr expr |
+    node1.isResultValue(expr.getABranch()) and
+    step.value() and
+    node2.isResultValue(expr)
+  )
+  or
   none() // Temporarily disable compilation errors from unsatisfiable types
 }
 

unified/ql/lib/codeql/unified/internal/dataflow/DataFlowInstantiation.qll

@@ -23,8 +23,7 @@ module DataFlowInput implements InputSig<Location> {
   // Parameter, argument, return, and out nodes and their positions/kinds
   //
   private predicate isReturnNode(Node n, ReturnKind kind) {
-    // TODO: Handle short-hand returns such as `func foo() -> Int { 5 }`
-    n.asExpr() = any(ReturnExpr r).getValue() and kind.isStandardReturnValue()
+    n.isReturnNode(_) and kind.isStandardReturnValue()
     or
     FlowSummaryOutput::summaryReturnNode(getSummaryNode(n), kind)
   }
@@ -240,7 +239,8 @@ module DataFlowInput implements InputSig<Location> {
 
   predicate neverSkipInPathGraph(Node n) {
     n.isIncomingValue(_) or // Never skip assignment target
-    n.asExpr() instanceof LocalVariableAccess // Never skip a variable reference
+    n.asExpr() instanceof LocalVariableAccess or // Never skip a variable reference
+    n.asExpr() = any(ReturnExpr e).getValue()
   }
 
   class DataFlowExpr = Expr;

unified/ql/lib/codeql/unified/internal/dataflow/DataFlowNode.qll

@@ -91,6 +91,9 @@ newtype TDataFlowNode =
   TLocalVariableRefNode(AstNode repr, LocalVariable var, VariableRefKind kind) {
     performsVariableAccess(repr, var, kind, _)
   } or
+  TReturnNode(DataFlowCallable callable) {
+    not exists(callable.asSummarizedCallable()) // represented by TSummaryNode
+  } or
   TCallableNode(DataFlowCallable callable) or
   TImplicitParameterNode(DataFlowCallable callable, ImplicitParameterPosition pos) {
     not exists(callable.asSummarizedCallable()) // represented by TSummaryNode
@@ -104,7 +107,7 @@ newtype TDataFlowNode =
 
 class TDataFlowNodeStage1 =
   TValueNode or TStrictlyIncomingValue or TExprPostUpdateNode or TLocalVariableRefNode or
-      TCallableNode or TImplicitParameterNode or TImplicitArgumentNode;
+      TReturnNode or TCallableNode or TImplicitParameterNode or TImplicitArgumentNode;
 
 class TDataFlowNodeStage2 = TDataFlowNodeStage1 or TLocalSsaNode;
 
@@ -152,6 +155,16 @@ class BuilderNode extends TDataFlowNodeStage1 {
   /** Holds if this represents the updated state of the value returned by `expr` after it has been mutated by the surrounding assignment or call. */
   predicate isPostUpdate(Expr expr) { this = TExprPostUpdateNode(expr) }
 
+  /** Holds if this node is the canonical representative for the return value of `callable`. */
+  predicate isReturnNode(Callable callable) { this.isReturnNodeEx(toDataFlowCallable(callable)) }
+
+  /**
+   * Holds if this node is the canonical representative for the return value of `callable`.
+   *
+   * Note that this does not hold for the return node in a summarized callable.
+   */
+  predicate isReturnNodeEx(DataFlowCallable callable) { this = TReturnNode(callable) }
+
   /** Holds if this node represents an implicit parameter of `callable`. */
   predicate isImplicitParameter(DataFlowCallable callable, ImplicitParameterPosition pos) {
     this = TImplicitParameterNode(callable, pos)
@@ -258,6 +271,11 @@ class BuilderNode extends TDataFlowNodeStage1 {
       result = "[callable] " + callable.toString()
     )
     or
+    exists(DataFlowCallable callable |
+      this.isReturnNodeEx(callable) and
+      result = "[return] " + callable.toString()
+    )
+    or
     exists(DataFlowCall call, ArgumentPosition pos |
       this.isImplicitArgument(call, pos, false) and
       result = "[" + pos + " arg] " + call.toString()
@@ -284,6 +302,11 @@ class BuilderNode extends TDataFlowNodeStage1 {
       result = callable.getLocation()
     )
     or
+    exists(DataFlowCallable callable |
+      this.isReturnNodeEx(callable) and
+      result = callable.getLocation()
+    )
+    or
     exists(DataFlowCall call |
       this.isImplicitArgument(call, _, _) and
       result = call.getLocation()
@@ -371,6 +394,9 @@ class Node extends TDataFlowNode {
   /** Holds if this is the canonical representative for the given `callable`. */
   predicate isCallable(Callable callable) { this.(BuilderNode).isCallable(callable) }
 
+  /** Holds if this node is the canonical representative for the return value of `callable`. */
+  predicate isReturnNode(Callable callable) { this.(BuilderNode).isReturnNode(callable) }
+
   /**
    * Gets the AST node wrapped by this data flow, if any.
    */
@@ -419,8 +445,6 @@ class Node extends TDataFlowNode {
       result.asSourceCallable() = node.getSourceVariable().getDeclaringCallable()
     )
     or
-    this.(BuilderNode).isImplicitParameter(result, _)
-    or
     exists(CaptureSsaOutput::SynthesizedCaptureNode node |
       this = TCaptureSsaNode(node) and
       result.asSourceCallable() = node.getEnclosingCallable()
@@ -433,6 +457,8 @@ class Node extends TDataFlowNode {
       result = call.getEnclosingCallable()
     )
     or
+    this.(BuilderNode).isReturnNodeEx(result)
+    or
     exists(DataFlowCallable callable |
       this.isCallableEx(callable) and
       result.asSourceCallable() = callable.asSourceCallable().getEnclosingCallable()

unified/ql/lib/codeql/unified/internal/dataflow/DataFlowPluginSwift.qll

@@ -76,5 +76,13 @@ private class SwiftDataFlowPlugin extends DataFlowPlugin {
       step.taint() and
       node2.isResultValue(call)
     )
+    or
+    // Singleton blocks propagate their value upwards, e.g. in `if expr { expr } else { expr }`
+    exists(Block block |
+      not exists(block.getStmt(1)) and
+      node1.isResultValue(block.getStmt(0)) and
+      step.value() and
+      node2.isResultValue(block)
+    )
   }
 }

unified/ql/lib/codeql/unified/internal/mad/LegacyMaD.qll

@@ -245,7 +245,7 @@ private DataFlow::Node getSourceFromCall(CallExpr call, AccessPathToken token) {
 bindingset[callable, token]
 private DataFlow::Node getSinkFromCallable(Callable callable, AccessPathToken token) {
   token = "ReturnValue" and
-  result.asExpr() = any(ReturnExpr r | r.getEnclosingCallable() = callable).getValue() // TODO: expose canonical return node in data flow
+  result.isReturnNode(callable)
 }
 
 /**

unified/ql/test/library-tests/dataflow/short-return.swift

@@ -0,0 +1,9 @@
+func source(_ s: String) -> String { return s }
+
+@discardableResult
+func sink(_ s: Any) -> String { return "" }
+
+func t1() {
+    func getValue() -> String { source("t1.1") }
+    sink(getValue())  // $ hasValueFlow=t1.1
+}