Safety data-flow

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

@@ -26,14 +26,17 @@
 import com.google.errorprone.matchers.Matcher;
 import com.google.errorprone.matchers.method.MethodMatchers;
 import com.google.errorprone.util.ASTHelpers;
+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.ExpressionTree;
 import com.sun.source.tree.MethodInvocationTree;
-import com.sun.source.tree.Tree;
-import com.sun.source.tree.TypeCastTree;
-import com.sun.tools.javac.code.Symbol;
+import com.sun.source.tree.MethodTree;
+import com.sun.source.tree.ReturnTree;
+import com.sun.source.tree.StatementTree;
+import com.sun.source.util.TreePath;
 import com.sun.tools.javac.code.Symbol.MethodSymbol;
 import com.sun.tools.javac.code.Symbol.VarSymbol;
-import com.sun.tools.javac.code.Type;
 import java.util.List;
 
 /**
@@ -52,19 +55,14 @@
         linkType = BugPattern.LinkType.CUSTOM,
         severity = BugPattern.SeverityLevel.ERROR,
         summary = "safe-logging annotations must agree between args and method parameters")
-public final class IllegalSafeLoggingArgument extends BugChecker implements BugChecker.MethodInvocationTreeMatcher {
-    private static final String SAFE = "com.palantir.logsafe.Safe";
-    private static final String UNSAFE = "com.palantir.logsafe.Unsafe";
-    private static final String DO_NOT_LOG = "com.palantir.logsafe.DoNotLog";
+public final class IllegalSafeLoggingArgument extends BugChecker
+        implements BugChecker.MethodInvocationTreeMatcher, BugChecker.ReturnTreeMatcher {
 
     private static final String UNSAFE_ARG = "com.palantir.logsafe.UnsafeArg";
     private static final Matcher<ExpressionTree> SAFE_ARG_OF_METHOD_MATCHER = MethodMatchers.staticMethod()
             .onClass("com.palantir.logsafe.SafeArg")
             .named("of");
 
-    private static final Matcher<ExpressionTree> TO_STRING =
-            MethodMatchers.instanceMethod().anyClass().named("toString").withNoParameters();
-
     @Override
     public Description matchMethodInvocation(MethodInvocationTree tree, VisitorState state) {
         List<? extends ExpressionTree> arguments = tree.getArguments();
@@ -78,16 +76,18 @@ public Description matchMethodInvocation(MethodInvocationTree tree, VisitorState
         List<VarSymbol> parameters = methodSymbol.getParameters();
         for (int i = 0; i < parameters.size(); i++) {
             VarSymbol parameter = parameters.get(i);
-            Safety parameterSafety = getSafety(parameter, state);
-            if (parameterSafety == Safety.UNKNOWN) {
-                // Fast path, avoid analysis when the value isn't provided to a safety-aware consumer
+            Safety parameterSafety = SafetyAnnotations.getSafety(parameter, state);
+            if (parameterSafety.allowsAll()) {
+                // Fast path: all types are accepted, there's no reason to do further analysis.
                 continue;
             }
 
             int limit = methodSymbol.isVarArgs() && i == parameters.size() - 1 ? arguments.size() : i + 1;
             for (int j = i; j < limit; j++) {
                 ExpressionTree argument = arguments.get(j);
-                Safety argumentSafety = getSafety(argument, state);
+
+                Safety argumentSafety = SafetyAnalysis.of(state.withPath(new TreePath(state.getPath(), argument)));
+
                 if (!parameterSafety.allowsValueWith(argumentSafety)) {
                     // use state.reportMatch to report all failing arguments if multiple are invalid
                     state.reportMatch(buildDescription(argument)
@@ -102,6 +102,36 @@ public Description matchMethodInvocation(MethodInvocationTree tree, VisitorState
         return Description.NO_MATCH;
     }
 
+    @Override
+    public Description matchReturn(ReturnTree tree, VisitorState state) {
+        if (tree.getExpression() == null) {
+            return Description.NO_MATCH;
+        }
+        TreePath path = state.getPath();
+        while (path != null && path.getLeaf() instanceof StatementTree) {
+            path = path.getParentPath();
+        }
+        if (path == null || !(path.getLeaf() instanceof MethodTree)) {
+            return Description.NO_MATCH;
+        }
+        MethodTree method = (MethodTree) path.getLeaf();
+        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;
+        }
+        Safety returnValueSafety =
+                SafetyAnalysis.of(state.withPath(new TreePath(state.getPath(), tree.getExpression())));
+        if (methodDeclaredSafety.allowsValueWith(returnValueSafety)) {
+            return Description.NO_MATCH;
+        }
+        return buildDescription(tree)
+                .setMessage(String.format(
+                        "Dangerous return value: result is '%s' but the method is annotated '%s'.",
+                        returnValueSafety, methodDeclaredSafety))
+                .build();
+    }
+
     private static SuggestedFix getSuggestedFix(MethodInvocationTree tree, VisitorState state, Safety argumentSafety) {
         if (SAFE_ARG_OF_METHOD_MATCHER.matches(tree, state) && Safety.UNSAFE.allowsValueWith(argumentSafety)) {
             SuggestedFix.Builder fix = SuggestedFix.builder();
@@ -112,83 +142,4 @@ private static SuggestedFix getSuggestedFix(MethodInvocationTree tree, VisitorSt
 
         return SuggestedFix.emptyFix();
     }
-
-    private static Safety getSafety(ExpressionTree tree, VisitorState state) {
-        Tree argument = ASTHelpers.stripParentheses(tree);
-        // Check annotations on the result type
-        Type resultType = ASTHelpers.getResultType(tree);
-        if (resultType != null) {
-            Safety resultTypeSafety = getSafety(resultType.tsym, state);
-            if (resultTypeSafety != Safety.UNKNOWN) {
-                return resultTypeSafety;
-            }
-        }
-        // Unwrap type-casts: 'Object value = (Object) unsafeType;' is still unsafe.
-        if (argument instanceof TypeCastTree) {
-            TypeCastTree typeCastTree = (TypeCastTree) argument;
-            return getSafety(typeCastTree.getExpression(), state);
-        }
-        // If the argument is a method invocation, check the method for safety annotations
-        if (argument instanceof MethodInvocationTree) {
-            MethodInvocationTree argumentInvocation = (MethodInvocationTree) argument;
-            MethodSymbol methodSymbol = ASTHelpers.getSymbol(argumentInvocation);
-            if (methodSymbol != null) {
-                Safety methodSafety = getSafety(methodSymbol, state);
-                // non-annotated toString inherits type-level safety.
-                if (methodSafety == Safety.UNKNOWN && TO_STRING.matches(argumentInvocation, state)) {
-                    return getSafety(ASTHelpers.getReceiver(argumentInvocation), state);
-                }
-                return methodSafety;
-            }
-        }
-        // Check the argument symbol itself:
-        Symbol argumentSymbol = ASTHelpers.getSymbol(argument);
-        return argumentSymbol != null ? getSafety(argumentSymbol, state) : Safety.UNKNOWN;
-    }
-
-    private static Safety getSafety(Symbol symbol, VisitorState state) {
-        if (ASTHelpers.hasAnnotation(symbol, DO_NOT_LOG, state)) {
-            return Safety.DO_NOT_LOG;
-        }
-        if (ASTHelpers.hasAnnotation(symbol, UNSAFE, state)) {
-            return Safety.UNSAFE;
-        }
-        if (ASTHelpers.hasAnnotation(symbol, SAFE, state)) {
-            return Safety.SAFE;
-        }
-        return Safety.UNKNOWN;
-    }
-
-    private enum Safety {
-        UNKNOWN() {
-            @Override
-            boolean allowsValueWith(Safety _valueSafety) {
-                // No constraints when safety isn't specified
-                return true;
-            }
-        },
-        DO_NOT_LOG() {
-            @Override
-            boolean allowsValueWith(Safety _valueSafety) {
-                // do-not-log on a parameter isn't meaningful for callers, only for the implementation
-                return true;
-            }
-        },
-        UNSAFE() {
-            @Override
-            boolean allowsValueWith(Safety valueSafety) {
-                // We allow safe data to be provided to an unsafe annotated parameter because that's safe, however
-                // we should separately flag and prompt migration of such UnsafeArgs to SafeArg.
-                return valueSafety != Safety.DO_NOT_LOG;
-            }
-        },
-        SAFE() {
-            @Override
-            boolean allowsValueWith(Safety valueSafety) {
-                return valueSafety == Safety.UNKNOWN || valueSafety == Safety.SAFE;
-            }
-        };
-
-        abstract boolean allowsValueWith(Safety valueSafety);
-    }
 }

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

@@ -0,0 +1,96 @@
+/*
+ * (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.safety;
+
+import org.checkerframework.errorprone.dataflow.analysis.AbstractValue;
+
+public enum Safety implements AbstractValue<Safety> {
+    UNKNOWN() {
+        @Override
+        public Safety leastUpperBound(Safety other) {
+            return other == SAFE ? this : other;
+        }
+
+        @Override
+        public boolean allowsValueWith(Safety _valueSafety) {
+            // No constraints when safety isn't specified
+            return true;
+        }
+    },
+    DO_NOT_LOG() {
+        @Override
+        public Safety leastUpperBound(Safety _other) {
+            return this;
+        }
+
+        @Override
+        public boolean allowsValueWith(Safety _valueSafety) {
+            // do-not-log on a parameter isn't meaningful for callers, only for the implementation
+            return true;
+        }
+    },
+    UNSAFE() {
+        @Override
+        public Safety leastUpperBound(Safety other) {
+            return other == DO_NOT_LOG ? other : this;
+        }
+
+        @Override
+        public boolean allowsValueWith(Safety valueSafety) {
+            // We allow safe data to be provided to an unsafe annotated parameter because that's safe, however
+            // we should separately flag and prompt migration of such UnsafeArgs to SafeArg.
+            return valueSafety != Safety.DO_NOT_LOG;
+        }
+    },
+    SAFE() {
+        @Override
+        public Safety leastUpperBound(Safety other) {
+            return other;
+        }
+
+        @Override
+        public boolean allowsValueWith(Safety valueSafety) {
+            return valueSafety == Safety.UNKNOWN || valueSafety == Safety.SAFE;
+        }
+    };
+
+    public abstract boolean allowsValueWith(Safety valueSafety);
+
+    public boolean allowsAll() {
+        return this == UNKNOWN || this == DO_NOT_LOG;
+    }
+
+    /**
+     * Merge Safety using {@link Safety#leastUpperBound(AbstractValue)} except that {@link Safety#UNKNOWN} assumes
+     * no confidence, preferring the other type if data is available.
+     * For example, casting from {@link Object} to a known-safe type should result in {@link Safety#SAFE}.
+     */
+    public static Safety mergeAssumingUnknownIsSame(Safety one, Safety two) {
+        if (one == UNKNOWN) {
+            return two;
+        }
+        if (two == UNKNOWN) {
+            return one;
+        }
+        return one.leastUpperBound(two);
+    }
+
+    public static Safety mergeAssumingUnknownIsSame(Safety one, Safety two, Safety three) {
+        Safety result = mergeAssumingUnknownIsSame(one, two);
+        return mergeAssumingUnknownIsSame(result, three);
+    }
+}

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

@@ -0,0 +1,49 @@
+/*
+ * (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.safety;
+
+import com.google.errorprone.VisitorState;
+import com.google.errorprone.dataflow.DataFlow;
+import com.palantir.baseline.errorprone.safety.SafetyPropagationTransfer.ClearVisitorState;
+import com.sun.source.util.TreePath;
+import com.sun.tools.javac.util.Context;
+
+public final class SafetyAnalysis {
+    private static final Context.Key<SafetyPropagationTransfer> SAFETY_PROPAGATION = new Context.Key<>();
+
+    /**
+     * Returns the safety of the item at the current path.
+     * Callers may need to use {@link VisitorState#withPath(TreePath)} to provide a more specific path.
+     */
+    public static Safety of(VisitorState state) {
+        SafetyPropagationTransfer propagation = instance(state.context);
+        try (ClearVisitorState ignored = propagation.setVisitorState(state)) {
+            return DataFlow.expressionDataflow(state.getPath(), state.context, propagation);
+        }
+    }
+
+    private static SafetyPropagationTransfer instance(Context context) {
+        SafetyPropagationTransfer instance = context.get(SAFETY_PROPAGATION);
+        if (instance == null) {
+            instance = new SafetyPropagationTransfer();
+            context.put(SAFETY_PROPAGATION, instance);
+        }
+        return instance;
+    }
+
+    private SafetyAnalysis() {}
+}