improvement: suggested fixes for IllegalSafeLoggingArgument check

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

@@ -20,6 +20,8 @@
 import com.google.errorprone.BugPattern;
 import com.google.errorprone.VisitorState;
 import com.google.errorprone.bugpatterns.BugChecker;
+import com.google.errorprone.fixes.SuggestedFix;
+import com.google.errorprone.fixes.SuggestedFixes;
 import com.google.errorprone.matchers.Description;
 import com.google.errorprone.matchers.Matcher;
 import com.google.errorprone.matchers.method.MethodMatchers;
@@ -38,8 +40,6 @@
  * Ensures that safe-logging annotated elements are handled correctly by annotated method parameters.
  * Potential future work:
  * <ul>
- *     <li>Suggest replacing {@code UnsafeArg.of(name, verifiableSafeValue)} with
- *     {@code SafeArg.of(name, verifiableSafeValue)}</li>
  *     <li>We could check return statements in methods annotated for
  *     safety to require consistency</li>
  *     <li>Enforce propagation of safety annotations from fields and types to types which encapsulate them.</li>
@@ -58,6 +58,11 @@ public final class IllegalSafeLoggingArgument extends BugChecker implements BugC
     private static final String UNSAFE = "com.palantir.logsafe.Unsafe";
     private static final String DO_NOT_LOG = "com.palantir.logsafe.DoNotLog";
 
+    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();
 
@@ -90,13 +95,25 @@ public Description matchMethodInvocation(MethodInvocationTree tree, VisitorState
                             .setMessage(String.format(
                                     "Dangerous argument value: arg is '%s' but the parameter requires '%s'.",
                                     argumentSafety, parameterSafety))
+                            .addFix(getSuggestedFix(tree, state, argumentSafety))
                             .build());
                 }
             }
         }
         return Description.NO_MATCH;
     }
 
+    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();
+            String unsafeQualifiedClassName = SuggestedFixes.qualifyType(state, fix, UNSAFE_ARG);
+            String replacement = String.format("%s.of", unsafeQualifiedClassName);
+            return fix.replace(tree.getMethodSelect(), replacement).build();
+        }
+
+        return SuggestedFix.emptyFix();
+    }
+
     private static Safety getSafety(ExpressionTree tree, VisitorState state) {
         Tree argument = ASTHelpers.stripParentheses(tree);
         // Check annotations on the result type

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

@@ -288,7 +288,75 @@ public void testIncomingParameter_toStringUsesObjectSafety() {
                 .doTest();
     }
 
+    @Test
+    public void testSafeArgOfUnsafe_recommendsUnsafeArgOf() {
+        refactoringHelper()
+                .addInputLines(
+                        "Test.java",
+                        "import com.palantir.logsafe.SafeArg;",
+                        "import com.palantir.logsafe.Unsafe;",
+                        "class Test {",
+                        "  @Unsafe static class UnsafeClass {}",
+                        "  void f() {",
+                        "    SafeArg.of(\"unsafe\", new UnsafeClass());",
+                        "  }",
+                        "}")
+                .addOutputLines(
+                        "Test.java",
+                        "import com.palantir.logsafe.SafeArg;",
+                        "import com.palantir.logsafe.Unsafe;",
+                        "import com.palantir.logsafe.UnsafeArg;",
+                        "class Test {",
+                        "  @Unsafe static class UnsafeClass {}",
+                        "  void f() {",
+                        "    UnsafeArg.of(\"unsafe\", new UnsafeClass());",
+                        "  }",
+                        "}")
+                .doTest();
+    }
+
+    @Test
+    public void testSafeArgOfDoNotLog_noRecommendation() {
+        refactoringHelper()
+                .addInputLines(
+                        "Test.java",
+                        "import com.palantir.logsafe.SafeArg;",
+                        "import com.palantir.logsafe.DoNotLog;",
+                        "class Test {",
+                        "  @DoNotLog static class DoNotLogClass {}",
+                        "  void f() {",
+                        "    // BUG: Diagnostic contains: Dangerous argument value: arg is 'DO_NOT_LOG'",
+                        "    SafeArg.of(\"unsafe\", new DoNotLogClass());",
+                        "  }",
+                        "}")
+                .expectUnchanged()
+                .doTest();
+    }
+
+    @Test
+    public void testUnsafeArgumentForSafeParameter_noRecommendation() {
+        refactoringHelper()
+                .addInputLines(
+                        "Test.java",
+                        "import com.palantir.logsafe.Safe;",
+                        "import com.palantir.logsafe.Unsafe;",
+                        "class Test {",
+                        "  @Unsafe static class UnsafeClass {}",
+                        "  void f() {",
+                        "    // BUG: Diagnostic contains: Dangerous argument value: arg is 'UNSAFE'",
+                        "    fun(new UnsafeClass());",
+                        "  }",
+                        "  private static void fun(@Safe Object obj) {}",
+                        "}")
+                .expectUnchanged()
+                .doTest();
+    }
+
     private CompilationTestHelper helper() {
         return CompilationTestHelper.newInstance(IllegalSafeLoggingArgument.class, getClass());
     }
+
+    private RefactoringValidator refactoringHelper() {
+        return RefactoringValidator.of(IllegalSafeLoggingArgument.class, getClass());
+    }
 }

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

@@ -0,0 +1,5 @@
+type: improvement
+improvement:
+  description: suggested fixes for IllegalSafeLoggingArgument check
+  links:
+  - https://github.com/palantir/gradle-baseline/pull/2133

gradle-baseline-java/src/main/groovy/com/palantir/baseline/extensions/BaselineErrorProneExtension.java

@@ -42,6 +42,7 @@ public class BaselineErrorProneExtension {
             "ExecutorSubmitRunnableFutureIgnored",
             "ExtendsErrorOrThrowable",
             "FinalClass",
+            "IllegalSafeLoggingArgument",
             "ImmutablesStyle",
             "ImplicitPublicBuilderConstructor",
             "JavaTimeDefaultTimeZone",