Lambdas/Anon-classes follow captured local variable safety

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

@@ -33,7 +33,10 @@
 import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.ClassTree;
 import com.sun.source.tree.ExpressionTree;
+import com.sun.source.tree.LambdaExpressionTree;
+import com.sun.source.tree.MethodTree;
 import com.sun.source.tree.Tree;
+import com.sun.source.tree.Tree.Kind;
 import com.sun.source.tree.TypeCastTree;
 import com.sun.source.tree.VariableTree;
 import com.sun.source.util.TreePath;
@@ -47,6 +50,7 @@
 import java.io.Closeable;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
@@ -70,6 +74,7 @@
 import org.checkerframework.errorprone.dataflow.analysis.TransferResult;
 import org.checkerframework.errorprone.dataflow.cfg.ControlFlowGraph;
 import org.checkerframework.errorprone.dataflow.cfg.UnderlyingAST;
+import org.checkerframework.errorprone.dataflow.cfg.UnderlyingAST.CFGStatement;
 import org.checkerframework.errorprone.dataflow.cfg.builder.CFGBuilder;
 import org.checkerframework.errorprone.dataflow.cfg.node.ArrayAccessNode;
 import org.checkerframework.errorprone.dataflow.cfg.node.ArrayCreationNode;
@@ -144,13 +149,15 @@
 import org.checkerframework.errorprone.dataflow.cfg.node.UnsignedRightShiftNode;
 import org.checkerframework.errorprone.dataflow.cfg.node.VariableDeclarationNode;
 import org.checkerframework.errorprone.dataflow.cfg.node.WideningConversionNode;
+import org.checkerframework.errorprone.javacutil.TreePathUtil;
 
 /**
  * Heavily modified fork from error-prone NullnessPropagationTransfer (apache 2).
  * @see <a href="https://github.com/google/error-prone/blob/v2.11.0/check_api/src/main/java/com/google/errorprone/dataflow/nullnesspropagation/NullnessPropagationTransfer.java">NullnessPropagationTransfer</a>
  */
 public final class SafetyPropagationTransfer implements ForwardTransferFunction<Safety, AccessPathStore<Safety>> {
 
+    private static final Matcher<Tree> THROWABLE_SUBTYPE = Matchers.isSubtypeOf(Throwable.class);
     private static final Matcher<ExpressionTree> TO_STRING =
             MethodMatchers.instanceMethod().anyClass().named("toString").withNoParameters();
 
@@ -685,12 +692,74 @@ public TransferResult<Safety, AccessPathStore<Safety>> visitAssignment(
     @Override
     public TransferResult<Safety, AccessPathStore<Safety>> visitLocalVariable(
             LocalVariableNode node, TransferInput<Safety, AccessPathStore<Safety>> input) {
-        Safety safety = hasNonNullConstantValue(node)
-                ? Safety.SAFE
-                : input.getRegularStore().valueOfAccessPath(AccessPath.fromLocalVariable(node), Safety.UNKNOWN);
+        if (hasNonNullConstantValue(node)) {
+            return noStoreChanges(Safety.SAFE, input);
+        }
+        AccessPath accessPath = AccessPath.fromLocalVariable(node);
+        Safety safety = input.getRegularStore().valueOfAccessPath(accessPath, null);
+        if (safety == null) {
+            // This may occur in one of several ways:
+            // 1. catch (SomeThrowable t)
+            // 2. referencing a captured local variable within a lambda
+            // 3. referencing a captured local variable within an anonymous class
+
+            // Cast a wide net for all throwables (covers catch statements)
+            if (THROWABLE_SUBTYPE.matches(node.getTree(), state)) {
+                safety = Safety.UNSAFE;
+            } else {
+                // No safety information found, likely a captured reference used within a lambda or anonymous class.
+                safety = getCapturedLocalVariableSafety(node);
+            }
+        }
         return noStoreChanges(safety, input);
     }
 
+    private Safety getCapturedLocalVariableSafety(LocalVariableNode node) {
+        Symbol symbol = ASTHelpers.getSymbol(node.getTree());
+        if (!(symbol instanceof VarSymbol)) {
+            return Safety.UNKNOWN;
+        }
+        VarSymbol variableSymbol = (VarSymbol) symbol;
+        JavacProcessingEnvironment javacEnv = JavacProcessingEnvironment.instance(state.context);
+        TreePath variableDefinition = Trees.instance(javacEnv).getPath(variableSymbol);
+        if (variableDefinition == null) {
+            return Safety.UNKNOWN;
+        }
+        TreePath enclosingPath =
+                TreePathUtil.pathTillOfKind(variableDefinition, EnumSet.of(Kind.METHOD, Kind.LAMBDA_EXPRESSION));
+        if (enclosingPath == null) {
+            return Safety.UNKNOWN;
+        }
+        if (!traversed.add(variableSymbol)) {
+            // Avoid infinite recursion between sub-analysis cycles
+            return Safety.UNKNOWN;
+        }
+        try {
+            UnderlyingAST ast = createAst(enclosingPath);
+            ControlFlowGraph cfg = CFGBuilder.build(state.getPath().getCompilationUnit(), ast, false, false, javacEnv);
+            Analysis<Safety, AccessPathStore<Safety>, SafetyPropagationTransfer> analysis =
+                    new ForwardAnalysisImpl<>(this);
+            analysis.performAnalysis(cfg);
+            Safety maybeResult = analysis.getValue(variableDefinition.getLeaf());
+            return maybeResult == null ? Safety.UNKNOWN : maybeResult;
+        } finally {
+            traversed.remove(variableSymbol);
+        }
+    }
+
+    private static UnderlyingAST createAst(TreePath path) {
+        Tree tree = path.getLeaf();
+        ClassTree enclosingClass = TreePathUtil.enclosingClass(path);
+        if (tree instanceof MethodTree) {
+            return new UnderlyingAST.CFGMethod((MethodTree) tree, enclosingClass);
+        }
+        if (tree instanceof LambdaExpressionTree) {
+            return new UnderlyingAST.CFGLambda(
+                    (LambdaExpressionTree) tree, enclosingClass, TreePathUtil.enclosingMethod(path));
+        }
+        return new CFGStatement(tree, enclosingClass);
+    }
+
     private static boolean hasNonNullConstantValue(LocalVariableNode node) {
         if (node.getElement() instanceof VariableElement) {
             VariableElement element = (VariableElement) node.getElement();

baseline-error-prone/src/test/java/com/palantir/baseline/errorprone/IllegalSafeLoggingLambdaTest.java

@@ -0,0 +1,92 @@
+/*
+ * (c) Copyright 2022 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.palantir.baseline.errorprone;
+
+import com.google.errorprone.CompilationTestHelper;
+import org.junit.jupiter.api.Test;
+
+class IllegalSafeLoggingLambdaTest {
+
+    @Test
+    void testLambdaReferencesUnsafeExternalData() {
+        helper().addSourceLines(
+                        "Test.java",
+                        "import com.palantir.logsafe.*;",
+                        "import java.util.*;",
+                        "class Test {",
+                        "  Runnable f(RuntimeException exception) {",
+                        "    String message = exception.getMessage();",
+                        "      // BUG: Diagnostic contains: Dangerous argument value: arg is 'UNSAFE' "
+                                + "but the parameter requires 'SAFE'.",
+                        "    return () -> fun(message);",
+                        "  }",
+                        "  void fun(@Safe Object in) {}",
+                        "}")
+                .doTest();
+    }
+
+    @Test
+    void testAnonymousClassReferencesUnsafeExternalData() {
+        helper().addSourceLines(
+                        "Test.java",
+                        "import com.palantir.logsafe.*;",
+                        "import java.util.*;",
+                        "class Test {",
+                        "  Runnable f(RuntimeException exception) {",
+                        "    String message = exception.getMessage();",
+                        "    return new Runnable() {",
+                        "      @Override public void run() {",
+                        "        // BUG: Diagnostic contains: Dangerous argument value: arg is 'UNSAFE' "
+                                + "but the parameter requires 'SAFE'.",
+                        "        fun(message);",
+                        "      }",
+                        "    };",
+                        "  }",
+                        "  void fun(@Safe Object in) {}",
+                        "}")
+                .doTest();
+    }
+
+    @Test
+    void testNestedAnonymousInLambdaUnsafeExternalData() {
+        helper().addSourceLines(
+                        "Test.java",
+                        "import com.palantir.logsafe.*;",
+                        "import java.util.*;",
+                        "import java.util.function.*;",
+                        "class Test {",
+                        "  Function<RuntimeException, Runnable> f() {",
+                        "    return exception -> {",
+                        "      String message = exception.getMessage();",
+                        "      return new Runnable() {",
+                        "        @Override public void run() {",
+                        "          // BUG: Diagnostic contains: Dangerous argument value: arg is 'UNSAFE' "
+                                + "but the parameter requires 'SAFE'.",
+                        "          fun(message);",
+                        "        }",
+                        "      };",
+                        "    };",
+                        "  }",
+                        "  void fun(@Safe Object in) {}",
+                        "}")
+                .doTest();
+    }
+
+    private CompilationTestHelper helper() {
+        return CompilationTestHelper.newInstance(IllegalSafeLoggingArgument.class, getClass());
+    }
+}

changelog/@unreleased/pr-2177.v2.yml

@@ -0,0 +1,5 @@
+type: improvement
+improvement:
+  description: Lambdas/Anon-classes follow captured local variable safety
+  links:
+  - https://github.com/palantir/gradle-baseline/pull/2177