Add additional Ratpack test and improve Promise based dataflow tracking

Author: JLLeitschuh

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

@@ -3,6 +3,8 @@
  */
 
 import java
+private import semmle.code.java.dataflow.DataFlow
+private import semmle.code.java.dataflow.FlowSteps
 
 /**
  * Ratpack methods that access user-supplied request data.
@@ -81,3 +83,31 @@ class RatpackUploadFileGetMethod extends RatpackGetRequestDataMethod {
     hasName("getFileName")
   }
 }
+
+class RatpackHeader extends RefType {
+  RatpackHeader() {
+    hasQualifiedName("ratpack.http", "Headers") or
+    hasQualifiedName("ratpack.core.http", "Headers")
+  }
+}
+
+private class RatpackHeaderTaintPropigatingMethod extends Method {
+  RatpackHeaderTaintPropigatingMethod() {
+    getDeclaringType() instanceof RatpackHeader and
+    hasName(["get", "getAll", "getNames", "asMultiValueMap"])
+  }
+}
+
+class TaintPropigatingHeaderMethod extends AdditionalTaintStep {
+  override predicate step(DataFlow::Node node1, DataFlow::Node node2) {
+    stepHeaderPropigatingTaint(node1, node2)
+  }
+
+  private predicate stepHeaderPropigatingTaint(DataFlow::Node node1, DataFlow::Node node2) {
+    exists(MethodAccess ma |
+      ma.getMethod() instanceof RatpackHeaderTaintPropigatingMethod and
+      node2.asExpr() = ma and
+      node1.asExpr() = ma.getQualifier()
+    )
+  }
+}

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

@@ -9,14 +9,20 @@ 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 { }
+
 class RatpackPromiseMapMethod extends Method {
   RatpackPromiseMapMethod() {
     getDeclaringType() instanceof RatpackPromise and
     hasName("map")
   }
 }
 
-class RatpackPromiseMapMethodAccess extends MethodAccess {
+class RatpackPromiseMapMethodAccess extends TaintFromQualifierToFunctionalArgumentMethodAccess {
   RatpackPromiseMapMethodAccess() { getMethod() instanceof RatpackPromiseMapMethod }
 }
 
@@ -27,10 +33,21 @@ class RatpackPromiseThenMethod extends Method {
   }
 }
 
-class RatpackPromiseThenMethodAccess extends MethodAccess {
+class RatpackPromiseThenMethodAccess extends TaintFromQualifierToFunctionalArgumentMethodAccess {
   RatpackPromiseThenMethodAccess() { getMethod() instanceof RatpackPromiseThenMethod }
 }
 
+class RatpackPromiseNextMethod extends FluentMethod {
+  RatpackPromiseNextMethod() {
+    getDeclaringType() instanceof RatpackPromise and
+    hasName("next")
+  }
+}
+
+class RatpackPromiseNextMethodAccess extends TaintFromQualifierToFunctionalArgumentMethodAccess {
+  RatpackPromiseNextMethodAccess() { getMethod() instanceof RatpackPromiseNextMethod }
+}
+
 private class RatpackPromiseTaintPreservingCallable extends AdditionalTaintStep {
   override predicate step(DataFlow::Node node1, DataFlow::Node node2) {
     stepFromFunctionalExpToPromise(node1, node2) or
@@ -51,12 +68,7 @@ private class RatpackPromiseTaintPreservingCallable extends AdditionalTaintStep
    * 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(RatpackPromiseMapMethodAccess ma |
-      node1.asExpr() = ma.getQualifier() and
-      ma.getArgument(0).(FunctionalExpr).asMethod().getParameter(0) = node2.asParameter()
-    )
-    or
-    exists(RatpackPromiseThenMethodAccess ma |
+    exists(TaintFromQualifierToFunctionalArgumentMethodAccess ma |
       node1.asExpr() = ma.getQualifier() and
       ma.getArgument(0).(FunctionalExpr).asMethod().getParameter(0) = node2.asParameter()
     )

java/ql/test/library-tests/frameworks/ratpack/resources/Resource.java

@@ -1,5 +1,7 @@
 import ratpack.core.handling.Context;
 import ratpack.core.http.TypedData;
+import ratpack.core.form.UploadedFile;
+import java.io.OutputStream;
 
 class Resource {
 
@@ -10,18 +12,45 @@ void test1(Context ctx) {
         sink(ctx.getRequest().getCookies()); //$hasTaintFlow
         sink(ctx.getRequest().oneCookie("Magic-Cookie")); //$hasTaintFlow
         sink(ctx.getRequest().getHeaders()); //$hasTaintFlow
+        sink(ctx.getRequest().getHeaders().get("questionable_header")); //$hasTaintFlow
+        sink(ctx.getRequest().getHeaders().getAll("questionable_header")); //$hasTaintFlow
+        sink(ctx.getRequest().getHeaders().getNames()); //$hasTaintFlow
+        sink(ctx.getRequest().getHeaders().asMultiValueMap()); //$hasTaintFlow
+        sink(ctx.getRequest().getHeaders().asMultiValueMap().get("questionable_header")); //$hasTaintFlow
         sink(ctx.getRequest().getPath()); //$hasTaintFlow
         sink(ctx.getRequest().getQuery()); //$hasTaintFlow
         sink(ctx.getRequest().getQueryParams()); //$hasTaintFlow
+        sink(ctx.getRequest().getQueryParams().get("questionable_parameter")); //$hasTaintFlow
         sink(ctx.getRequest().getRawUri()); //$hasTaintFlow
         sink(ctx.getRequest().getUri()); //$hasTaintFlow
     }
 
     void test2(TypedData td) {
         sink(td.getText()); //$hasTaintFlow
+        sink(td.getBuffer()); //$hasTaintFlow
+        sink(td.getBytes()); //$hasTaintFlow
+        sink(td.getContentType()); //$hasTaintFlow
+        sink(td.getInputStream()); //$hasTaintFlow
     }
 
-    void test2(Context ctx) {
+    void test3(TypedData td, OutputStream os) throws java.io.IOException {
+        sink(os);
+        td.writeTo(os);
+        sink(os); //$hasTaintFlow
+    }
+
+    void test4(UploadedFile uf) {
+        sink(uf.getFileName()); //$hasTaintFlow
+    }
+
+    void test5(Context ctx) {
+        sink(ctx.getRequest().getBody().map(TypedData::getText)); //$hasTaintFlow
         ctx.getRequest().getBody().map(TypedData::getText).then(this::sink); //$hasTaintFlow
+        ctx
+            .getRequest()
+            .getBody()
+            .map(TypedData::getText)
+            .next(this::sink) //$hasTaintFlow
+            .then(this::sink); //$hasTaintFlow
     }
 }