Fix loading ATs from classpath mods

Author: mezz

src/main/resources/net/minecraftforge/gradle/GradleForgeHacks.java

@@ -19,17 +19,21 @@
  */
 package net.minecraftforge.gradle;
 
+import javax.annotation.Nullable;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.lang.reflect.Field;
+import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.net.URISyntaxException;
 import java.net.URL;
 import java.net.URLClassLoader;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.jar.Attributes;
 import java.util.jar.JarFile;
 import java.util.jar.Manifest;
 
@@ -61,7 +65,7 @@ public class GradleForgeHacks
 
     public static final Map<String, File> coreMap        = Maps.newHashMap();
 
-    public static void searchCoremods(GradleStartCommon common) throws Exception
+    public static void searchCoremods(GradleStartCommon common)
     {
         // check for argument
         if (common.extras.contains(NO_CORE_SEARCH))
@@ -75,69 +79,135 @@ public static void searchCoremods(GradleStartCommon common) throws Exception
             return;
         }
 
-        // intialize AT hack Method
-        Method atRegistrar = null;
-        try
+        // initialize AT hack Method
+        AtRegistrar atRegistrar = new AtRegistrar();
+
+        URLClassLoader urlClassLoader = (URLClassLoader) GradleStartCommon.class.getClassLoader();
+        for (URL url : urlClassLoader.getURLs())
         {
-            atRegistrar = Class.forName(MOD_ATD_CLASS).getDeclaredMethod(MOD_AT_METHOD, JarFile.class);
+            try
+            {
+                searchCoremodAtUrl(url, atRegistrar);
+            }
+            catch (IOException | InvocationTargetException | IllegalAccessException | URISyntaxException e)
+            {
+                GradleStartCommon.LOGGER.warn("GradleForgeHacks failed to search for coremod at url {}", url, e);
+            }
         }
-        catch (Throwable t)
-        {}
 
-        for (URL url : ((URLClassLoader) GradleStartCommon.class.getClassLoader()).getURLs())
+        // set property.
+        Set<String> coremodsSet = Sets.newHashSet();
+        if (!Strings.isNullOrEmpty(System.getProperty(COREMOD_VAR)))
+            coremodsSet.addAll(Splitter.on(',').splitToList(System.getProperty(COREMOD_VAR)));
+        coremodsSet.addAll(coreMap.keySet());
+        System.setProperty(COREMOD_VAR, Joiner.on(',').join(coremodsSet));
+
+        // ok.. tweaker hack now.
+        if (!Strings.isNullOrEmpty(common.getTweakClass()))
         {
-            if (!url.getProtocol().startsWith("file")) // because file urls start with file://
-                continue; //         this isnt a file
+            common.extras.add("--tweakClass");
+            common.extras.add("net.minecraftforge.gradle.tweakers.CoremodTweaker");
+        }
+    }
+
+    private static void searchCoremodAtUrl(URL url, AtRegistrar atRegistrar) throws IOException, InvocationTargetException, IllegalAccessException, URISyntaxException
+    {
+        if (!url.getProtocol().startsWith("file")) // because file urls start with file://
+            return; //         this isn't a file
 
-            File coreMod = new File(url.toURI().getPath());
-            Manifest manifest = null;
+        File coreMod = new File(url.toURI().getPath());
+        Manifest manifest = null;
 
-            if (!coreMod.exists())
-                continue;
+        if (!coreMod.exists())
+            return;
 
-            if (coreMod.isDirectory())
+        if (coreMod.isDirectory())
+        {
+            File manifestMF = new File(coreMod, "META-INF/MANIFEST.MF");
+            if (manifestMF.exists())
+            {
+                FileInputStream stream = new FileInputStream(manifestMF);
+                manifest = new Manifest(stream);
+                stream.close();
+            }
+        }
+        else if (coreMod.getName().endsWith("jar")) // its a jar
+        {
+            try (JarFile jar = new JarFile(coreMod))
             {
-                File manifestMF = new File(coreMod, "META-INF/MANIFEST.MF");
-                if (manifestMF.exists())
+                manifest = jar.getManifest();
+                if (manifest != null)
                 {
-                    FileInputStream stream = new FileInputStream(manifestMF);
-                    manifest = new Manifest(stream);
-                    stream.close();
+                    atRegistrar.addJar(jar, manifest);
                 }
             }
-            else if (coreMod.getName().endsWith("jar")) // its a jar
+        }
+
+        // we got the manifest? use it.
+        if (manifest != null)
+        {
+            String clazz = manifest.getMainAttributes().getValue(COREMOD_MF);
+            if (!Strings.isNullOrEmpty(clazz))
             {
-                JarFile jar = new JarFile(coreMod);
-                manifest = jar.getManifest();
-                if (atRegistrar != null && manifest != null)
-                    atRegistrar.invoke(null, jar);
-                jar.close();
+                GradleStartCommon.LOGGER.info("Found and added coremod: " + clazz);
+                coreMap.put(clazz, coreMod);
             }
+        }
+    }
+
+    /**
+     * Hack to register jar ATs with Minecraft Forge
+     */
+    private static final class AtRegistrar
+    {
+        private static final Attributes.Name FMLAT = new Attributes.Name("FMLAT");
+
+        @Nullable
+        private Method newMethod = null;
+        @Nullable
+        private Method oldMethod = null;
 
-            // we got the manifest? use it.
-            if (manifest != null)
+        private AtRegistrar()
+        {
+            try
             {
-                String clazz = manifest.getMainAttributes().getValue(COREMOD_MF);
-                if (!Strings.isNullOrEmpty(clazz))
+                Class<?> modAtdClass = Class.forName(MOD_ATD_CLASS);
+                try
+                {
+                    newMethod = modAtdClass.getDeclaredMethod(MOD_AT_METHOD, JarFile.class, String.class);
+                }
+                catch (NoSuchMethodException | SecurityException ignored)
                 {
-                    GradleStartCommon.LOGGER.info("Found and added coremod: " + clazz);
-                    coreMap.put(clazz, coreMod);
+                    try
+                    {
+                        oldMethod = modAtdClass.getDeclaredMethod(MOD_AT_METHOD, JarFile.class);
+                    }
+                    catch (NoSuchMethodException | SecurityException ignored2)
+                    {
+                        GradleStartCommon.LOGGER.error("Failed to find method {}.{}", MOD_ATD_CLASS, MOD_AT_METHOD);
+                    }
                 }
             }
+            catch (ClassNotFoundException e)
+            {
+                GradleStartCommon.LOGGER.error("Failed to find class {}", MOD_ATD_CLASS);
+            }
         }
 
-        // set property.
-        Set<String> coremodsSet = Sets.newHashSet();
-        if (!Strings.isNullOrEmpty(System.getProperty(COREMOD_VAR)))
-            coremodsSet.addAll(Splitter.on(',').splitToList(System.getProperty(COREMOD_VAR)));
-        coremodsSet.addAll(coreMap.keySet());
-        System.setProperty(COREMOD_VAR, Joiner.on(',').join(coremodsSet));
-
-        // ok.. tweaker hack now.
-        if (!Strings.isNullOrEmpty(common.getTweakClass()))
+        public void addJar(JarFile jarFile, Manifest manifest) throws InvocationTargetException, IllegalAccessException
         {
-            common.extras.add("--tweakClass");
-            common.extras.add("net.minecraftforge.gradle.tweakers.CoremodTweaker");
+            if (newMethod != null)
+            {
+                String ats = manifest.getMainAttributes().getValue(FMLAT);
+                if (ats != null && !ats.isEmpty())
+                {
+                    newMethod.invoke(null, jarFile, ats);
+                }
+            }
+            else if (oldMethod != null)
+            {
+                oldMethod.invoke(null, jarFile);
+            }
         }
     }