Fix JarInfer integration test on Java 11 (#529)

Switch to using a service loader to load stubx files from jars on the classpath, to avoid doing our own classpath scanning (which is a pain to get working on Java 9+).

7 files changed, 61 insertions, 110 deletions

diff --git a/.github/workflows/continuous-integration.yml b/.github/workflows/continuous-integration.yml
index 34f291a..552088d 100644
--- a/.github/workflows/continuous-integration.yml
+++ b/.github/workflows/continuous-integration.yml
@@ -54,7 +54,7 @@ jobs:
       - name: Build and test using Gradle and Java 11
         uses: eskatos/gradle-command-action@v1
         with:
-          arguments: build -x :sample-app:build -x :jar-infer:nullaway-integration-test:build
+          arguments: build -x :sample-app:build
         if: matrix.java == '11'
       - name: Build and test using Gradle and Java 17
         uses: eskatos/gradle-command-action@v1
diff --git a/code-coverage-report/build.gradle b/code-coverage-report/build.gradle
index acd5a74..0862c69 100644
--- a/code-coverage-report/build.gradle
+++ b/code-coverage-report/build.gradle
@@ -83,8 +83,6 @@ coverallsTask.configure {
 dependencies {
     implementation project(':nullaway')
     implementation project(':jar-infer:jar-infer-lib')
-    // We cannot include the JarInfer NullAway integration test yet since it does not run on JDK 11.
-    // To include it we would have to also add dependencies on the code being tested
-    //implementation project(':jar-infer:nullaway-integration-test')
+    implementation project(':jar-infer:nullaway-integration-test')
     implementation project(':jdk17-unit-tests')
 }
diff --git a/jar-infer/nullaway-integration-test/build.gradle b/jar-infer/nullaway-integration-test/build.gradle
index c69a6c8..f47fe58 100644
--- a/jar-infer/nullaway-integration-test/build.gradle
+++ b/jar-infer/nullaway-integration-test/build.gradle
@@ -15,6 +15,7 @@
  */
 plugins {
   id "java-library"
+  id "nullaway.jacoco-conventions"
 }
 
 sourceCompatibility = "1.8"
diff --git a/jar-infer/test-java-lib-jarinfer/build.gradle b/jar-infer/test-java-lib-jarinfer/build.gradle
index 0ed2865..c38c688 100644
--- a/jar-infer/test-java-lib-jarinfer/build.gradle
+++ b/jar-infer/test-java-lib-jarinfer/build.gradle
@@ -23,7 +23,7 @@ evaluationDependsOn(":jar-infer:jar-infer-cli")
 sourceCompatibility = "1.8"
 targetCompatibility = "1.8"
 
-def astubxPath = "META-INF/nullaway/jarinfer.astubx"
+def astubxPath = "com/uber/nullaway/jarinfer/provider/jarinfer.astubx"
 
 jar {
     manifest {
@@ -48,6 +48,9 @@ jar.doLast {
 }
 
 dependencies {
+    compileOnly deps.apt.autoService
+    annotationProcessor deps.apt.autoService
+    compileOnly project(":nullaway")
     implementation deps.build.jsr305Annotations
 }
 
diff --git a/jar-infer/test-java-lib-jarinfer/src/main/java/com/uber/nullaway/jarinfer/provider/TestProvider.java b/jar-infer/test-java-lib-jarinfer/src/main/java/com/uber/nullaway/jarinfer/provider/TestProvider.java
new file mode 100644
index 0000000..712028e
--- /dev/null
+++ b/jar-infer/test-java-lib-jarinfer/src/main/java/com/uber/nullaway/jarinfer/provider/TestProvider.java
@@ -0,0 +1,14 @@
+package com.uber.nullaway.jarinfer.provider;
+
+import com.google.auto.service.AutoService;
+import com.uber.nullaway.jarinfer.JarInferStubxProvider;
+import java.util.Collections;
+import java.util.List;
+
+@AutoService(JarInferStubxProvider.class)
+public class TestProvider implements JarInferStubxProvider {
+  @Override
+  public List<String> pathsToStubxFiles() {
+    return Collections.singletonList("jarinfer.astubx");
+  }
+}
diff --git a/nullaway/src/main/java/com/uber/nullaway/handlers/InferredJARModelsHandler.java b/nullaway/src/main/java/com/uber/nullaway/handlers/InferredJARModelsHandler.java
index 6a2ac38..3e34e3a 100644
--- a/nullaway/src/main/java/com/uber/nullaway/handlers/InferredJARModelsHandler.java
+++ b/nullaway/src/main/java/com/uber/nullaway/handlers/InferredJARModelsHandler.java
@@ -38,20 +38,16 @@ import com.uber.nullaway.Config;
 import com.uber.nullaway.NullAway;
 import com.uber.nullaway.dataflow.AccessPath;
 import com.uber.nullaway.dataflow.AccessPathNullnessPropagation;
+import com.uber.nullaway.jarinfer.JarInferStubxProvider;
 import java.io.DataInputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import java.net.JarURLConnection;
-import java.net.URL;
-import java.net.URLClassLoader;
-import java.net.URLConnection;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.ServiceLoader;
 import java.util.Set;
-import java.util.jar.JarEntry;
-import java.util.jar.JarFile;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.type.TypeKind;
 import org.checkerframework.nullaway.dataflow.cfg.node.MethodInvocationNode;
@@ -68,7 +64,6 @@ public class InferredJARModelsHandler extends BaseNoOpHandler {
   }
 
   private static final int VERSION_0_FILE_MAGIC_NUMBER = 691458791;
-  private static final String DEFAULT_ASTUBX_LOCATION = "META-INF/nullaway/jarinfer.astubx";
   private static final String ANDROID_ASTUBX_LOCATION = "jarinfer.astubx";
   private static final String ANDROID_MODEL_CLASS =
       "com.uber.nullaway.jarinfer.AndroidJarInferModels";
@@ -76,8 +71,6 @@ public class InferredJARModelsHandler extends BaseNoOpHandler {
   private static final int RETURN = -1; // '-1' indexes Return type in the Annotation Cache
 
   private final Map<String, Map<String, Map<Integer, Set<String>>>> argAnnotCache;
-  private final Map<String, Set<String>> mapModelJarLocations;
-  private final Set<String> loadedJars;
 
   private final Config config;
 
@@ -85,8 +78,7 @@ public class InferredJARModelsHandler extends BaseNoOpHandler {
     super();
     this.config = config;
     argAnnotCache = new LinkedHashMap<>();
-    mapModelJarLocations = new LinkedHashMap<>();
-    loadedJars = new LinkedHashSet<>();
+    loadStubxFiles();
     // Load Android SDK JarInfer models
     try {
       InputStream androidStubxIS =
@@ -109,21 +101,25 @@ public class InferredJARModelsHandler extends BaseNoOpHandler {
     }
   }
 
-  /*
-   * Scan Java Classpath for JarInfer model jars and map their names to locations
+  /**
+   * Loads all stubx files discovered in the classpath. Stubx files are discovered via
+   * implementations of {@link JarInferStubxProvider} loaded using a {@link ServiceLoader}
    */
-  private void processClassPath() {
-    URL[] classLoaderUrls =
-        ((URLClassLoader) Thread.currentThread().getContextClassLoader()).getURLs();
-    for (URL url : classLoaderUrls) {
-      String path = url.getFile();
-      if (path.matches(config.getJarInferRegexStripModelJarName())) {
-        String name = path.replaceAll(config.getJarInferRegexStripModelJarName(), "$1");
-        LOG(DEBUG, "DEBUG", "model jar name: " + name + "\tjar path: " + path);
-        if (!mapModelJarLocations.containsKey(name)) {
-          mapModelJarLocations.put(name, new LinkedHashSet<>());
+  private void loadStubxFiles() {
+    Iterable<JarInferStubxProvider> astubxProviders =
+        ServiceLoader.load(
+            JarInferStubxProvider.class, InferredJARModelsHandler.class.getClassLoader());
+    for (JarInferStubxProvider provider : astubxProviders) {
+      for (String astubxPath : provider.pathsToStubxFiles()) {
+        Class<? extends JarInferStubxProvider> providerClass = provider.getClass();
+        InputStream stubxInputStream = providerClass.getResourceAsStream(astubxPath);
+        String stubxLocation = providerClass + ":" + astubxPath;
+        try {
+          parseStubStream(stubxInputStream, stubxLocation);
+          LOG(DEBUG, "DEBUG", "loaded stubx file " + stubxLocation);
+        } catch (IOException e) {
+          throw new RuntimeException("could not parse stubx file " + stubxLocation, e);
         }
-        mapModelJarLocations.get(name).add(path);
       }
     }
   }
@@ -135,9 +131,6 @@ public class InferredJARModelsHandler extends BaseNoOpHandler {
       Symbol.MethodSymbol methodSymbol,
       List<? extends ExpressionTree> actualParams,
       ImmutableSet<Integer> nonNullPositions) {
-    if (mapModelJarLocations.isEmpty()) {
-      processClassPath();
-    }
     Symbol.ClassSymbol classSymbol = methodSymbol.enclClass();
     String className = classSymbol.getQualifiedName().toString();
     if (methodSymbol.getModifiers().contains(Modifier.ABSTRACT)) {
@@ -147,7 +140,7 @@ public class InferredJARModelsHandler extends BaseNoOpHandler {
           "Skipping abstract method: " + className + " : " + methodSymbol.getQualifiedName());
       return nonNullPositions;
     }
-    if (!lookupAndBuildCache(classSymbol)) {
+    if (!argAnnotCache.containsKey(className)) {
       return nonNullPositions;
     }
     String methodSign = getMethodSignature(methodSymbol);
@@ -200,7 +193,7 @@ public class InferredJARModelsHandler extends BaseNoOpHandler {
       Preconditions.checkNotNull(methodSymbol);
       Symbol.ClassSymbol classSymbol = methodSymbol.enclClass();
       String className = classSymbol.getQualifiedName().toString();
-      if (lookupAndBuildCache(classSymbol)) {
+      if (argAnnotCache.containsKey(className)) {
         String methodSign = getMethodSignature(methodSymbol);
         Map<Integer, Set<String>> methodArgAnnotations = lookupMethodInCache(className, methodSign);
         if (methodArgAnnotations != null) {
@@ -217,82 +210,6 @@ public class InferredJARModelsHandler extends BaseNoOpHandler {
     return false;
   }
 
-  private boolean lookupAndBuildCache(Symbol.ClassSymbol klass) {
-    String className = klass.getQualifiedName().toString();
-    try {
-      LOG(DEBUG, "DEBUG", "Looking for class: " + className);
-      if (klass.classfile == null) {
-        LOG(VERBOSE, "Warn", "Cannot resolve source for class: " + className);
-        return false;
-      }
-      // Annotation cache
-      String jarPath = "";
-      if (!argAnnotCache.containsKey(className)) {
-        // this works for aar !
-        URLConnection uc = klass.classfile.toUri().toURL().openConnection();
-        if (!(uc instanceof JarURLConnection)) {
-          return false;
-        }
-        JarURLConnection juc = (JarURLConnection) uc;
-        jarPath = juc.getJarFileURL().getPath();
-        LOG(DEBUG, "DEBUG", "Found source of class: " + className + ", jar: " + jarPath);
-        // Avoid reloading for classes w/o any stubs from already loaded jars.
-        if (!loadedJars.contains(jarPath)) {
-          loadedJars.add(jarPath);
-          String jarName = jarPath.replaceAll(config.getJarInferRegexStripCodeJarName(), "$1");
-          LOG(DEBUG, "DEBUG", "code jar name: " + jarName + "\tjar path: " + jarPath);
-          // Lookup model jar locations for jar name
-          if (!mapModelJarLocations.containsKey(jarName)) {
-            LOG(
-                VERBOSE,
-                "Warn",
-                "Cannot find model jar for class: " + className + ", jar: " + jarName);
-            return false;
-          }
-          // Load model jars
-          for (String modelJarPath : mapModelJarLocations.get(jarName)) {
-            JarFile jar = new JarFile(modelJarPath);
-            LOG(DEBUG, "DEBUG", "Found model jar at: " + modelJarPath);
-            JarEntry astubxJE = jar.getJarEntry(DEFAULT_ASTUBX_LOCATION);
-            if (astubxJE == null) {
-              LOG(VERBOSE, "Warn", "Cannot find jarinfer.astubx in jar: " + modelJarPath);
-              return false;
-            }
-            InputStream astubxIS = jar.getInputStream(astubxJE);
-            if (astubxIS == null) {
-              LOG(VERBOSE, "Warn", "Cannot load jarinfer.astubx in jar: " + modelJarPath);
-              return false;
-            }
-            parseStubStream(astubxIS, modelJarPath + ": " + DEFAULT_ASTUBX_LOCATION);
-            LOG(
-                DEBUG,
-                "DEBUG",
-                "Loaded "
-                    + argAnnotCache.keySet().size()
-                    + " astubx for class: "
-                    + className
-                    + " from jar: "
-                    + modelJarPath);
-          }
-        } else {
-          LOG(DEBUG, "DEBUG", "Skipping already loaded jar: " + jarPath);
-        }
-      } else {
-        LOG(DEBUG, "DEBUG", "Hit annotation cache for class: " + className);
-      }
-      if (!argAnnotCache.containsKey(className)) {
-        LOG(
-            VERBOSE,
-            "Warn",
-            "Cannot find Annotation Cache for class: " + className + ", jar: " + jarPath);
-        return false;
-      }
-    } catch (IOException e) {
-      throw new Error(e);
-    }
-    return true;
-  }
-
   private Map<Integer, Set<String>> lookupMethodInCache(String className, String methodSign) {
     if (!argAnnotCache.containsKey(className)) {
       return null;
diff --git a/nullaway/src/main/java/com/uber/nullaway/jarinfer/JarInferStubxProvider.java b/nullaway/src/main/java/com/uber/nullaway/jarinfer/JarInferStubxProvider.java
new file mode 100644
index 0000000..b4a4b02
--- /dev/null
+++ b/nullaway/src/main/java/com/uber/nullaway/jarinfer/JarInferStubxProvider.java
@@ -0,0 +1,18 @@
+package com.uber.nullaway.jarinfer;
+
+import java.util.List;
+
+/** Interface to be implemented by providers of stubx files to be read by NullAway */
+public interface JarInferStubxProvider {
+
+  /**
+   * Returns a list of paths to stubx files to be loaded. Paths can be relative to the location of
+   * the implementing provider class (recommended), or absolute (starting with {@code /}).
+   *
+   * <p><b>NOTE:</b> NullAway does <em>not</em> have any special handling for cases where there are
+   * multiple jars in its classpath containing a stubx file at the same path. So, providers are
+   * strongly encouraged to use paths to stubx files that are unlikely to clash with those in other
+   * jars.
+   */
+  List<String> pathsToStubxFiles();
+}