Refactors Ratpack lambda taint tracking to use generic API

Author: JLLeitschuh

java/ql/src/semmle/code/java/frameworks/ratpack/RatpackExec.qll

@@ -9,68 +9,74 @@ class RatpackPromise extends RefType {
   }
 }
 
-/**
- * Taint flows from the qualifier to the first argument of the lambda passed to this method access.
- * Eg. `tainted.map(stillTainted -> ..)`
- */
-abstract private class TaintFromQualifierToFunctionalArgumentMethodAccess extends MethodAccess { }
+abstract private class SimpleFluentLambdaMethod extends Method {
+  SimpleFluentLambdaMethod() { getNumberOfParameters() = 1 }
 
-class RatpackPromiseMapMethod extends Method {
+  /**
+   * Holds if this lambda consumes taint from the quaifier when `arg` is tainted.
+   * Eg. `tainted.map(stillTainted -> ..)`
+   */
+  abstract predicate consumesTaint(int arg);
+
+  /**
+   * Holds if the lambda passed produces taint that taints the result of this method.
+   * Eg. `var tainted = CompletableFuture.supplyAsync(() -> taint());`
+   */
+  predicate doesReturnTaint() { none() }
+}
+
+class RatpackPromiseMapMethod extends SimpleFluentLambdaMethod {
   RatpackPromiseMapMethod() {
     getDeclaringType() instanceof RatpackPromise and
     hasName("map")
   }
-}
 
-class RatpackPromiseMapMethodAccess extends TaintFromQualifierToFunctionalArgumentMethodAccess {
-  RatpackPromiseMapMethodAccess() { getMethod() instanceof RatpackPromiseMapMethod }
+  override predicate consumesTaint(int arg) { arg = 0 }
+
+  override predicate doesReturnTaint() { any() }
 }
 
-class RatpackPromiseThenMethod extends Method {
+class RatpackPromiseThenMethod extends SimpleFluentLambdaMethod {
   RatpackPromiseThenMethod() {
     getDeclaringType() instanceof RatpackPromise and
     hasName("then")
   }
-}
 
-class RatpackPromiseThenMethodAccess extends TaintFromQualifierToFunctionalArgumentMethodAccess {
-  RatpackPromiseThenMethodAccess() { getMethod() instanceof RatpackPromiseThenMethod }
+  override predicate consumesTaint(int arg) { arg = 0 }
 }
 
-class RatpackPromiseNextMethod extends FluentMethod {
+class RatpackPromiseNextMethod extends FluentMethod, SimpleFluentLambdaMethod {
   RatpackPromiseNextMethod() {
     getDeclaringType() instanceof RatpackPromise and
     hasName("next")
   }
-}
 
-class RatpackPromiseNextMethodAccess extends TaintFromQualifierToFunctionalArgumentMethodAccess {
-  RatpackPromiseNextMethodAccess() { getMethod() instanceof RatpackPromiseNextMethod }
+  override predicate consumesTaint(int arg) { arg = 0 }
 }
 
 private class RatpackPromiseTaintPreservingCallable extends AdditionalTaintStep {
   override predicate step(DataFlow::Node node1, DataFlow::Node node2) {
-    stepFromFunctionalExpToPromise(node1, node2) or
-    stepFromPromiseToFunctionalArgument(node1, node2)
+    stepIntoLambda(node1, node2) or
+    stepOutOfLambda(node1, node2)
   }
 
-  /**
-   * Tracks taint from return from lambda function to the outer `Promise`.
-   */
-  private predicate stepFromFunctionalExpToPromise(DataFlow::Node node1, DataFlow::Node node2) {
-    exists(FunctionalExpr fe |
-      fe.asMethod().getBody().getAStmt().(ReturnStmt).getResult() = node1.asExpr() and
-      node2.asExpr().(RatpackPromiseMapMethodAccess).getArgument(0) = fe
+  predicate stepIntoLambda(DataFlow::Node node1, DataFlow::Node node2) {
+    exists(MethodAccess ma, SimpleFluentLambdaMethod sflm, int arg |
+      ma.getMethod() = sflm and
+      sflm.consumesTaint(arg) and
+      node1.asExpr() = ma.getQualifier() and
+      ma.getArgument(0).(FunctionalExpr).asMethod().getParameter(arg) = node2.asParameter()
     )
   }
 
-  /**
-   * Tracks taint from the previous `Promise` to the first argument of lambda passed to `map` or `then`.
-   */
-  private predicate stepFromPromiseToFunctionalArgument(DataFlow::Node node1, DataFlow::Node node2) {
-    exists(TaintFromQualifierToFunctionalArgumentMethodAccess ma |
-      node1.asExpr() = ma.getQualifier() and
-      ma.getArgument(0).(FunctionalExpr).asMethod().getParameter(0) = node2.asParameter()
+  predicate stepOutOfLambda(DataFlow::Node node1, DataFlow::Node node2) {
+    exists(SimpleFluentLambdaMethod sflm, MethodAccess ma, FunctionalExpr fe |
+      sflm.doesReturnTaint()
+    |
+      fe.asMethod().getBody().getAStmt().(ReturnStmt).getResult() = node1.asExpr() and
+      ma.getMethod() = sflm and
+      node2.asExpr() = ma and
+      ma.getArgument(0) = fe
     )
   }
 }