    001    package cpw.mods.fml.common.asm.transformers;
    002
    003    import java.io.BufferedInputStream;
    004    import java.io.BufferedOutputStream;
    005    import java.io.ByteArrayOutputStream;
    006    import java.io.File;
    007    import java.io.FileInputStream;
    008    import java.io.FileNotFoundException;
    009    import java.io.FileOutputStream;
    010    import java.io.IOException;
    011    import java.net.URL;
    012    import java.util.List;
    013    import java.util.zip.ZipEntry;
    014    import java.util.zip.ZipInputStream;
    015    import java.util.zip.ZipOutputStream;
    016
    017    import org.objectweb.asm.ClassReader;
    018    import org.objectweb.asm.ClassWriter;
    019    import org.objectweb.asm.tree.ClassNode;
    020
    021    import com.google.common.base.Charsets;
    022    import com.google.common.base.Splitter;
    023    import com.google.common.collect.ArrayListMultimap;
    024    import com.google.common.collect.Iterables;
    025    import com.google.common.collect.ListMultimap;
    026    import com.google.common.collect.Lists;
    027    import com.google.common.io.LineProcessor;
    028    import com.google.common.io.Resources;
    029
    030    import cpw.mods.fml.relauncher.IClassTransformer;
    031
    032    public class MarkerTransformer implements IClassTransformer
    033    {
    034        private ListMultimap<String, String> markers = ArrayListMultimap.create();
    035
    036        public MarkerTransformer() throws IOException
    037        {
    038            this("fml_marker.cfg");
    039        }
    040        protected MarkerTransformer(String rulesFile) throws IOException
    041        {
    042            readMapFile(rulesFile);
    043        }
    044
    045        private void readMapFile(String rulesFile) throws IOException
    046        {
    047            File file = new File(rulesFile);
    048            URL rulesResource;
    049            if (file.exists())
    050            {
    051                rulesResource = file.toURI().toURL();
    052            }
    053            else
    054            {
    055                rulesResource = Resources.getResource(rulesFile);
    056            }
    057            Resources.readLines(rulesResource, Charsets.UTF_8, new LineProcessor<Void>()
    058            {
    059                @Override
    060                public Void getResult()
    061                {
    062                    return null;
    063                }
    064
    065                @Override
    066                public boolean processLine(String input) throws IOException
    067                {
    068                    String line = Iterables.getFirst(Splitter.on('#').limit(2).split(input), "").trim();
    069                    if (line.length()==0)
    070                    {
    071                        return true;
    072                    }
    073                    List<String> parts = Lists.newArrayList(Splitter.on(" ").trimResults().split(line));
    074                    if (parts.size()!=2)
    075                    {
    076                        throw new RuntimeException("Invalid config file line "+ input);
    077                    }
    078                    List<String> markerInterfaces = Lists.newArrayList(Splitter.on(",").trimResults().split(parts.get(1)));
    079                    for (String marker : markerInterfaces)
    080                    {
    081                        markers.put(parts.get(0), marker);
    082                    }
    083                    return true;
    084                }
    085            });
    086        }
    087
    088        @SuppressWarnings("unchecked")
    089        @Override
    090        public byte[] transform(String name, byte[] bytes)
    091        {
    092            if (bytes == null) { return null; }
    093            if (!markers.containsKey(name)) { return bytes; }
    094
    095            ClassNode classNode = new ClassNode();
    096            ClassReader classReader = new ClassReader(bytes);
    097            classReader.accept(classNode, 0);
    098
    099            for (String marker : markers.get(name))
    100            {
    101                classNode.interfaces.add(marker);
    102            }
    103
    104            ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS);
    105            classNode.accept(writer);
    106            return writer.toByteArray();
    107        }
    108
    109        public static void main(String[] args)
    110        {
    111            if (args.length < 2)
    112            {
    113                System.out.println("Usage: MarkerTransformer <JarPath> <MapFile> [MapFile2]... ");
    114                return;
    115            }
    116
    117            boolean hasTransformer = false;
    118            MarkerTransformer[] trans = new MarkerTransformer[args.length - 1];
    119            for (int x = 1; x < args.length; x++)
    120            {
    121                try
    122                {
    123                    trans[x - 1] = new MarkerTransformer(args[x]);
    124                    hasTransformer = true;
    125                }
    126                catch (IOException e)
    127                {
    128                    System.out.println("Could not read Transformer Map: " + args[x]);
    129                    e.printStackTrace();
    130                }
    131            }
    132
    133            if (!hasTransformer)
    134            {
    135                System.out.println("Culd not find a valid transformer to perform");
    136                return;
    137            }
    138
    139            File orig = new File(args[0]);
    140            File temp = new File(args[0] + ".ATBack");
    141            if (!orig.exists() && !temp.exists())
    142            {
    143                System.out.println("Could not find target jar: " + orig);
    144                return;
    145            }
    146    /*
    147            if (temp.exists())
    148            {
    149                if (orig.exists() && !orig.renameTo(new File(args[0] + (new SimpleDateFormat(".yyyy.MM.dd.HHmmss")).format(new Date()))))
    150                {
    151                    System.out.println("Could not backup existing file: " + orig);
    152                    return;
    153                }
    154                if (!temp.renameTo(orig))
    155                {
    156                    System.out.println("Could not restore backup from previous run: " + temp);
    157                    return;
    158                }
    159            }
    160    */
    161            if (!orig.renameTo(temp))
    162            {
    163                System.out.println("Could not rename file: " + orig + " -> " + temp);
    164                return;
    165            }
    166
    167            try
    168            {
    169                processJar(temp, orig, trans);
    170            }
    171            catch (IOException e)
    172            {
    173                e.printStackTrace();
    174            }
    175
    176            if (!temp.delete())
    177            {
    178                System.out.println("Could not delete temp file: " + temp);
    179            }
    180        }
    181
    182        private static void processJar(File inFile, File outFile, MarkerTransformer[] transformers) throws IOException
    183        {
    184            ZipInputStream inJar = null;
    185            ZipOutputStream outJar = null;
    186
    187            try
    188            {
    189                try
    190                {
    191                    inJar = new ZipInputStream(new BufferedInputStream(new FileInputStream(inFile)));
    192                }
    193                catch (FileNotFoundException e)
    194                {
    195                    throw new FileNotFoundException("Could not open input file: " + e.getMessage());
    196                }
    197
    198                try
    199                {
    200                    outJar = new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(outFile)));
    201                }
    202                catch (FileNotFoundException e)
    203                {
    204                    throw new FileNotFoundException("Could not open output file: " + e.getMessage());
    205                }
    206
    207                ZipEntry entry;
    208                while ((entry = inJar.getNextEntry()) != null)
    209                {
    210                    if (entry.isDirectory())
    211                    {
    212                        outJar.putNextEntry(entry);
    213                        continue;
    214                    }
    215
    216                    byte[] data = new byte[4096];
    217                    ByteArrayOutputStream entryBuffer = new ByteArrayOutputStream();
    218
    219                    int len;
    220                    do
    221                    {
    222                        len = inJar.read(data);
    223                        if (len > 0)
    224                        {
    225                            entryBuffer.write(data, 0, len);
    226                        }
    227                    }
    228                    while (len != -1);
    229
    230                    byte[] entryData = entryBuffer.toByteArray();
    231
    232                    String entryName = entry.getName();
    233
    234                    if (entryName.endsWith(".class") && !entryName.startsWith("."))
    235                    {
    236                        ClassNode cls = new ClassNode();
    237                        ClassReader rdr = new ClassReader(entryData);
    238                        rdr.accept(cls, 0);
    239                        String name = cls.name.replace('/', '.').replace('\\', '.');
    240
    241                        for (MarkerTransformer trans : transformers)
    242                        {
    243                            entryData = trans.transform(name, entryData);
    244                        }
    245                    }
    246
    247                    ZipEntry newEntry = new ZipEntry(entryName);
    248                    outJar.putNextEntry(newEntry);
    249                    outJar.write(entryData);
    250                }
    251            }
    252            finally
    253            {
    254                if (outJar != null)
    255                {
    256                    try
    257                    {
    258                        outJar.close();
    259                    }
    260                    catch (IOException e)
    261                    {
    262                    }
    263                }
    264
    265                if (inJar != null)
    266                {
    267                    try
    268                    {
    269                        inJar.close();
    270                    }
    271                    catch (IOException e)
    272                    {
    273                    }
    274                }
    275            }
    276        }
    277    }