Require single main class

Update run tasks to ensure that only a single main class is required
when performing a class search.

See gh-886

spring-boot-tools/spring-boot-gradle-plugin/src/main/groovy/org/springframework/boot/gradle/task/ComputeMain.java

@@ -29,7 +29,7 @@ import org.springframework.boot.loader.tools.MainClassFinder;
 
 /**
  * Add a main class if one is missing from the build
- * 
+ *
  * @author Dave Syer
  */
 public class ComputeMain implements Action<Task> {
@@ -68,7 +68,7 @@ public class ComputeMain implements Action<Task> {
 		project.getLogger().debug(
 				"Looking for main in: " + main.getOutput().getClassesDir());
 		try {
-			String mainClass = MainClassFinder.findMainClass(main.getOutput()
+			String mainClass = MainClassFinder.findSingleMainClass(main.getOutput()
 					.getClassesDir());
 			project.getLogger().info("Computed main class: " + mainClass);
 			return mainClass;

spring-boot-tools/spring-boot-gradle-plugin/src/main/groovy/org/springframework/boot/gradle/task/RunApp.java

@@ -91,7 +91,7 @@ public class RunApp extends DefaultTask {
 		}
 		getLogger().info("Looking for main in: " + main.getOutput().getClassesDir());
 		try {
-			return MainClassFinder.findMainClass(main.getOutput().getClassesDir());
+			return MainClassFinder.findSingleMainClass(main.getOutput().getClassesDir());
 		}
 		catch (IOException ex) {
 			throw new IllegalStateException("Cannot find main class", ex);

spring-boot-tools/spring-boot-loader-tools/src/main/java/org/springframework/boot/loader/tools/MainClassFinder.java

@@ -29,7 +29,9 @@ import java.util.Collections;
 import java.util.Comparator;
 import java.util.Deque;
 import java.util.Enumeration;
+import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Set;
 import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
 
@@ -85,6 +87,18 @@ public abstract class MainClassFinder {
 		});
 	}
 
+	/**
+	 * Find a single main class from a given folder.
+	 * @param rootFolder the root folder to search
+	 * @return the main class or {@code null}
+	 * @throws IOException
+	 */
+	public static String findSingleMainClass(File rootFolder) throws IOException {
+		MainClassesCallback callback = new MainClassesCallback();
+		MainClassFinder.doWithMainClasses(rootFolder, callback);
+		return callback.getMainClass();
+	}
+
 	/**
 	 * Perform the given callback operation on all main classes from the given root
 	 * folder.
@@ -93,7 +107,7 @@ public abstract class MainClassFinder {
 	 * @return the first callback result or {@code null}
 	 * @throws IOException
 	 */
-	public static <T> T doWithMainClasses(File rootFolder, ClassNameCallback<T> callback)
+	static <T> T doWithMainClasses(File rootFolder, ClassNameCallback<T> callback)
 			throws IOException {
 		if (!rootFolder.exists()) {
 			return null; // nothing to do
@@ -160,6 +174,20 @@ public abstract class MainClassFinder {
 				});
 	}
 
+	/**
+	 * Find a single main class in a given jar file.
+	 * @param jarFile the jar file to search
+	 * @param classesLocation the location within the jar containing classes
+	 * @return the main class or {@code null}
+	 * @throws IOException
+	 */
+	public static String findSingleMainClass(JarFile jarFile, String classesLocation)
+			throws IOException {
+		MainClassesCallback callback = new MainClassesCallback();
+		MainClassFinder.doWithMainClasses(jarFile, classesLocation, callback);
+		return callback.getMainClass();
+	}
+
 	/**
 	 * Perform the given callback operation on all main classes from the given jar.
 	 * @param jarFile the jar file to search
@@ -167,7 +195,7 @@ public abstract class MainClassFinder {
 	 * @return the first callback result or {@code null}
 	 * @throws IOException
 	 */
-	public static <T> T doWithMainClasses(JarFile jarFile, String classesLocation,
+	static <T> T doWithMainClasses(JarFile jarFile, String classesLocation,
 			ClassNameCallback<T> callback) throws IOException {
 		List<JarEntry> classEntries = getClassEntries(jarFile, classesLocation);
 		Collections.sort(classEntries, new ClassEntryComparator());
@@ -293,4 +321,30 @@ public abstract class MainClassFinder {
 		T doWith(String className);
 
 	}
+
+	/**
+	 * Find a single main class, throwing an {@link IllegalStateException} if multiple
+	 * candidates exist.
+	 */
+	private static class MainClassesCallback implements ClassNameCallback<Object> {
+
+		private final Set<String> classNames = new LinkedHashSet<String>();
+
+		@Override
+		public Object doWith(String className) {
+			this.classNames.add(className);
+			return null;
+		}
+
+		public String getMainClass() {
+			if (this.classNames.size() > 1) {
+				throw new IllegalStateException(
+						"Unable to find a single main class from the following candidates "
+								+ this.classNames);
+			}
+			return this.classNames.isEmpty() ? null : this.classNames.iterator().next();
+		}
+
+	}
+
 }

spring-boot-tools/spring-boot-loader-tools/src/main/java/org/springframework/boot/loader/tools/Repackager.java

@@ -20,13 +20,9 @@ import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import java.util.ArrayList;
-import java.util.List;
 import java.util.jar.JarFile;
 import java.util.jar.Manifest;
 
-import org.springframework.boot.loader.tools.MainClassFinder.ClassNameCallback;
-
 /**
  * Utility class that can be used to repackage an archive so that it can be executed using
  * '{@literal java -jar}'.
@@ -228,10 +224,8 @@ public class Repackager {
 	}
 
 	protected String findMainMethod(JarFile source) throws IOException {
-		MainClassesCallback callback = new MainClassesCallback();
-		MainClassFinder.doWithMainClasses(source, this.layout.getClassesLocation(),
-				callback);
-		return callback.getMainClass();
+		return MainClassFinder.findSingleMainClass(source,
+				this.layout.getClassesLocation());
 	}
 
 	private void renameFile(File file, File dest) {
@@ -247,24 +241,4 @@ public class Repackager {
 		}
 	}
 
-	private static class MainClassesCallback implements ClassNameCallback<Object> {
-
-		private final List<String> classNames = new ArrayList<String>();
-
-		@Override
-		public Object doWith(String className) {
-			this.classNames.add(className);
-			return null;
-		}
-
-		public String getMainClass() {
-			if (this.classNames.size() > 1) {
-				throw new IllegalStateException(
-						"Unable to find a single main class from the following candidates "
-								+ this.classNames);
-			}
-			return this.classNames.isEmpty() ? null : this.classNames.get(0);
-		}
-
-	}
 }

spring-boot-tools/spring-boot-loader-tools/src/test/java/org/springframework/boot/loader/tools/MainClassFinderTests.java

@@ -23,6 +23,7 @@ import java.util.List;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 import org.junit.rules.TemporaryFolder;
 import org.springframework.boot.loader.tools.MainClassFinder.ClassNameCallback;
 import org.springframework.boot.loader.tools.sample.ClassWithMainMethod;
@@ -41,6 +42,9 @@ public class MainClassFinderTests {
 	@Rule
 	public TemporaryFolder temporaryFolder = new TemporaryFolder();
 
+	@Rule
+	public ExpectedException thrown = ExpectedException.none();
+
 	private TestJarFile testJarFile;
 
 	@Before
@@ -73,6 +77,16 @@ public class MainClassFinderTests {
 		assertThat(actual, equalTo("a.B"));
 	}
 
+	@Test
+	public void findSingleJarSearch() throws Exception {
+		this.testJarFile.addClass("a/B.class", ClassWithMainMethod.class);
+		this.testJarFile.addClass("a/b/c/E.class", ClassWithMainMethod.class);
+		this.thrown.expect(IllegalStateException.class);
+		this.thrown.expectMessage("Unable to find a single main class "
+				+ "from the following candidates [a.B, a.b.c.E]");
+		MainClassFinder.findSingleMainClass(this.testJarFile.getJarFile(), "");
+	}
+
 	@Test
 	public void findMainClassInJarSubLocation() throws Exception {
 		this.testJarFile.addClass("a/B.class", ClassWithMainMethod.class);
@@ -108,6 +122,16 @@ public class MainClassFinderTests {
 		assertThat(actual, equalTo("a.B"));
 	}
 
+	@Test
+	public void findSingleFolderSearch() throws Exception {
+		this.testJarFile.addClass("a/B.class", ClassWithMainMethod.class);
+		this.testJarFile.addClass("a/b/c/E.class", ClassWithMainMethod.class);
+		this.thrown.expect(IllegalStateException.class);
+		this.thrown.expectMessage("Unable to find a single main class "
+				+ "from the following candidates [a.B, a.b.c.E]");
+		MainClassFinder.findSingleMainClass(this.testJarFile.getJarSource());
+	}
+
 	@Test
 	public void doWithFolderMainMethods() throws Exception {
 		this.testJarFile.addClass("a/b/c/D.class", ClassWithMainMethod.class);

spring-boot-tools/spring-boot-maven-plugin/src/main/java/org/springframework/boot/maven/RunMojo.java

@@ -217,7 +217,7 @@ public class RunMojo extends AbstractDependencyFilterMojo {
 		String mainClass = this.mainClass;
 		if (mainClass == null) {
 			try {
-				mainClass = MainClassFinder.findMainClass(this.classesDirectory);
+				mainClass = MainClassFinder.findSingleMainClass(this.classesDirectory);
 			}
 			catch (IOException ex) {
 				throw new MojoExecutionException(ex.getMessage(), ex);