    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 (!markers.containsKey(name)) { return bytes; }
    093
    094            ClassNode classNode = new ClassNode();
    095            ClassReader classReader = new ClassReader(bytes);
    096            classReader.accept(classNode, 0);
    097
    098            for (String marker : markers.get(name))
    099            {
    100                classNode.interfaces.add(marker);
    101            }
    102
    103            ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS);
    104            classNode.accept(writer);
    105            return writer.toByteArray();
    106        }
    107
    108        public static void main(String[] args)
    109        {
    110            if (args.length < 2)
    111            {
    112                System.out.println("Usage: MarkerTransformer <JarPath> <MapFile> [MapFile2]... ");
    113                return;
    114            }
    115
    116            boolean hasTransformer = false;
    117            MarkerTransformer[] trans = new MarkerTransformer[args.length - 1];
    118            for (int x = 1; x < args.length; x++)
    119            {
    120                try
    121                {
    122                    trans[x - 1] = new MarkerTransformer(args[x]);
    123                    hasTransformer = true;
    124                }
    125                catch (IOException e)
    126                {
    127                    System.out.println("Could not read Transformer Map: " + args[x]);
    128                    e.printStackTrace();
    129                }
    130            }
    131
    132            if (!hasTransformer)
    133            {
    134                System.out.println("Culd not find a valid transformer to perform");
    135                return;
    136            }
    137
    138            File orig = new File(args[0]);
    139            File temp = new File(args[0] + ".ATBack");
    140            if (!orig.exists() && !temp.exists())
    141            {
    142                System.out.println("Could not find target jar: " + orig);
    143                return;
    144            }
    145    /*
    146            if (temp.exists())
    147            {
    148                if (orig.exists() && !orig.renameTo(new File(args[0] + (new SimpleDateFormat(".yyyy.MM.dd.HHmmss")).format(new Date()))))
    149                {
    150                    System.out.println("Could not backup existing file: " + orig);
    151                    return;
    152                }
    153                if (!temp.renameTo(orig))
    154                {
    155                    System.out.println("Could not restore backup from previous run: " + temp);
    156                    return;
    157                }
    158            }
    159    */
    160            if (!orig.renameTo(temp))
    161            {
    162                System.out.println("Could not rename file: " + orig + " -> " + temp);
    163                return;
    164            }
    165
    166            try
    167            {
    168                processJar(temp, orig, trans);
    169            }
    170            catch (IOException e)
    171            {
    172                e.printStackTrace();
    173            }
    174
    175            if (!temp.delete())
    176            {
    177                System.out.println("Could not delete temp file: " + temp);
    178            }
    179        }
    180
    181        private static void processJar(File inFile, File outFile, MarkerTransformer[] transformers) throws IOException
    182        {
    183            ZipInputStream inJar = null;
    184            ZipOutputStream outJar = null;
    185
    186            try
    187            {
    188                try
    189                {
    190                    inJar = new ZipInputStream(new BufferedInputStream(new FileInputStream(inFile)));
    191                }
    192                catch (FileNotFoundException e)
    193                {
    194                    throw new FileNotFoundException("Could not open input file: " + e.getMessage());
    195                }
    196
    197                try
    198                {
    199                    outJar = new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(outFile)));
    200                }
    201                catch (FileNotFoundException e)
    202                {
    203                    throw new FileNotFoundException("Could not open output file: " + e.getMessage());
    204                }
    205
    206                ZipEntry entry;
    207                while ((entry = inJar.getNextEntry()) != null)
    208                {
    209                    if (entry.isDirectory())
    210                    {
    211                        outJar.putNextEntry(entry);
    212                        continue;
    213                    }
    214
    215                    byte[] data = new byte[4096];
    216                    ByteArrayOutputStream entryBuffer = new ByteArrayOutputStream();
    217
    218                    int len;
    219                    do
    220                    {
    221                        len = inJar.read(data);
    222                        if (len > 0)
    223                        {
    224                            entryBuffer.write(data, 0, len);
    225                        }
    226                    }
    227                    while (len != -1);
    228
    229                    byte[] entryData = entryBuffer.toByteArray();
    230
    231                    String entryName = entry.getName();
    232
    233                    if (entryName.endsWith(".class") && !entryName.startsWith("."))
    234                    {
    235                        ClassNode cls = new ClassNode();
    236                        ClassReader rdr = new ClassReader(entryData);
    237                        rdr.accept(cls, 0);
    238                        String name = cls.name.replace('/', '.').replace('\\', '.');
    239
    240                        for (MarkerTransformer trans : transformers)
    241                        {
    242                            entryData = trans.transform(name, entryData);
    243                        }
    244                    }
    245
    246                    ZipEntry newEntry = new ZipEntry(entryName);
    247                    outJar.putNextEntry(newEntry);
    248                    outJar.write(entryData);
    249                }
    250            }
    251            finally
    252            {
    253                if (outJar != null)
    254                {
    255                    try
    256                    {
    257                        outJar.close();
    258                    }
    259                    catch (IOException e)
    260                    {
    261                    }
    262                }
    263
    264                if (inJar != null)
    265                {
    266                    try
    267                    {
    268                        inJar.close();
    269                    }
    270                    catch (IOException e)
    271                    {
    272                    }
    273                }
    274            }
    275        }
    276    }