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