From 9500c9c8bc601682908f679751ff7569a3b3ebef Mon Sep 17 00:00:00 2001
From: Benjamin Muskalla <bmuskalla@github.com>
Date: Tue, 26 Oct 2021 15:57:27 +0200
Subject: [PATCH] Support lambda flow for source models

Also rely on public API to detect the source node
---
 .../dataflow/internal/FlowSummaryImpl.qll     | 15 +++-------
 .../model-generator/CaptureSourceModels.ql    | 30 ++++++++++++-------
 .../CaptureSourceModels.expected              |  4 ++-
 .../test/utils/model-generator/p/Sources.java | 11 +++++++
 4 files changed, 38 insertions(+), 22 deletions(-)

diff --git a/java/ql/lib/semmle/code/java/dataflow/internal/FlowSummaryImpl.qll b/java/ql/lib/semmle/code/java/dataflow/internal/FlowSummaryImpl.qll
index 9d7588ab87e..5955285bd6f 100644
--- a/java/ql/lib/semmle/code/java/dataflow/internal/FlowSummaryImpl.qll
+++ b/java/ql/lib/semmle/code/java/dataflow/internal/FlowSummaryImpl.qll
@@ -889,17 +889,10 @@ module Private {
      * model.
      */
     predicate isSourceNode(InterpretNode node, string kind) {
-      exists(InterpretNode ref, string output | isSourceNode(ref, node, output, kind))
-    }
-
-    // TODO: I wonder if this is actually the interface we want to expose.
-    predicate isSourceNode(InterpretNode node, string output, string kind) {
-      exists(InterpretNode ref | isSourceNode(ref, node, output, kind))
-    }
-
-    predicate isSourceNode(InterpretNode ref, InterpretNode node, string output, string kind) {
-      sourceElementRef(ref, output, kind) and
-      interpretOutput(output, 0, ref, node)
+      exists(InterpretNode ref, string output |
+        sourceElementRef(ref, output, kind) and
+        interpretOutput(output, 0, ref, node)
+      )
     }
 
     /**
diff --git a/java/ql/src/utils/model-generator/CaptureSourceModels.ql b/java/ql/src/utils/model-generator/CaptureSourceModels.ql
index 93017b2eaf8..8386c54a44b 100644
--- a/java/ql/src/utils/model-generator/CaptureSourceModels.ql
+++ b/java/ql/src/utils/model-generator/CaptureSourceModels.ql
@@ -13,6 +13,7 @@ private import ModelGeneratorUtils
 private import semmle.code.java.dataflow.internal.FlowSummaryImplSpecific
 private import semmle.code.java.dataflow.internal.FlowSummaryImpl
 private import semmle.code.java.dataflow.internal.DataFlowImplCommon
+import semmle.code.java.dataflow.internal.DataFlowNodes::Private
 
 class FromSourceConfiguration extends TaintTracking::Configuration {
   FromSourceConfiguration() { this = "FromSourceConfiguration" }
@@ -26,24 +27,33 @@ class FromSourceConfiguration extends TaintTracking::Configuration {
       c.isPublic() and
       c.fromSource()
     )
+    or
+    exists(MethodAccess c | sink.asExpr() = c.getAnArgument())
   }
 }
 
-// TODO: better way than rely on internals to capture kind?
-cached
-predicate specificSourceNode(DataFlow::Node node, string output, string kind) {
-  exists(InterpretNode n | Private::External::isSourceNode(n, output, kind) and n.asNode() = node)
+string asOutput(DataFlow::Node node) {
+  if node instanceof ReturnNodeExt
+  then result = "ReturnValue"
+  else
+    result =
+      "Parameter[" +
+        node.(ArgumentNode)
+            .getCall()
+            .asCall()
+            .getQualifier()
+            .(VarAccess)
+            .getVariable()
+            .(Parameter)
+            .getPosition() + "]"
 }
 
 string captureSource(Callable api) {
-  exists(
-    DataFlow::Node src, DataFlow::Node sink, FromSourceConfiguration config, string kind,
-    string output
-  |
+  exists(DataFlow::Node src, DataFlow::Node sink, FromSourceConfiguration config, string kind |
     config.hasFlow(src, sink) and
-    specificSourceNode(sink, output, kind) and
+    sourceNode(sink, kind) and
     api = src.getEnclosingCallable() and
-    result = asSourceModel(api, output, kind)
+    result = asSourceModel(api, asOutput(sink), kind)
   )
 }
 
diff --git a/java/ql/test/utils/model-generator/CaptureSourceModels.expected b/java/ql/test/utils/model-generator/CaptureSourceModels.expected
index 533038c1521..5ae892061c4 100644
--- a/java/ql/test/utils/model-generator/CaptureSourceModels.expected
+++ b/java/ql/test/utils/model-generator/CaptureSourceModels.expected
@@ -1 +1,3 @@
-| p;Sources;true;readUrl;(URL);;ReturnValue;remote; |
\ No newline at end of file
+| p;Sources;true;consumeSource;(int,Consumer);;Parameter[1];remote; |
+| p;Sources;true;readUrl;(URL);;ReturnValue;remote; |
+| p;Sources;true;socketStream;();;ReturnValue;remote; |
\ No newline at end of file
diff --git a/java/ql/test/utils/model-generator/p/Sources.java b/java/ql/test/utils/model-generator/p/Sources.java
index fe6145a4efc..c41f4a499f6 100644
--- a/java/ql/test/utils/model-generator/p/Sources.java
+++ b/java/ql/test/utils/model-generator/p/Sources.java
@@ -2,7 +2,9 @@ package p;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.net.ServerSocket;
 import java.net.URL;
+import java.util.function.Consumer;
 
 
 public class Sources {
@@ -11,4 +13,13 @@ public class Sources {
         return url.openConnection().getInputStream();
     }
 
+    public InputStream socketStream() throws IOException {
+        ServerSocket socket = new ServerSocket(123);
+        return socket.accept().getInputStream();
+    }
+
+    public void consumeSource(int port, Consumer<InputStream> consumer) throws IOException {
+        ServerSocket socket = new ServerSocket(port);
+        consumer.accept(socket.accept().getInputStream());
+    }
 }