Capture argument->return value flows

java/ql/src/utils/model-generator/CaptureSummaryModels.ql

@@ -6,8 +6,13 @@
 
 import java
 import ModelGeneratorUtils
+import semmle.code.java.dataflow.TaintTracking
+import semmle.code.java.dataflow.internal.DataFlowImplCommon
 
-string captureFlow(Callable api) { result = captureQualifierFlow(api) }
+string captureFlow(Callable api) {
+  result = captureQualifierFlow(api) or
+  result = captureParameterFlowToReturnValue(api)
+}
 
 string captureQualifierFlow(Callable api) {
   exists(ReturnStmt rtn |
@@ -17,6 +22,44 @@ string captureQualifierFlow(Callable api) {
   result = asValueModel(api, "Argument[-1]", "ReturnValue")
 }
 
+class ParameterToReturnValueTaintConfig extends TaintTracking::Configuration {
+  ParameterToReturnValueTaintConfig() { this = "ParameterToReturnValueTaintConfig" }
+
+  override predicate isSource(DataFlow::Node source) {
+    exists(Parameter p, Callable api |
+      p = source.asParameter() and
+      api = p.getCallable() and
+      (
+        not api.getReturnType() instanceof PrimitiveType and
+        not p.getType() instanceof PrimitiveType
+      ) and
+      (
+        not api.getReturnType() instanceof TypeClass and
+        not p.getType() instanceof TypeClass
+      )
+    )
+  }
+
+  override predicate isSink(DataFlow::Node sink) { sink instanceof ReturnNodeExt }
+}
+
+// TODO: rtn -> Node as ReturnNodeExt is also PostUpdateNode, might be able to merge with p2p flow
+predicate paramFlowToReturnValueExists(Parameter p) {
+  exists(ParameterToReturnValueTaintConfig config, ReturnStmt rtn |
+    rtn.getEnclosingCallable() = p.getCallable() and
+    config.hasFlow(DataFlow::parameterNode(p), DataFlow::exprNode(rtn.getResult()))
+  )
+}
+
+string captureParameterFlowToReturnValue(Callable api) {
+  exists(Parameter p |
+    p = api.getAParameter() and
+    paramFlowToReturnValueExists(p)
+  |
+    result = asTaintModel(api, parameterAccess(p), "ReturnValue")
+  )
+}
+
 // TODO: handle cases like Ticker
 // TODO: "com.google.common.base;Converter;true;convertAll;(Iterable);;Element of Argument[0];Element of ReturnValue;taint",
 // TODO: infer interface from multiple implementations? e.g. UriComponentsContributor

java/ql/src/utils/model-generator/ModelGeneratorUtils.qll

@@ -1,5 +1,6 @@
 import java
 import semmle.code.java.dataflow.ExternalFlow
+import semmle.code.java.dataflow.internal.ContainerFlow
 
 string isExtensible(RefType ref) { if ref.isFinal() then result = "false" else result = "true" }
 
@@ -26,3 +27,12 @@ string asSummaryModel(Callable api, string input, string output, string kind) {
       + output + ";" //
       + kind + ";" //
 }
+
+string parameterAccess(Parameter p) {
+  if p.getType() instanceof Array
+  then result = "ArrayElement of Argument[" + p.getPosition() + "]"
+  else
+    if p.getType() instanceof ContainerType
+    then result = "Element of Argument[" + p.getPosition() + "]"
+    else result = "Argument[" + p.getPosition() + "]"
+}

java/ql/test/utils/model-generator/CaptureSummaryModels.expected

@@ -2,3 +2,12 @@
 | p;FluentAPI;false;returnsThis;(String);;Argument[-1];ReturnValue;value; |
 | p;InnerClasses$CaptureMe;true;yesCm;(String);;Argument[0];ReturnValue;taint; |
 | p;InnerClasses;true;yes;(String);;Argument[0];ReturnValue;taint; |
+| p;ParamFlow;true;addTo;(String,List);;Argument[0];Element of Argument[1];taint; |
+| p;ParamFlow;true;returnArrayElement;(String[]);;ArrayElement of Argument[0];ReturnValue;taint; |
+| p;ParamFlow;true;returnCollectionElement;(List);;Element of Argument[0];ReturnValue;taint; |
+| p;ParamFlow;true;returnIterableElement;(Iterable);;Element of Argument[0];ReturnValue;taint; |
+| p;ParamFlow;true;returnIteratorElement;(Iterator);;Element of Argument[0];ReturnValue;taint; |
+| p;ParamFlow;true;returnMultipleParameters;(String,String);;Argument[0];ReturnValue;taint; |
+| p;ParamFlow;true;returnMultipleParameters;(String,String);;Argument[1];ReturnValue;taint; |
+| p;ParamFlow;true;returnVarArgElement;(String[]);;ArrayElement of Argument[0];ReturnValue;taint; |
+| p;ParamFlow;true;returnsInput;(String);;Argument[0];ReturnValue;taint; |

java/ql/test/utils/model-generator/p/ParamFlow.java

@@ -0,0 +1,59 @@
+package p;
+
+import java.util.Iterator;
+import java.util.List;
+import java.io.IOException;
+import java.io.OutputStream;
+
+
+public class ParamFlow {
+
+    public String returnsInput(String input) {
+        return input;
+    }
+
+    public int ignorePrimitiveReturnValue(String input) {
+        return input.length();
+    }
+
+    public String returnMultipleParameters(String one, String two) {
+        if (System.currentTimeMillis() > 100) {
+            return two;
+        }
+        return one;
+    }
+
+    public String returnArrayElement(String[] input) {
+        return input[0];
+    }
+
+    public String returnVarArgElement(String... input) {
+        return input[0];
+    }
+
+    public String returnCollectionElement(List<String> input) {
+        return input.get(0);
+    }
+
+    public String returnIteratorElement(Iterator<String> input) {
+        return input.next();
+    }
+
+    public String returnIterableElement(Iterable<String> input) {
+        return input.iterator().next();
+    }
+
+    public Class<?> mapType(Class<?> input) {
+        return input;
+    }
+
+    public void writeChunked(byte[] data, OutputStream output)
+            throws IOException {
+        output.write(data, 0, data.length);
+    }
+
+    public void addTo(String data, List<String> target) {
+        target.add(data);
+    }
+
+}