Swift: flow out of calls via return statements

Author: Robert Marsh

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

@@ -155,7 +155,19 @@ abstract class ReturnNode extends Node {
   abstract ReturnKind getKind();
 }
 
-private module ReturnNodes { }
+private module ReturnNodes { 
+  class ReturnReturnNode extends ReturnNode, ExprNode {
+    ReturnReturnNode() {
+      exists(ReturnStmt stmt |
+        stmt.getResult() = this.asExpr()
+      )
+    }
+
+    override ReturnKind getKind() {
+      result instanceof NormalReturnKind
+    }
+  }
+}
 
 import ReturnNodes
 
@@ -165,7 +177,13 @@ abstract class OutNode extends Node {
   abstract DataFlowCall getCall(ReturnKind kind);
 }
 
-private module OutNodes { }
+private module OutNodes { 
+  class CallOutNode extends OutNode, DataFlowCall {
+    override DataFlowCall getCall(ReturnKind kind) {
+      result = this and kind instanceof NormalReturnKind
+    }
+  }
+}
 
 import OutNodes
 

swift/ql/test/library-tests/dataflow/dataflow/DataFlow.expected

@@ -6,6 +6,7 @@ edges
 | test.swift:25:20:25:27 | CallExpr :  | test.swift:28:18:28:21 | WriteDef :  |
 | test.swift:28:18:28:21 | ParamDecl :  | test.swift:29:15:29:15 | DeclRefExpr |
 | test.swift:28:18:28:21 | WriteDef :  | test.swift:29:15:29:15 | DeclRefExpr |
+| test.swift:33:12:33:19 | CallExpr :  | test.swift:37:15:37:29 | CallExpr |
 nodes
 | test.swift:6:19:6:26 | CallExpr :  | semmle.label | CallExpr :  |
 | test.swift:7:15:7:15 | DeclRefExpr | semmle.label | DeclRefExpr |
@@ -17,9 +18,12 @@ nodes
 | test.swift:28:18:28:21 | WriteDef :  | semmle.label | ParamDecl :  |
 | test.swift:28:18:28:21 | WriteDef :  | semmle.label | WriteDef :  |
 | test.swift:29:15:29:15 | DeclRefExpr | semmle.label | DeclRefExpr |
+| test.swift:33:12:33:19 | CallExpr :  | semmle.label | CallExpr :  |
+| test.swift:37:15:37:29 | CallExpr | semmle.label | CallExpr |
 subpaths
 #select
 | test.swift:6:19:6:26 | CallExpr :  | test.swift:7:15:7:15 | DeclRefExpr |
 | test.swift:6:19:6:26 | CallExpr :  | test.swift:9:15:9:15 | DeclRefExpr |
 | test.swift:6:19:6:26 | CallExpr :  | test.swift:10:15:10:15 | DeclRefExpr |
 | test.swift:25:20:25:27 | CallExpr :  | test.swift:29:15:29:15 | DeclRefExpr |
+| test.swift:33:12:33:19 | CallExpr :  | test.swift:37:15:37:29 | CallExpr |