    001package cpw.mods.fml.common.asm.transformers;
    002
    003import java.io.BufferedInputStream;
    004import java.io.BufferedOutputStream;
    005import java.io.ByteArrayOutputStream;
    006import java.io.File;
    007import java.io.FileInputStream;
    008import java.io.FileNotFoundException;
    009import java.io.FileOutputStream;
    010import java.io.IOException;
    011import java.net.URL;
    012import java.util.List;
    013import java.util.zip.ZipEntry;
    014import java.util.zip.ZipInputStream;
    015import java.util.zip.ZipOutputStream;
    016
    017import org.objectweb.asm.ClassReader;
    018import org.objectweb.asm.ClassWriter;
    019import org.objectweb.asm.tree.ClassNode;
    020
    021import com.google.common.base.Charsets;
    022import com.google.common.base.Splitter;
    023import com.google.common.collect.ArrayListMultimap;
    024import com.google.common.collect.Iterables;
    025import com.google.common.collect.ListMultimap;
    026import com.google.common.collect.Lists;
    027import com.google.common.io.LineProcessor;
    028import com.google.common.io.Resources;
    029
    030import cpw.mods.fml.relauncher.IClassTransformer;
    031
    032public 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}