baseline-error-prone/src/main/java/com/palantir/baseline/errorprone/IllegalSafeLoggingArgument.java

@@ -26,13 +26,15 @@
 import com.google.errorprone.matchers.Matcher;
 import com.google.errorprone.matchers.method.MethodMatchers;
 import com.google.errorprone.util.ASTHelpers;
+import com.google.errorprone.util.ASTHelpers.TargetType;
 import com.palantir.baseline.errorprone.safety.Safety;
 import com.palantir.baseline.errorprone.safety.SafetyAnalysis;
 import com.palantir.baseline.errorprone.safety.SafetyAnnotations;
 import com.sun.source.tree.AssignmentTree;
 import com.sun.source.tree.ClassTree;
 import com.sun.source.tree.CompoundAssignmentTree;
 import com.sun.source.tree.ExpressionTree;
+import com.sun.source.tree.LambdaExpressionTree;
 import com.sun.source.tree.MethodInvocationTree;
 import com.sun.source.tree.MethodTree;
 import com.sun.source.tree.NewClassTree;
@@ -78,7 +80,8 @@ public final class IllegalSafeLoggingArgument extends BugChecker
                 BugChecker.MethodTreeMatcher,
                 BugChecker.VariableTreeMatcher,
                 BugChecker.NewClassTreeMatcher,
-                BugChecker.ClassTreeMatcher {
+                BugChecker.ClassTreeMatcher,
+                BugChecker.LambdaExpressionTreeMatcher {
 
     private static final String UNSAFE_ARG = "com.palantir.logsafe.UnsafeArg";
     private static final Matcher<ExpressionTree> SAFE_ARG_OF_METHOD_MATCHER = MethodMatchers.staticMethod()
@@ -260,11 +263,20 @@ public Description matchReturn(ReturnTree tree, VisitorState state) {
         while (path != null && path.getLeaf() instanceof StatementTree) {
             path = path.getParentPath();
         }
-        if (path == null || !(path.getLeaf() instanceof MethodTree)) {
-            return Description.NO_MATCH;
+        if (path != null) {
+            if (path.getLeaf() instanceof MethodTree method) {
+                return handleMethodReturn(tree, method, state);
+            } else if (path.getLeaf() instanceof LambdaExpressionTree lambda) {
+                return handleLambdaReturn(tree, lambda, state);
+            }
         }
-        MethodTree method = (MethodTree) path.getLeaf();
+
+        return Description.NO_MATCH;
+    }
+
+    private Description handleMethodReturn(ReturnTree tree, MethodTree method, VisitorState state) {
         Safety methodDeclaredSafety = SafetyAnnotations.getSafety(ASTHelpers.getSymbol(method), state);
+
         if (methodDeclaredSafety.allowsAll()) {
             // Fast path, all types are accepted, there's no reason to do further analysis.
             return Description.NO_MATCH;
@@ -281,6 +293,25 @@ public Description matchReturn(ReturnTree tree, VisitorState state) {
                 .build();
     }
 
+    private Description handleLambdaReturn(ReturnTree tree, LambdaExpressionTree lambda, VisitorState state) {
+        Safety requiredSafety = getLambdaRequiredReturnSafety(lambda, state);
+
+        if (requiredSafety.allowsAll()) {
+            // Fast path, all types are accepted, there's no reason to do further analysis.
+            return Description.NO_MATCH;
+        }
+        Safety returnValueSafety =
+                SafetyAnalysis.of(state.withPath(new TreePath(state.getPath(), tree.getExpression())));
+        if (requiredSafety.allowsValueWith(returnValueSafety)) {
+            return Description.NO_MATCH;
+        }
+        return buildDescription(tree)
+                .setMessage(String.format(
+                        "Dangerous return value: result is '%s' but the lambda expects return '%s'.",
+                        returnValueSafety, requiredSafety))
+                .build();
+    }
+
     @Override
     public Description matchAssignment(AssignmentTree tree, VisitorState state) {
         return handleAssignment(tree, tree.getVariable(), tree.getExpression(), state);
@@ -427,4 +458,43 @@ public Description matchClass(ClassTree tree, VisitorState state) {
                         "Dangerous type: annotated '%s' but ancestors declare '%s'.", directSafety, ancestorSafety))
                 .build();
     }
+
+    @Override
+    public Description matchLambdaExpression(LambdaExpressionTree tree, VisitorState state) {
+        Safety requiredReturnSafety = getLambdaRequiredReturnSafety(tree, state);
+
+        if (requiredReturnSafety.allowsAll()) {
+            // Short-circuit if the return type allows all values
+            return Description.NO_MATCH;
+        }
+
+        Safety resultSafety = Safety.UNKNOWN;
+        switch (tree.getBodyKind()) {
+            case EXPRESSION:
+                resultSafety = SafetyAnalysis.of(state.withPath(new TreePath(state.getPath(), tree.getBody())));
+                break;
+            case STATEMENT:
+                // Shortcut - statement lambdas get their return type checked in the return statement matcher
+                // This also allows us to indicate which return statement is bad (if any) rather than the lambda itself
+                return Description.NO_MATCH;
+        }
+
+        if (requiredReturnSafety.allowsValueWith(resultSafety)) {
+            return Description.NO_MATCH;
+        }
+
+        return buildDescription(tree)
+                .setMessage(String.format(
+                        "Dangerous return value: result is '%s' but the lambda expects return '%s'.",
+                        resultSafety, requiredReturnSafety))
+                .build();
+    }
+
+    private Safety getLambdaRequiredReturnSafety(LambdaExpressionTree tree, VisitorState state) {
+        TargetType returnType = ASTHelpers.targetType(state.withPath(new TreePath(state.getPath(), tree)));
+        if (returnType == null) {
+            return Safety.UNKNOWN;
+        }
+        return SafetyAnnotations.getSafety(returnType.type(), state);
+    }
 }

baseline-error-prone/src/main/java/com/palantir/baseline/errorprone/safety/SafetyPropagationTransfer.java

@@ -438,16 +438,52 @@ public final class SafetyPropagationTransfer implements ForwardTransferFunction<
     private final Set<VarSymbol> traversed = new HashSet<>();
 
     @Override
-    public AccessPathStore<Safety> initialStore(UnderlyingAST _underlyingAst, List<LocalVariableNode> parameters) {
+    public AccessPathStore<Safety> initialStore(UnderlyingAST underlyingAst, List<LocalVariableNode> parameters) {
         if (parameters == null) {
             return AccessPathStore.empty();
         }
-        AccessPathStore.Builder<Safety> result = AccessPathStore.<Safety>empty().toBuilder();
 
-        for (LocalVariableNode param : parameters) {
-            Safety declared = SafetyAnnotations.getSafety(param.getTree(), state);
-            result.setInformation(AccessPath.fromLocalVariable(param), declared);
+        if (underlyingAst.getKind() == UnderlyingAST.Kind.LAMBDA) {
+            // Special-case lambda types as the parameter annotations are not propagated in
+            //   ForwardAnalysisImpl#getParameters, which is passed as the argument here
+            return initialStoreLambda(underlyingAst, parameters);
+        } else {
+            AccessPathStore.Builder<Safety> result = AccessPathStore.<Safety>empty().toBuilder();
+
+            for (LocalVariableNode param : parameters) {
+                Safety declared = SafetyAnnotations.getSafety(param.getTree(), state);
+                result.setInformation(AccessPath.fromLocalVariable(param), declared);
+            }
+
+            return result.build();
         }
+    }
+
+    private AccessPathStore<Safety> initialStoreLambda(
+            UnderlyingAST underlyingAst, List<LocalVariableNode> parameters) {
+        AccessPathStore.Builder<Safety> result = AccessPathStore.<Safety>empty().toBuilder();
+        LambdaExpressionTree lambda = ((UnderlyingAST.CFGLambda) underlyingAst).getLambdaTree();
+
+        // We grab the inferred type of the lambda, then the argument types of the functional interface
+        //  and match them one by one with the actual parameters
+        com.sun.tools.javac.util.List<Type> parameterTypes =
+                state.getTypes().findDescriptorType(ASTHelpers.getType(lambda)).getParameterTypes();
+
+        // This is expected to be true, but let's be defensive here
+        if (parameterTypes.size() == parameters.size()) {
+            for (int i = 0; i < parameters.size(); i++) {
+                Type type = parameterTypes.get(i);
+                LocalVariableNode param = parameters.get(i);
+
+                // If the lambda itself has a safety annotation, which doesn't match the expected one from the
+                //   type it is used as, combine both
+                Safety declared = SafetyAnnotations.getSafety(param.getTree(), state);
+                Safety safety = SafetyAnnotations.getSafety(type, state);
+                result.setInformation(
+                        AccessPath.fromLocalVariable(param), Safety.mergeAssumingUnknownIsSame(declared, safety));
+            }
+        }
+
         return result.build();
     }