Fix bugs in reading varargs annotations from bytecodes (#1055)

msridhar · web-flow · commit 91cf25db19cc · 2024-10-14T23:43:32.000Z
For the case of declaration annotations, we were using `Flags.VARARGS`
to check if a `Symbol` represented a varargs parameter, but this only
works if the method is defined in source code. We add new methods for
varargs parameters that may be defined in bytecodes, and always check
for declaration annotations in such cases.

For type use annotations on parameters, for a method from bytecodes the
annotation info is stored on the method's symbol, not the parameter's.
We borrow some Error Prone code to handle this case.
diff --git a/nullaway/src/main/java/com/uber/nullaway/NullabilityUtil.java b/nullaway/src/main/java/com/uber/nullaway/NullabilityUtil.java
@@ -55,6 +55,7 @@
 import java.util.stream.Stream;
 import javax.lang.model.element.AnnotationMirror;
 import javax.lang.model.element.AnnotationValue;
+import javax.lang.model.element.ElementKind;
 import javax.lang.model.element.ExecutableElement;
 import org.checkerframework.nullaway.javacutil.AnnotationUtils;
 import org.jspecify.annotations.Nullable;
@@ -301,9 +302,27 @@ public static Stream<? extends AnnotationMirror> getAllAnnotationsForParameter(
    * Gets the type use annotations on a symbol, ignoring annotations on components of the type (type
    * arguments, wildcards, etc.)
    */
-  public static Stream<? extends AnnotationMirror> getTypeUseAnnotations(
-      Symbol symbol, Config config) {
-    Stream<Attribute.TypeCompound> rawTypeAttributes = symbol.getRawTypeAttributes().stream();
+  public static Stream<Attribute.TypeCompound> getTypeUseAnnotations(Symbol symbol, Config config) {
+    return getTypeUseAnnotations(symbol, config, /* onlyDirect= */ true);
+  }
+
+  /**
+   * Gets the type use annotations on a symbol
+   *
+   * @param symbol the symbol
+   * @param config NullAway configuration
+   * @param onlyDirect if true, only return annotations that are directly on the type, not on
+   *     components of the type (type arguments, wildcards, array contents, etc.)
+   * @return the type use annotations on the symbol
+   */
+  private static Stream<Attribute.TypeCompound> getTypeUseAnnotations(
+      Symbol symbol, Config config, boolean onlyDirect) {
+    // Adapted from Error Prone's MoreAnnotations class:
+    // https://github.com/google/error-prone/blob/5f71110374e63f3c35b661f538295fa15b5c1db2/check_api/src/main/java/com/google/errorprone/util/MoreAnnotations.java#L84-L91
+    Symbol typeAnnotationOwner =
+        symbol.getKind().equals(ElementKind.PARAMETER) ? symbol.owner : symbol;
+    Stream<Attribute.TypeCompound> rawTypeAttributes =
+        typeAnnotationOwner.getRawTypeAttributes().stream();
     if (symbol instanceof Symbol.MethodSymbol) {
       // for methods, we want annotations on the return type
       return rawTypeAttributes.filter(
@@ -313,7 +332,45 @@ public static Stream<? extends AnnotationMirror> getTypeUseAnnotations(
     } else {
       // filter for annotations directly on the type
       return rawTypeAttributes.filter(
-          t -> NullabilityUtil.isDirectTypeUseAnnotation(t, symbol, config));
+          t ->
+              targetTypeMatches(symbol, t.position)
+                  && (!onlyDirect || NullabilityUtil.isDirectTypeUseAnnotation(t, symbol, config)));
+    }
+  }
+
+  // Adapted from Error Prone MoreAnnotations:
+  // https://github.com/google/error-prone/blob/5f71110374e63f3c35b661f538295fa15b5c1db2/check_api/src/main/java/com/google/errorprone/util/MoreAnnotations.java#L128
+  private static boolean targetTypeMatches(Symbol sym, TypeAnnotationPosition position) {
+    switch (sym.getKind()) {
+      case LOCAL_VARIABLE:
+        return position.type == TargetType.LOCAL_VARIABLE;
+      case FIELD:
+      // treated like a field
+      case ENUM_CONSTANT:
+        return position.type == TargetType.FIELD;
+      case CONSTRUCTOR:
+      case METHOD:
+        return position.type == TargetType.METHOD_RETURN;
+      case PARAMETER:
+        if (position.type.equals(TargetType.METHOD_FORMAL_PARAMETER)) {
+          int parameterIndex = position.parameter_index;
+          if (position.onLambda != null) {
+            com.sun.tools.javac.util.List<JCTree.JCVariableDecl> lambdaParams =
+                position.onLambda.params;
+            return parameterIndex < lambdaParams.size()
+                && lambdaParams.get(parameterIndex).sym.equals(sym);
+          } else {
+            return ((Symbol.MethodSymbol) sym.owner).getParameters().indexOf(sym) == parameterIndex;
+          }
+        } else {
+          return false;
+        }
+      case CLASS:
+        // There are no type annotations on the top-level type of the class being declared, only
+        // on other types in the signature (e.g. `class Foo extends Bar<@A Baz> {}`).
+        return false;
+      default:
+        throw new AssertionError("unsupported element kind " + sym.getKind() + " symbol " + sym);
     }
   }
 
@@ -488,12 +545,28 @@ public static boolean isArrayElementNullable(Symbol arraySymbol, Config config)
     }
     // For varargs symbols we also consider the elements to be @Nullable if there is a @Nullable
     // declaration annotation on the parameter
+    // NOTE this flag check does not work for the varargs parameter of a method defined in bytecodes
     if ((arraySymbol.flags() & Flags.VARARGS) != 0) {
       return Nullness.hasNullableDeclarationAnnotation(arraySymbol, config);
     }
     return false;
   }
 
+  /**
+   * Checks if the given varargs symbol has a {@code @Nullable} annotation for its elements. Works
+   * for both source and bytecode.
+   *
+   * @param varargsSymbol the symbol of the varargs parameter
+   * @param config NullAway configuration
+   * @return true if the varargs symbol has a {@code @Nullable} annotation for its elements, false
+   *     otherwise
+   */
+  public static boolean nullableVarargsElementsForSourceOrBytecode(
+      Symbol varargsSymbol, Config config) {
+    return isArrayElementNullable(varargsSymbol, config)
+        || Nullness.hasNullableDeclarationAnnotation(varargsSymbol, config);
+  }
+
   /**
    * Checks if the given array symbol has a {@code @NonNull} annotation for its elements.
    *
@@ -503,23 +576,44 @@ public static boolean isArrayElementNullable(Symbol arraySymbol, Config config)
    *     otherwise
    */
   public static boolean isArrayElementNonNull(Symbol arraySymbol, Config config) {
-    for (Attribute.TypeCompound t : arraySymbol.getRawTypeAttributes()) {
-      for (TypeAnnotationPosition.TypePathEntry entry : t.position.location) {
-        if (entry.tag == TypeAnnotationPosition.TypePathEntryKind.ARRAY) {
-          if (Nullness.isNonNullAnnotation(t.type.toString(), config)) {
-            return true;
-          }
-        }
-      }
+    if (getTypeUseAnnotations(arraySymbol, config, /* onlyDirect= */ false)
+        .anyMatch(
+            t -> {
+              for (TypeAnnotationPosition.TypePathEntry entry : t.position.location) {
+                if (entry.tag == TypeAnnotationPosition.TypePathEntryKind.ARRAY) {
+                  if (Nullness.isNonNullAnnotation(t.type.toString(), config)) {
+                    return true;
+                  }
+                }
+              }
+              return false;
+            })) {
+      return true;
     }
     // For varargs symbols we also consider the elements to be @NonNull if there is a @NonNull
     // declaration annotation on the parameter
+    // NOTE this flag check does not work for the varargs parameter of a method defined in bytecodes
     if ((arraySymbol.flags() & Flags.VARARGS) != 0) {
       return Nullness.hasNonNullDeclarationAnnotation(arraySymbol, config);
     }
     return false;
   }
 
+  /**
+   * Checks if the given varargs symbol has a {@code @NonNull} annotation for its elements. Works
+   * for both source and bytecode.
+   *
+   * @param varargsSymbol the symbol of the varargs parameter
+   * @param config NullAway configuration
+   * @return true if the varargs symbol has a {@code @NonNull} annotation for its elements, false
+   *     otherwise
+   */
+  public static boolean nonnullVarargsElementsForSourceOrBytecode(
+      Symbol varargsSymbol, Config config) {
+    return isArrayElementNonNull(varargsSymbol, config)
+        || Nullness.hasNonNullDeclarationAnnotation(varargsSymbol, config);
+  }
+
   /**
    * Does the given symbol have a JetBrains @NotNull declaration annotation? Useful for workarounds
    * in light of https://github.com/uber/NullAway/issues/720
diff --git a/nullaway/src/main/java/com/uber/nullaway/Nullness.java b/nullaway/src/main/java/com/uber/nullaway/Nullness.java
@@ -227,9 +227,8 @@ public static boolean paramHasNullableAnnotation(
     if (symbol.isVarArgs()
         && paramInd == symbol.getParameters().size() - 1
         && !config.isLegacyAnnotationLocation()) {
-      // individual arguments passed in the varargs positions can be @Nullable if the array element
-      // type of the parameter is @Nullable
-      return NullabilityUtil.isArrayElementNullable(symbol.getParameters().get(paramInd), config);
+      return NullabilityUtil.nullableVarargsElementsForSourceOrBytecode(
+          symbol.getParameters().get(paramInd), config);
     } else {
       return hasNullableAnnotation(
           NullabilityUtil.getAllAnnotationsForParameter(symbol, paramInd, config), config);
@@ -269,9 +268,8 @@ public static boolean paramHasNonNullAnnotation(
     if (symbol.isVarArgs()
         && paramInd == symbol.getParameters().size() - 1
         && !config.isLegacyAnnotationLocation()) {
-      // individual arguments passed in the varargs positions must be @NonNull if the array element
-      // type of the parameter is @NonNull
-      return NullabilityUtil.isArrayElementNonNull(symbol.getParameters().get(paramInd), config);
+      return NullabilityUtil.nonnullVarargsElementsForSourceOrBytecode(
+          symbol.getParameters().get(paramInd), config);
     } else {
       return hasNonNullAnnotation(
           NullabilityUtil.getAllAnnotationsForParameter(symbol, paramInd, config), config);
diff --git a/nullaway/src/main/java/com/uber/nullaway/dataflow/CoreNullnessStoreInitializer.java b/nullaway/src/main/java/com/uber/nullaway/dataflow/CoreNullnessStoreInitializer.java
@@ -67,6 +67,8 @@ private static NullnessStore methodInitialStore(
     for (LocalVariableNode param : parameters) {
       Symbol paramSymbol = (Symbol) param.getElement();
       Nullness assumed;
+      // Using this flag to check for a varargs parameter works since we know paramSymbol represents
+      // a parameter defined in source code
       if ((paramSymbol.flags() & Flags.VARARGS) != 0) {
         assumed = Nullness.varargsArrayIsNullable(paramSymbol, config) ? NULLABLE : NONNULL;
       } else {
diff --git a/nullaway/src/test/java/com/uber/nullaway/Java8Tests.java b/nullaway/src/test/java/com/uber/nullaway/Java8Tests.java
@@ -47,4 +47,25 @@ public void methodReferenceOnNullableVariable() {
             "}")
         .doTest();
   }
+
+  /** test that we can properly read an explicit type-use annotation on a lambda parameter */
+  @Test
+  public void testNullableLambdaParamTypeUse() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "class Test {",
+            "    @FunctionalInterface",
+            "    interface NullableParamFunctionTypeUse<T, U> {",
+            "      U takeVal(@Nullable T x);",
+            "    }",
+            "    static void testParamTypeUse() {",
+            "      NullableParamFunctionTypeUse n3 = (@Nullable Object x) -> (x == null) ? \"null\" : x.toString();",
+            "      NullableParamFunctionTypeUse n4 = (x) -> (x == null) ? \"null\" : x.toString();",
+            "    }",
+            "}")
+        .doTest();
+  }
 }
diff --git a/nullaway/src/test/java/com/uber/nullaway/VarargsTests.java b/nullaway/src/test/java/com/uber/nullaway/VarargsTests.java
@@ -71,6 +71,41 @@ public void testNullableVarargs() {
         .doTest();
   }
 
+  /** Test for a @Nullable declaration annotation on a varargs parameter defined in bytecode */
+  @Test
+  public void nullableDeclarationVarArgsFromBytecode() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import com.uber.lib.Varargs;",
+            "public class Test {",
+            "  public void testDeclaration() {",
+            "    String x = null;",
+            "    Varargs s = new Varargs(x);",
+            "    String y = \"hello\", z = null;",
+            "    Varargs s2 = new Varargs(y, z);",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void nullableTypeUseVarArgsFromBytecode() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import com.uber.lib.Varargs;",
+            "public class Test {",
+            "  public void testTypeUse() {",
+            "    String[] x = null;",
+            "    Varargs.typeUse(x);",
+            "  }",
+            "}")
+        .doTest();
+  }
+
   @Test
   public void nullableTypeUseVarargs() {
     defaultCompilationHelper
@@ -515,4 +550,37 @@ public void testVarargsRestrictive() {
             "}")
         .doTest();
   }
+
+  /**
+   * Test for a restrictive @NonNull declaration annotation on a varargs parameter defined in
+   * bytecode
+   */
+  @Test
+  public void testVarargsRestrictiveBytecodes() {
+    makeTestHelperWithArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber",
+                "-XepOpt:NullAway:UnannotatedSubPackages=com.uber.lib.unannotated",
+                "-XepOpt:NullAway:AcknowledgeRestrictiveAnnotations=true"))
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import com.uber.lib.unannotated.RestrictivelyAnnotatedVarargs;",
+            "public class Test {",
+            "  public void testDeclaration() {",
+            "    String x = null;",
+            "    String[] y = null;",
+            "    // BUG: Diagnostic contains: passing @Nullable parameter 'x'",
+            "    RestrictivelyAnnotatedVarargs.test(x);",
+            "    RestrictivelyAnnotatedVarargs.test(y);",
+            "    // BUG: Diagnostic contains: passing @Nullable parameter 'x'",
+            "    RestrictivelyAnnotatedVarargs.testTypeUse(x);",
+            "    // TODO should report an error; requires a fuller fix for #1027",
+            "    RestrictivelyAnnotatedVarargs.testTypeUse(y);",
+            "  }",
+            "}")
+        .doTest();
+  }
 }
diff --git a/test-java-lib/build.gradle b/test-java-lib/build.gradle
@@ -22,6 +22,7 @@ plugins {
 
 dependencies {
     annotationProcessor project(":nullaway")
+    implementation deps.test.jsr305Annotations
     implementation deps.build.jspecify
 
     compileOnly deps.build.javaxValidation
diff --git a/test-java-lib/src/main/java/com/uber/lib/Varargs.java b/test-java-lib/src/main/java/com/uber/lib/Varargs.java
@@ -0,0 +1,10 @@
+package com.uber.lib;
+
+import javax.annotation.Nullable;
+
+public class Varargs {
+
+  public Varargs(@Nullable String... args) {}
+
+  public static void typeUse(String @org.jspecify.annotations.Nullable ... args) {}
+}
diff --git a/test-java-lib/src/main/java/com/uber/lib/unannotated/RestrictivelyAnnotatedVarargs.java b/test-java-lib/src/main/java/com/uber/lib/unannotated/RestrictivelyAnnotatedVarargs.java
@@ -0,0 +1,11 @@
+package com.uber.lib.unannotated;
+
+import javax.annotation.Nonnull;
+
+public class RestrictivelyAnnotatedVarargs {
+
+  public static void test(@Nonnull String... args) {}
+
+  public static void testTypeUse(
+      @org.jspecify.annotations.NonNull String @org.jspecify.annotations.NonNull ... args) {}
+}
