Java - Taint propagation across visitor pattern + lambdas

[chmodxxx]

Hello,

I'm dealing with this flow in java which includes the following code :

    @Override
    public void doSomething(SourceType source, Optional<Scope> maybeScope) {
        maybeScope.ifPresentOrElse(
                scope -> scope.accept(Scope.Visitor.<Void>builder()
                        .something(somethingScope -> {
                            sink1(source, somethingScope.getId());
                            return null;
                        })
                        .throwOnUnknown()
                        .build()),
                () -> sink2(source));
    }

This code doesn't propagate source to sink1.

I have confirmed that data flows through ifPresentOrElse, I think data is not flowing when we deal with visitors + lambdas .
Please correct me if i'm wrong

[aschackmull]

That ought to work.
The first thing to clarify: Is the implementation of something(..) included in the analysed code? If so, how many dispatch targets are there for the call, i.e. how many potential implementations? If this number is more than 2, try upping the fieldFlowBranchLimit to at least this number in your data flow configuration and see if that helps.
If the implementation of something(..) isn't in the db, is the parameter type of the static call target of the .something(..) call identified as a functional interface?

[chmodxxx]

Yes something() is included in the db.
I have increasing fieldFlowBranchLimit and no luck.

I used partialflow (forward) to debug where the flow is cutting from the source and it looks like it reaches this functionalExpr/Lambda :

somethingScope -> {
                            sink1(source, somethingScope.getId());
                            return null;
                        }

After that It doesn't go through the lambda body statements exprs.

[aschackmull]

After that It doesn't go through the lambda body statements exprs

And it should not - at least not at first. Instead I'd expect flow to go into the something call target and somewhere inside that call target reach a point where the callback taking the value of the lambda is called. At that point - and only then - should we see dispatch to the lambda body.
This is why I was asking about the nature of the call target of the .something(..) call. If partial flow reaches the lambda that's the argument of this call then it should also flow further into the call target assuming that it's not cut of by a low fieldFlowBranchLimit.

[chmodxxx]

Ok let me share what I have :

The something method is included in a generated class (annotated with @generated ) and definition roughly looks like this :

public interface Visitor<T> {
        T visitSomething(SomethingScope value);

        T visitUnknown(@Safe String unknownType);

        static <T> SomethingStageVisitorBuilder<T> builder() {
            return new VisitorBuilder<T>();
        }
    }

private static final class VisitorBuilder<T>
            implements SomethingStageVisitorBuilder<T>, UnknownStageVisitorBuilder<T>, Completed_StageVisitorBuilder<T> {
        private Function<SomethingScope, T> somethingVisitor;

        private Function<String, T> unknownVisitor;

        @Override
        public UnknownStageVisitorBuilder<T> something(@Nonnull Function<somethingScope, T> somethingVisitor) {
            Preconditions.checkNotNull(somethingVisitor, "somethingVisitor cannot be null");
            this.somethingVisitor = somethingVisitor;
            return this;
        }

After reaching the lambda, we do get inside the something method parameter and the last step of the partial flow is return this with a distance of 3 steps to the source , after that we don't dispatch through the calling lambda.

fieldFlowBranchLimit is 50 and explorationlimit is set to 10.

I can try to replicate this code in a class and attach it if that helps you reproduce.

[aschackmull]

Ah, right. That explains things. There's actually no call to the somethingVisitor inside the something method. That means that merely executing the something call won't execute the lambda, and that's why we never enter the lambda body. Now we continue to track the lambda back out through the return this and at that point we keep track of it being stored in the somethingVisitor field. I'd expect flow to then enter the throwUnknown call, which I'd also expect to return this, and hence allow us to track flow further into the build call. This probably doesn't call the visitor callback either, so at that point we're relying on first tracking further flow of the callback field to wherever it ends up being stored in the object returned by the build call, and then we rely on tracking that into the scope.accept call, which could presumably end up calling the callback stored in the visitor.
So all in all, there's a rather complex path to track before getting to an actual call to the lambda, so one option is to dig further into that and find out where the missing link is.
Another option is to somehow cheat and short-circuit the flow by making some context-dependent flow-model that would tell the data flow library that the callback happens early - say, already in the something call - when we can deduce from the context that it's going to happen eventually anyway. This would likely require some custom QL modeling. It would be potentially interesting to consider whether this could be suitably generalised if this can be considered a general problem. Thoughts, @atorralba ?

[aschackmull]

Fun fact: If we hadn't extracted the source of the something method, then it would just work, since we have a heuristic that says that any method call to a library method taking a parameter with a type that's a functional interface will likely call the given callback. (Here "library method" means a method for which we haven't extracted it's source implementation).

[chmodxxx]

This is a reproduction code if it helps you :
DebugClass.java :

public final class DebugClass {

    public String someMethod(AuthHeader auth) {

        Integer _val = new SomethingScope().something(arg -> {
            sink(auth);
            return null;
        });
        return "foo";
    }


    private String sink(AuthHeader _auth){
        return "foo";
    }

}

SomethingScope :

import java.util.function.Function;

public class SomethingScope {

    @SuppressWarnings({"NullAway","StrictUnusedVariable"})
    private Function<Integer, Integer> foo; 

    public Integer something(Function<Integer, Integer> func){
        this.foo = func;
        return 2;
    }
    
} 

CodeQL :

import semmle.code.java.dataflow.DataFlow
import semmle.code.java.dataflow.TaintTracking
import lib.frameworks.Tokens

module DebugConfig implements DataFlow::ConfigSig {
  predicate isSource(DataFlow::Node source) {
    exists(Method endpoint |
      endpoint.getName() = "someMethod" and
      endpoint.getParameter(0) = source.asParameter() and
      source.asParameter().getType() instanceof AuthHeader
    )
  }

  predicate isSink(DataFlow::Node sink) {
    exists(MethodCall ma | ma.getMethod().getName() = "sink" and sink.asExpr() = ma.getAnArgument())
  }

  int fieldFlowBranchLimit() { result = 10 }
}

module DebugFlow = TaintTracking::Global<DebugConfig>; // 2: TaintTracking selected

from DataFlow::Node src, DataFlow::Node snk
where DebugFlow::flow(src, snk)
select src, snk

[chmodxxx]

@aschackmull the hack I wanted to do was something along these lines :

  predicate isAdditionalFlowStep(DataFlow::Node source, DataFlow::Node sink) {
    exists(FunctionalExpr expr |
      expr = source.asExpr() and
      expr.asMethod().getBody().getAStmt().(ExprStmt).getExpr() = sink.asExpr()
    )
  }

It still doesn't work, but I'm wondering if this what you had in mind about shortcircuiting the lambda call and immediately sending the flow to the exprs of the lambda

[atorralba]

Another option is to somehow cheat and short-circuit the flow by making some context-dependent flow-model that would tell the data flow library that the callback happens early - say, already in the something call - when we can deduce from the context that it's going to happen eventually anyway. This would likely require some custom QL modeling. It would be potentially interesting to consider whether this could be suitably generalised if this can be considered a general problem. Thoughts, @atorralba ?

Not for lambdas, but in the past I've had the need to tell "this callable is implicitly called in this node" (similarly to "a content read happens in this node"), but I eventually solved it with regular flow steps to not complicate things.

So while I can definitely think of use cases for such a feature, maybe it isn't that common (or it can be worked around) to justify the effort.

[chmodxxx]

@atorralba

Would be amazing if you can correct my additionalTaintStep in the above comment, much appreciated 🙇

[atorralba]

I don't know how to do what @aschackmull says:

some context-dependent flow-model that would tell the data flow library that the callback happens early

But this is how I'd hack a flow step to get what you want (flow from captured variables into their uses inside the lambda argument of something), based on InputStreamWrapperCapturedLocalStep in the PR I linked above:

  predicate isAdditionalFlowStep(DataFlow::Node n1, DataFlow::Node n2) {
    exists(MethodCall something, SsaVariable captured, SsaImplicitInit capturer, LambdaExpr l |
      something.getMethod().hasName("something") and
      capturer.captures(captured) and
      l = something.getArgument(0) and
      n2.asExpr() = capturer.getAFirstUse() and
      l.asMethod() = n2.getEnclosingCallable()
    |
      captured.(SsaImplicitInit).isParameterDefinition(n1.asParameter()) or
      n1.asExpr() = captured.(SsaExplicitUpdate).getDefiningExpr().(VariableAssign).getSource()
    )
  }

That should add the result you want.

Note: you'll need import semmle.code.java.dataflow.SSA at the top of your query as well.

[chmodxxx]

Thanks @atorralba amazing that indeed seems to be working, I have two questions :
(1) : Do you mind explaining how it works, I have never used SsaImplicitInit/SsaVariable before
(2) : Do you think I could be able to generalize this for more internal use cases where the method is not necessarily something but it could be anything ( I could remove the getMethod().hasName() but not sure about the performance implications of that)

[atorralba]

(1) : Do you mind explaining how it works, I have never used SsaImplicitInit/SsaVariable before

The SSA module gives you classes and predicates to analyze how values are assigned to variables in a more powerful way than the regular Variable class. In this case, we mainly used it because we needed to reason about variable capture.

(2) : Do you think I could be able to generalize this for more internal use cases where the method is not necessarily something but it could be anything ( I could remove the getMethod().hasName() but not sure about the performance implications of that)

The idea of the flow step is that you have some specific context that allows you to tell that any lambda passed as an argument of that method call can be considered called, so you want to connect the variables captured by it to their uses inside the lambda. Are you sure you can say the same thing for all method calls receiving lambdas as arguments? Take into account what @aschackmull explained about how lambda flow is tracked to the actual call site of the lambda — you may create duplicate alerts and decrease precision in general, aside from potential performance problems if your step creates a lot of dataflow edges.

Ideally, you should be able to narrow down which callables implicitly call their lambda arguments, and restrict the flow step to those.