Unified: Support shorthand returns

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
+}

unified/ql/test/library-tests/dataflow/test.expected

@@ -290,6 +290,7 @@ edges
 | implicit-self.swift:85:9:85:14 | [incoming] ... .x | implicit-self.swift:85:9:85:12 | [post] self [x] | provenance |  |
 | implicit-self.swift:85:29:85:43 | source(...) | implicit-self.swift:85:9:85:14 | [incoming] ... .x | provenance |  |
 | implicit-self.swift:87:14:87:17 | self [x] | implicit-self.swift:87:14:87:19 | ... .x | provenance |  |
+| short-return.swift:7:33:7:46 | source(...) | short-return.swift:8:10:8:19 | getValue(...) | provenance |  |
 | test.swift:17:10:17:23 | source(...) | test.swift:17:10:17:32 | ... + ... | provenance |  |
 | test.swift:18:19:18:32 | source(...) | test.swift:18:10:18:32 | ... + ... | provenance |  |
 | test.swift:20:13:20:26 | source(...) | test.swift:20:10:20:33 | StringInterpolationExpr | provenance |  |
@@ -452,6 +453,11 @@ edges
 | test.swift:264:6:264:16 | [post] TypeCastExpr [field] | test.swift:265:10:265:12 | box [field] | provenance |  |
 | test.swift:264:27:264:41 | source(...) | test.swift:264:5:264:23 | ... .field | provenance |  |
 | test.swift:265:10:265:12 | box [field] | test.swift:265:10:265:18 | ... .field | provenance |  |
+| test.swift:269:14:269:29 | source(...) | test.swift:269:10:269:39 | IfExpr | provenance |  |
+| test.swift:270:24:270:39 | source(...) | test.swift:270:10:270:39 | IfExpr | provenance |  |
+| test.swift:271:9:271:9 | x | test.swift:272:10:272:10 | x | provenance |  |
+| test.swift:271:20:271:35 | source(...) | test.swift:271:9:271:9 | x | provenance |  |
+| test.swift:273:30:273:45 | source(...) | test.swift:274:10:274:12 | f(...) | provenance |  |
 nodes
 | calls.swift:7:19:7:19 | x | semmle.label | x |
 | calls.swift:8:14:8:14 | x | semmle.label | x |
@@ -799,6 +805,8 @@ nodes
 | implicit-self.swift:86:14:86:14 | x | semmle.label | x |
 | implicit-self.swift:87:14:87:17 | self [x] | semmle.label | self [x] |
 | implicit-self.swift:87:14:87:19 | ... .x | semmle.label | ... .x |
+| short-return.swift:7:33:7:46 | source(...) | semmle.label | source(...) |
+| short-return.swift:8:10:8:19 | getValue(...) | semmle.label | getValue(...) |
 | test.swift:9:10:9:21 | source(...) | semmle.label | source(...) |
 | test.swift:17:10:17:23 | source(...) | semmle.label | source(...) |
 | test.swift:17:10:17:32 | ... + ... | semmle.label | ... + ... |
@@ -1000,6 +1008,15 @@ nodes
 | test.swift:264:27:264:41 | source(...) | semmle.label | source(...) |
 | test.swift:265:10:265:12 | box [field] | semmle.label | box [field] |
 | test.swift:265:10:265:18 | ... .field | semmle.label | ... .field |
+| test.swift:269:10:269:39 | IfExpr | semmle.label | IfExpr |
+| test.swift:269:14:269:29 | source(...) | semmle.label | source(...) |
+| test.swift:270:10:270:39 | IfExpr | semmle.label | IfExpr |
+| test.swift:270:24:270:39 | source(...) | semmle.label | source(...) |
+| test.swift:271:9:271:9 | x | semmle.label | x |
+| test.swift:271:20:271:35 | source(...) | semmle.label | source(...) |
+| test.swift:272:10:272:10 | x | semmle.label | x |
+| test.swift:273:30:273:45 | source(...) | semmle.label | source(...) |
+| test.swift:274:10:274:12 | f(...) | semmle.label | f(...) |
 subpaths
 | calls.swift:31:17:31:30 | source(...) | calls.swift:28:19:28:19 | x | calls.swift:29:16:29:24 | ... + ... | calls.swift:31:10:31:31 | target(...) |
 | calls.swift:32:17:32:30 | source(...) | calls.swift:28:19:28:19 | x | calls.swift:29:16:29:24 | ... + ... | calls.swift:32:10:32:31 | target(...) |
@@ -1103,6 +1120,7 @@ testFailures
 | implicit-self.swift:80:14:80:19 | ... .x | implicit-self.swift:78:24:78:38 | source(...) | implicit-self.swift:80:14:80:19 | ... .x | $@ | implicit-self.swift:78:24:78:38 | source(...) | source(...) |
 | implicit-self.swift:86:14:86:14 | x | implicit-self.swift:85:29:85:43 | source(...) | implicit-self.swift:86:14:86:14 | x | $@ | implicit-self.swift:85:29:85:43 | source(...) | source(...) |
 | implicit-self.swift:87:14:87:19 | ... .x | implicit-self.swift:85:29:85:43 | source(...) | implicit-self.swift:87:14:87:19 | ... .x | $@ | implicit-self.swift:85:29:85:43 | source(...) | source(...) |
+| short-return.swift:8:10:8:19 | getValue(...) | short-return.swift:7:33:7:46 | source(...) | short-return.swift:8:10:8:19 | getValue(...) | $@ | short-return.swift:7:33:7:46 | source(...) | source(...) |
 | test.swift:9:10:9:21 | source(...) | test.swift:9:10:9:21 | source(...) | test.swift:9:10:9:21 | source(...) | $@ | test.swift:9:10:9:21 | source(...) | source(...) |
 | test.swift:17:10:17:32 | ... + ... | test.swift:17:10:17:23 | source(...) | test.swift:17:10:17:32 | ... + ... | $@ | test.swift:17:10:17:23 | source(...) | source(...) |
 | test.swift:18:10:18:32 | ... + ... | test.swift:18:19:18:32 | source(...) | test.swift:18:10:18:32 | ... + ... | $@ | test.swift:18:19:18:32 | source(...) | source(...) |
@@ -1152,3 +1170,7 @@ testFailures
 | test.swift:255:10:255:35 | ... ?? ... | test.swift:253:30:253:44 | source(...) | test.swift:255:10:255:35 | ... ?? ... | $@ | test.swift:253:30:253:44 | source(...) | source(...) |
 | test.swift:256:10:256:32 | ... ?? ... | test.swift:253:30:253:44 | source(...) | test.swift:256:10:256:32 | ... ?? ... | $@ | test.swift:253:30:253:44 | source(...) | source(...) |
 | test.swift:265:10:265:18 | ... .field | test.swift:264:27:264:41 | source(...) | test.swift:265:10:265:18 | ... .field | $@ | test.swift:264:27:264:41 | source(...) | source(...) |
+| test.swift:269:10:269:39 | IfExpr | test.swift:269:14:269:29 | source(...) | test.swift:269:10:269:39 | IfExpr | $@ | test.swift:269:14:269:29 | source(...) | source(...) |
+| test.swift:270:10:270:39 | IfExpr | test.swift:270:24:270:39 | source(...) | test.swift:270:10:270:39 | IfExpr | $@ | test.swift:270:24:270:39 | source(...) | source(...) |
+| test.swift:272:10:272:10 | x | test.swift:271:20:271:35 | source(...) | test.swift:272:10:272:10 | x | $@ | test.swift:271:20:271:35 | source(...) | source(...) |
+| test.swift:274:10:274:12 | f(...) | test.swift:273:30:273:45 | source(...) | test.swift:274:10:274:12 | f(...) | $@ | test.swift:273:30:273:45 | source(...) | source(...) |

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

@@ -264,3 +264,12 @@ func t29() {
     (box as! Box).field = source("t29.1")
     sink(box.field)  // $ hasValueFlow=t29.1
 }
+
+func conditional(_ b: Bool) {
+    sink(b ? source("cond.1") : "clean")  // $ hasValueFlow=cond.1
+    sink(b ? "clean" : source("cond.2"))  // $ hasValueFlow=cond.2
+    let x = if b { source("cond.3") } else { "clean" }
+    sink(x)  // $ hasValueFlow=cond.3
+    func f() -> String { b ? source("cond.4") : "clean" }
+    sink(f())  // $ hasValueFlow=cond.4
+}

unified/ql/test/library-tests/type-inference/pattern_matching.swift

@@ -31,3 +31,7 @@ func createValue() -> PatternMatch<Bool> {
   var y = PatternMatch<String>.novalue  // $ type=y@PatternMatch<T>:String field=PatternMatch.novalue
   return PatternMatch.novalue  // $ type=.novalue@PatternMatch<T>:Bool field=PatternMatch.novalue
 }
+
+func createValueShorthand() -> PatternMatch<Bool> {
+  PatternMatch.novalue  // $ type=.novalue@PatternMatch<T>:Bool field=PatternMatch.novalue
+}

unified/ql/test/library-tests/type-inference/type-inference.expected

@@ -2595,6 +2595,10 @@ inferType
 | pattern_matching.swift:31:11:31:38 | ... .novalue | T | {EXTERNAL LOCATION} | String |
 | pattern_matching.swift:32:10:32:29 | ... .novalue |  | pattern_matching.swift:7:1:11:1 | PatternMatch |
 | pattern_matching.swift:32:10:32:29 | ... .novalue | T | {EXTERNAL LOCATION} | Bool |
+| pattern_matching.swift:35:51:37:1 | Block |  | pattern_matching.swift:7:1:11:1 | PatternMatch |
+| pattern_matching.swift:35:51:37:1 | Block | T | {EXTERNAL LOCATION} | Bool |
+| pattern_matching.swift:36:3:36:22 | ... .novalue |  | pattern_matching.swift:7:1:11:1 | PatternMatch |
+| pattern_matching.swift:36:3:36:22 | ... .novalue | T | {EXTERNAL LOCATION} | Bool |
 | protocols.swift:8:7:8:12 | radius |  | {EXTERNAL LOCATION} | Double |
 | protocols.swift:10:8:10:13 | radius |  | {EXTERNAL LOCATION} | Double |
 | protocols.swift:10:24:12:3 | Block |  | {EXTERNAL LOCATION} | Double |