[Java] Tracking Taint through Fluent Libraries

[JLLeitschuh]

I'm attempting to track taint as it flows through a fluent async promise-based library built into the ratpack framework.

In particular, I'm attempting to track taint in the following situation:

void test2(Context ctx) {
    ctx.getRequest().getBody().map(TypedData::source).then(this::sink); //$hasTaintFlow
}

The source is TypedData::source and the sink is this::sink.

This is query that I've come up with so far:

import java
private import semmle.code.java.dataflow.DataFlow
private import semmle.code.java.dataflow.FlowSteps

/** A reference type that extends a parameterization the Promise type. */
class RatpackPromise extends RefType {
  RatpackPromise() {
    getSourceDeclaration().getASourceSupertype*().hasQualifiedName("ratpack.exec", "Promise")
  }
}

class RatpackPromiseMapMethod extends Method {
  RatpackPromiseMapMethod() {
    getDeclaringType() instanceof RatpackPromise and
    hasName("map")
  }
}

class RatpackPromiseMapMethodAccess extends MethodAccess {
  RatpackPromiseMapMethodAccess() { getMethod() instanceof RatpackPromiseMapMethod }
}

class RatpackPromiseThenMethod extends Method {
  RatpackPromiseThenMethod() {
    getDeclaringType() instanceof RatpackPromise and
    hasName("then")
  }
}

class RatpackPromiseThenMethodAccess extends MethodAccess {
  RatpackPromiseThenMethodAccess() { getMethod() instanceof RatpackPromiseThenMethod }
}

private class RatpackPromiseTaintPreservingCallable extends AdditionalTaintStep {
  override predicate step(DataFlow::Node node1, DataFlow::Node node2) {
    stepFromFunctionalExpToPromise(node1, node2) or
    stepFromPromiseToFunctionalArgument(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
    )
  }

  /**
   * 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 |
      node1.asExpr() = ma.getQualifier() and
      ma.getArgument(0).(FunctionalExpr).asMethod().getParameter(0) = node2.asParameter()
    )
  }
}

If it helps, this is the definition of Promise:

public interface Promise<T> {
  <O> Promise<O> map(Function<? super T, ? extends O> transformer);

  void then(Action<? super T> then);
}

Independently running stepFromFunctionalExpToPromise & stepFromPromiseToFunctionalArgument with "quick evaluate" I see that I can correctly capture the data flow nodes for the return value TypedData::source expression to the map MethodAccess. I can also see data flow from the qualifier Promise to the first argument parameter of map or then. However, in aggregate, I'm not seeing dataflow from TypedData::source to the call to sink. I'm wondering what I'm missing.

---

Relevant PR: #4991

[intrigus-lgtm]

I'd say this is a perfect question for @smowton as he (afaik) recently added support for modeling fluent methods in Java.

[smowton]

CC'd from Slack for future people finding this question:

The modern way (tm) is to extend SummaryModelCsv, like this: https://github.com/github/codeql/blob/main/java/ql/src/semmle/code/java/frameworks/apache/Lang.qll#L601

The last entry in these “CSV” (really SSV, semicolon-separated-values) rows is either taint or value depending on whether the model propagates taint or the exact value of an argument or qualifier (edited)

Fluent APIs are modelled by a value-flow step from Argument[-1] (the qualifier) to ReturnValue, plus optionally taint steps from Argument[something] to Argument[-1] if this fluent operation can add taint.

[JLLeitschuh]

I think the difference here is that the Apache methods you're pointing out here have taint flow through them from the qualifier.

What I'm trying to capture is data flow that bounces back and forth from the qualifier, to the lambda argument.

map((TypedData td) -> td.getText() /* source of taint */) /* taint transferred to `map` return value */.then((String taintedString) -> sink(taintedString));

The documentation you're pointing me to is about how taint propagates purely through a fluent API. Not a fluent API that uses async/lambda/stream based data flow.

[JLLeitschuh]

I figured it out. My query is actually correct. I just goofed and had failed to import private import semmle.code.java.frameworks.ratpack.RatpackExec into the Framework module so my additional taint steps weren't being loaded. My overall approach here was correct though.