name.remal.reflection.MainClassUtils Maven / Gradle / Ivy

package name.remal.reflection;

import static java.lang.ClassLoader.getSystemClassLoader;
import static java.util.jar.Attributes.Name.MAIN_CLASS;
import static name.remal.SneakyThrow.sneakyThrow;
import static name.remal.UncheckedCast.uncheckedCast;

import java.io.InputStream;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.URL;
import java.util.Enumeration;
import java.util.jar.Manifest;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

public class MainClassUtils {

    public static void invokeMainMethod(@NotNull ClassLoader classLoader, @NotNull String... args) {
        Class mainClass = getMainClass(classLoader);
        if (mainClass == null) throw new IllegalStateException("Main class can't be found for: " + classLoader);
        try {
            Method mainMethod = mainClass.getMethod("main", String[].class);
            mainMethod.setAccessible(true);
            mainMethod.invoke(null, (Object) args);

        } catch (@NotNull InvocationTargetException e) {
            throw sneakyThrow(e.getTargetException());
        } catch (@NotNull NoSuchMethodException | IllegalAccessException e) {
            throw sneakyThrow(e);
        }
    }

    @Nullable
    public static Class getMainClass(@NotNull ClassLoader classLoader) {
        String mainClassName = getMainClassName(classLoader);
        if (mainClassName != null) {
            try {
                return classLoader.loadClass(mainClassName);
            } catch (ClassNotFoundException e) {
                throw sneakyThrow(e);
            }
        }

        return null;
    }

    private static final String MANIFEST_RESOURCE_NAME = "META-INF/MANIFEST.MF";

    private static final Method findResourcesMethod;

    static {
        try {
            findResourcesMethod = ClassLoader.class.getDeclaredMethod("findResources", String.class);
            findResourcesMethod.setAccessible(true);
        } catch (NoSuchMethodException e) {
            throw sneakyThrow(e);
        }
    }

    @Nullable
    public static String getMainClassName(@NotNull ClassLoader classLoader) {
        try {
            Enumeration manifestURLs = uncheckedCast(findResourcesMethod.invoke(classLoader, MANIFEST_RESOURCE_NAME));
            while (manifestURLs.hasMoreElements()) {
                URL manifestURL = manifestURLs.nextElement();
                try (InputStream inputStream = manifestURL.openStream()) {
                    Manifest manifest = new Manifest(inputStream);
                    String mainClassName = manifest.getMainAttributes().getValue(MAIN_CLASS);
                    if (mainClassName != null) return mainClassName;
                }
            }

            ClassLoader parentClassLoader = classLoader.getParent();
            if (parentClassLoader != null && getSystemClassLoader() != parentClassLoader) {
                return getMainClassName(parentClassLoader);
            }

            return null;

        } catch (Exception e) {
            throw sneakyThrow(e);
        }
    }

}