Fixes main class finding with external dependencies. (#355)

jib-core/build.gradle

@@ -38,6 +38,7 @@ dependencies {
     compile 'com.google.guava:guava:23.5-jre'
     compile 'com.fasterxml.jackson.core:jackson-databind:2.9.2'
     compile 'org.slf4j:slf4j-api:1.7.25'
+    compile 'org.javassist:javassist:3.22.0-GA'
 
     testCompile 'junit:junit:4.12'
     testCompile 'org.mockito:mockito-core:2.12.0'

jib-core/src/main/java/com/google/cloud/tools/jib/frontend/MainClassFinder.java

@@ -21,87 +21,25 @@
 import com.google.cloud.tools.jib.filesystem.DirectoryWalker;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
-import com.google.common.base.Splitter;
 import java.io.IOException;
-import java.lang.reflect.Method;
+import java.io.InputStream;
 import java.lang.reflect.Modifier;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.util.ArrayDeque;
 import java.util.ArrayList;
-import java.util.Deque;
-import java.util.HashMap;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
+import javassist.ClassPool;
+import javassist.CtClass;
+import javassist.CtMethod;
+import javassist.NotFoundException;
 import javax.annotation.Nullable;
 
 /** Infers the main class in an application. */
 public class MainClassFinder {
 
-  /** Helper for loading a .class file. */
-  @VisibleForTesting
-  static class ClassFileLoader extends ClassLoader {
-
-    private final Map<Path, Class<?>> definedClasses = new HashMap<>();
-    private final Path rootDirectory;
-
-    @VisibleForTesting
-    ClassFileLoader(Path rootDirectory) {
-      this.rootDirectory = rootDirectory;
-    }
-
-    @Nullable
-    @Override
-    public Class<?> findClass(String name) {
-      return findClass(name, getPathFromClassName(name));
-    }
-
-    /**
-     * @param name the name of the class. This should be {@code null} when we call it manually with
-     *     a file, and not {@code null} when it is called internally.
-     * @param file the .class file defining the class.
-     * @return the {@link Class} defined by the file, or {@code null} if the class could not be
-     *     defined.
-     */
-    @Nullable
-    private Class<?> findClass(@Nullable String name, Path file) {
-      if (definedClasses.containsKey(file)) {
-        return definedClasses.get(file);
-      }
-
-      if (!Files.exists(file)) {
-        // TODO: Log search class failure?
-        return null;
-      }
-
-      try {
-        byte[] bytes = Files.readAllBytes(file);
-        Class<?> definedClass = defineClass(name, bytes, 0, bytes.length);
-        definedClasses.put(file, definedClass);
-        return definedClass;
-      } catch (IOException | ClassFormatError | SecurityException | NoClassDefFoundError ignored) {
-        // Not a valid class file
-        // TODO: Log search class failure when NoClassDefFoundError/SecurityException is caught?
-        return null;
-      }
-    }
-
-    /** Converts a class name (pack.ClassName) to a Path (rootDirectory/pack/ClassName.class). */
-    @VisibleForTesting
-    Path getPathFromClassName(String className) {
-      Path path = rootDirectory;
-      Deque<String> folders = new ArrayDeque<>(Splitter.on('.').splitToList(className));
-      String fileName = folders.removeLast() + ".class";
-      for (String folder : folders) {
-        path = path.resolve(folder);
-      }
-      path = path.resolve(fileName);
-      return path;
-    }
-  }
-
   /**
    * If {@code mainClass} is {@code null}, tries to infer main class in this order:
    *
@@ -139,7 +77,7 @@ public static String resolveMainClass(
               continue;
             }
             visitedRoots.add(root);
-            mainClasses.addAll(findMainClasses(root));
+            mainClasses.addAll(findMainClasses(root, logger));
           }
 
           if (mainClasses.size() == 1) {
@@ -177,47 +115,58 @@ public static String resolveMainClass(
   }
 
   /**
-   * Searches for a .class file containing a main method in a root directory.
+   * Finds the classes with {@code public static void main(String[] args)} in {@code rootDirectory}.
    *
-   * @return the name of the class if one is found, null if no class is found.
-   * @throws IOException if searching/reading files fails.
+   * @param rootDirectory directory containing the {@code .class} files
    */
   @VisibleForTesting
-  static List<String> findMainClasses(Path rootDirectory) throws IOException {
-    List<String> classNames = new ArrayList<>();
-
-    // Make sure rootDirectory is valid
+  static List<String> findMainClasses(Path rootDirectory, BuildLogger buildLogger)
+      throws IOException {
+    // Makes sure rootDirectory is valid.
     if (!Files.exists(rootDirectory) || !Files.isDirectory(rootDirectory)) {
-      return classNames;
+      return Collections.emptyList();
     }
 
-    // Get all .class files
-    ClassFileLoader classFileLoader = new ClassFileLoader(rootDirectory);
-    new DirectoryWalker(rootDirectory)
-        .filter(Files::isRegularFile)
-        .filter(path -> path.toString().endsWith(".class"))
-        .walk(
-            classFile -> {
-              Class<?> fileClass = classFileLoader.findClass(null, classFile);
-              if (fileClass == null) {
-                return;
-              }
-              try {
-                // Check if class contains {@code public static void main(String[] args)}
-                Method main = fileClass.getMethod("main", String[].class);
-                if (main != null
-                    && main.getReturnType() == void.class
-                    && Modifier.isStatic(main.getModifiers())
-                    && Modifier.isPublic(main.getModifiers())) {
-                  classNames.add(fileClass.getName());
+    List<String> classNames = new ArrayList<>();
+
+    ClassPool classPool = new ClassPool();
+    classPool.appendSystemPath();
+
+    try {
+      CtClass[] mainMethodParams = new CtClass[] {classPool.get("java.lang.String[]")};
+
+      new DirectoryWalker(rootDirectory)
+          .filter(Files::isRegularFile)
+          .filter(path -> path.toString().endsWith(".class"))
+          .walk(
+              classFile -> {
+                try (InputStream classFileInputStream = Files.newInputStream(classFile)) {
+                  CtClass fileClass = classPool.makeClass(classFileInputStream);
+
+                  // Check if class contains 'public static void main(String[] args)'.
+                  CtMethod mainMethod = fileClass.getDeclaredMethod("main", mainMethodParams);
+
+                  if (CtClass.voidType.equals(mainMethod.getReturnType())
+                      && Modifier.isStatic(mainMethod.getModifiers())
+                      && Modifier.isPublic(mainMethod.getModifiers())) {
+                    classNames.add(fileClass.getName());
+                  }
+
+                } catch (NotFoundException ex) {
+                  // Ignores main method not found.
+
+                } catch (IOException ex) {
+                  // Could not read class file.
+                  buildLogger.warn("Could not read class file: " + classFile);
                 }
-              } catch (NoSuchMethodException | NoClassDefFoundError ignored) {
-                // main method not found
-                // TODO: Log search class failure when NoClassDefFoundError is caught?
-              }
-            });
+              });
+
+      return classNames;
 
-    return classNames;
+    } catch (NotFoundException ex) {
+      // Thrown if 'java.lang.String' is not found in classPool.
+      throw new RuntimeException(ex);
+    }
   }
 
   private MainClassFinder() {}

jib-core/src/test/java/com/google/cloud/tools/jib/frontend/MainClassFinderTest.java

@@ -59,7 +59,7 @@ public void setup() {
   @Test
   public void testFindMainClass_simple() throws URISyntaxException, IOException {
     Path rootDirectory = Paths.get(Resources.getResource("class-finder-tests/simple").toURI());
-    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory);
+    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory, mockBuildLogger);
     Assert.assertEquals(1, mainClasses.size());
     Assert.assertTrue(mainClasses.contains("HelloWorld"));
   }
@@ -68,22 +68,22 @@ public void testFindMainClass_simple() throws URISyntaxException, IOException {
   public void testFindMainClass_subdirectories() throws URISyntaxException, IOException {
     Path rootDirectory =
         Paths.get(Resources.getResource("class-finder-tests/subdirectories").toURI());
-    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory);
+    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory, mockBuildLogger);
     Assert.assertEquals(1, mainClasses.size());
     Assert.assertTrue(mainClasses.contains("multi.layered.HelloWorld"));
   }
 
   @Test
   public void testFindMainClass_noClass() throws URISyntaxException, IOException {
     Path rootDirectory = Paths.get(Resources.getResource("class-finder-tests/no-main").toURI());
-    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory);
+    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory, mockBuildLogger);
     Assert.assertTrue(mainClasses.isEmpty());
   }
 
   @Test
   public void testFindMainClass_multiple() throws URISyntaxException, IOException {
     Path rootDirectory = Paths.get(Resources.getResource("class-finder-tests/multiple").toURI());
-    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory);
+    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory, mockBuildLogger);
     Assert.assertEquals(2, mainClasses.size());
     Assert.assertTrue(mainClasses.contains("multi.layered.HelloMoon"));
     Assert.assertTrue(mainClasses.contains("HelloWorld"));
@@ -92,16 +92,25 @@ public void testFindMainClass_multiple() throws URISyntaxException, IOException
   @Test
   public void testFindMainClass_extension() throws URISyntaxException, IOException {
     Path rootDirectory = Paths.get(Resources.getResource("class-finder-tests/extension").toURI());
-    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory);
+    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory, mockBuildLogger);
     Assert.assertEquals(1, mainClasses.size());
     Assert.assertTrue(mainClasses.contains("main.MainClass"));
   }
 
   @Test
-  public void testFindMainClass_externalmethod() throws URISyntaxException, IOException {
+  public void testFindMainClass_importedMethods() throws URISyntaxException, IOException {
     Path rootDirectory =
-        Paths.get(Resources.getResource("class-finder-tests/externalmethod").toURI());
-    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory);
+        Paths.get(Resources.getResource("class-finder-tests/imported-methods").toURI());
+    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory, mockBuildLogger);
+    Assert.assertEquals(1, mainClasses.size());
+    Assert.assertTrue(mainClasses.contains("main.MainClass"));
+  }
+
+  @Test
+  public void testFindMainClass_externalClasses() throws URISyntaxException, IOException {
+    Path rootDirectory =
+        Paths.get(Resources.getResource("class-finder-tests/external-classes").toURI());
+    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory, mockBuildLogger);
     Assert.assertEquals(1, mainClasses.size());
     Assert.assertTrue(mainClasses.contains("main.MainClass"));
   }
@@ -110,7 +119,7 @@ public void testFindMainClass_externalmethod() throws URISyntaxException, IOExce
   public void testFindMainClass_innerClasses() throws URISyntaxException, IOException {
     Path rootDirectory =
         Paths.get(Resources.getResource("class-finder-tests/inner-classes").toURI());
-    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory);
+    List<String> mainClasses = MainClassFinder.findMainClasses(rootDirectory, mockBuildLogger);
     Assert.assertEquals(1, mainClasses.size());
     Assert.assertTrue(mainClasses.contains("HelloWorld$InnerClass"));
   }
@@ -191,20 +200,4 @@ public void testResolveMainClass_noneInferredWithoutBackup() {
           .getMainClassHelpfulSuggestions("Main class was not found");
     }
   }
-
-  @Test
-  public void testGetPathFromClassName() {
-    Path root = Paths.get("test").resolve("root");
-    MainClassFinder.ClassFileLoader classFileLoader = new MainClassFinder.ClassFileLoader(root);
-    Assert.assertEquals(
-        root.resolve("test").resolve("ClassName.class"),
-        classFileLoader.getPathFromClassName("test.ClassName"));
-  }
-
-  @Test
-  public void testGetPathFromClassName_emptyClassName() {
-    MainClassFinder.ClassFileLoader classFileLoader =
-        new MainClassFinder.ClassFileLoader(Paths.get(""));
-    Assert.assertEquals(Paths.get(".class"), classFileLoader.getPathFromClassName(""));
-  }
 }

jib-gradle-plugin/build.gradle

@@ -61,6 +61,8 @@ dependencies {
     compile 'com.google.guava:guava:23.5-jre'
     compile 'com.fasterxml.jackson.core:jackson-databind:2.9.2'
     compile 'org.slf4j:slf4j-api:1.7.25'
+    compile 'org.javassist:javassist:3.22.0-GA'
+
     testCompile 'junit:junit:4.12'
     testCompile 'org.mockito:mockito-core:2.12.0'
 

jib-maven-plugin/pom.xml

@@ -77,6 +77,12 @@
       <version>1.7.25</version>
       <scope>compile</scope>
     </dependency>
+    <dependency>
+      <groupId>org.javassist</groupId>
+      <artifactId>javassist</artifactId>
+      <version>3.22.0-GA</version>
+      <scope>compile</scope>
+    </dependency>
     <!-- End dependencies from jib-core -->
 
     <dependency>