Data flow: Synthesize parameter return nodes

java/ql/src/utils/modelgenerator/internal/CaptureModels.qll

@@ -6,6 +6,29 @@
 private import CaptureModelsSpecific
 private import CaptureModelsPrinting
 
+/**
+ * A node from which flow can return to the caller. This is either a regular
+ * `ReturnNode` or a `PostUpdateNode` corresponding to the value of a parameter.
+ */
+private class ReturnNodeExt extends DataFlow::Node {
+  private DataFlowImplCommon::ReturnKindExt kind;
+
+  ReturnNodeExt() {
+    kind = DataFlowImplCommon::getValueReturnPosition(this).getKind() or
+    kind = DataFlowImplCommon::getParamReturnPosition(this, _).getKind()
+  }
+
+  string getOutput() {
+    kind instanceof DataFlowImplCommon::ValueReturnKind and
+    result = "ReturnValue"
+    or
+    exists(ParameterPosition pos |
+      pos = kind.(DataFlowImplCommon::ParamUpdateReturnKind).getPosition() and
+      result = paramReturnNodeAsOutput(returnNodeEnclosingCallable(this), pos)
+    )
+  }
+}
+
 class DataFlowTargetApi extends TargetApiSpecific {
   DataFlowTargetApi() { not isUninterestingForDataFlowModels(this) }
 }
@@ -65,7 +88,7 @@ string asInputArgument(DataFlow::Node source) { result = asInputArgumentSpecific
  * Gets the summary model of `api`, if it follows the `fluent` programming pattern (returns `this`).
  */
 string captureQualifierFlow(TargetApiSpecific api) {
-  exists(DataFlowImplCommon::ReturnNodeExt ret |
+  exists(ReturnNodeExt ret |
     api = returnNodeEnclosingCallable(ret) and
     isOwnInstanceAccessNode(ret)
   ) and
@@ -130,7 +153,7 @@ module ThroughFlowConfig implements DataFlow::StateConfigSig {
   }
 
   predicate isSink(DataFlow::Node sink, FlowState state) {
-    sink instanceof DataFlowImplCommon::ReturnNodeExt and
+    sink instanceof ReturnNodeExt and
     not isOwnInstanceAccessNode(sink) and
     not exists(captureQualifierFlow(sink.asExpr().getEnclosingCallable())) and
     (state instanceof TaintRead or state instanceof TaintStore)
@@ -171,14 +194,11 @@ private module ThroughFlow = TaintTracking::GlobalWithState<ThroughFlowConfig>;
  * Gets the summary model(s) of `api`, if there is flow from parameters to return value or parameter.
  */
 string captureThroughFlow(DataFlowTargetApi api) {
-  exists(
-    DataFlow::ParameterNode p, DataFlowImplCommon::ReturnNodeExt returnNodeExt, string input,
-    string output
-  |
+  exists(DataFlow::ParameterNode p, ReturnNodeExt returnNodeExt, string input, string output |
     ThroughFlow::flow(p, returnNodeExt) and
     returnNodeExt.(DataFlow::Node).getEnclosingCallable() = api and
     input = parameterNodeAsInput(p) and
-    output = returnNodeAsOutput(returnNodeExt) and
+    output = returnNodeExt.getOutput() and
     input != output and
     result = ModelPrinting::asTaintModel(api, input, output)
   )
@@ -196,7 +216,7 @@ module FromSourceConfig implements DataFlow::ConfigSig {
 
   predicate isSink(DataFlow::Node sink) {
     exists(DataFlowTargetApi c |
-      sink instanceof DataFlowImplCommon::ReturnNodeExt and
+      sink instanceof ReturnNodeExt and
       sink.getEnclosingCallable() = c
     )
   }
@@ -214,12 +234,12 @@ private module FromSource = TaintTracking::Global<FromSourceConfig>;
  * Gets the source model(s) of `api`, if there is flow from an existing known source to the return of `api`.
  */
 string captureSource(DataFlowTargetApi api) {
-  exists(DataFlow::Node source, DataFlow::Node sink, string kind |
+  exists(DataFlow::Node source, ReturnNodeExt sink, string kind |
     FromSource::flow(source, sink) and
     ExternalFlow::sourceNode(source, kind) and
     api = sink.getEnclosingCallable() and
     isRelevantSourceKind(kind) and
-    result = ModelPrinting::asSourceModel(api, returnNodeAsOutput(sink), kind)
+    result = ModelPrinting::asSourceModel(api, sink.getOutput(), kind)
   )
 }
 

java/ql/src/utils/modelgenerator/internal/CaptureModelsSpecific.qll

@@ -13,6 +13,7 @@ private import semmle.code.java.dataflow.TaintTracking as Tt
 import semmle.code.java.dataflow.ExternalFlow as ExternalFlow
 import semmle.code.java.dataflow.internal.DataFlowImplCommon as DataFlowImplCommon
 import semmle.code.java.dataflow.internal.DataFlowPrivate as DataFlowPrivate
+import semmle.code.java.dataflow.internal.DataFlowDispatch as DataFlowDispatch
 
 module DataFlow = Df::DataFlow;
 
@@ -202,26 +203,23 @@ string parameterAccess(J::Parameter p) {
 
 class InstanceParameterNode = DataFlow::InstanceParameterNode;
 
+class ParameterPosition = DataFlowDispatch::ParameterPosition;
+
 /**
- * Gets the MaD string represention of the the return node `node`.
+ * Gets the MaD string represention of return through parameter at position
+ * `pos` of callable `c`.
  */
-string returnNodeAsOutput(DataFlowImplCommon::ReturnNodeExt node) {
-  if node.getKind() instanceof DataFlowImplCommon::ValueReturnKind
-  then result = "ReturnValue"
-  else
-    exists(int pos |
-      pos = node.getKind().(DataFlowImplCommon::ParamUpdateReturnKind).getPosition()
-    |
-      result = parameterAccess(node.(DataFlow::Node).getEnclosingCallable().getParameter(pos))
-      or
-      result = qualifierString() and pos = -1
-    )
+bindingset[c]
+string paramReturnNodeAsOutput(Callable c, ParameterPosition pos) {
+  result = parameterAccess(c.getParameter(pos))
+  or
+  result = qualifierString() and pos = -1
 }
 
 /**
  * Gets the enclosing callable of `ret`.
  */
-Callable returnNodeEnclosingCallable(DataFlowImplCommon::ReturnNodeExt ret) {
+Callable returnNodeEnclosingCallable(DataFlow::Node ret) {
   result = DataFlowImplCommon::getNodeEnclosingCallable(ret).asCallable()
 }