    001package cpw.mods.fml.common.asm.transformers;
    002
    003import static org.objectweb.asm.Opcodes.ACC_FINAL;
    004import static org.objectweb.asm.Opcodes.ACC_PRIVATE;
    005import static org.objectweb.asm.Opcodes.ACC_PROTECTED;
    006import static org.objectweb.asm.Opcodes.ACC_PUBLIC;
    007
    008import java.io.BufferedInputStream;
    009import java.io.BufferedOutputStream;
    010import java.io.ByteArrayOutputStream;
    011import java.io.File;
    012import java.io.FileInputStream;
    013import java.io.FileNotFoundException;
    014import java.io.FileOutputStream;
    015import java.io.IOException;
    016import java.net.URL;
    017import java.util.Collection;
    018import java.util.List;
    019import java.util.zip.ZipEntry;
    020import java.util.zip.ZipInputStream;
    021import java.util.zip.ZipOutputStream;
    022
    023import org.objectweb.asm.ClassReader;
    024import org.objectweb.asm.ClassWriter;
    025import org.objectweb.asm.tree.ClassNode;
    026import org.objectweb.asm.tree.FieldNode;
    027import org.objectweb.asm.tree.MethodNode;
    028
    029import com.google.common.base.Charsets;
    030import com.google.common.base.Splitter;
    031import com.google.common.collect.ArrayListMultimap;
    032import com.google.common.collect.Iterables;
    033import com.google.common.collect.Lists;
    034import com.google.common.collect.Multimap;
    035import com.google.common.io.LineProcessor;
    036import com.google.common.io.Resources;
    037
    038import cpw.mods.fml.relauncher.IClassTransformer;
    039
    040public class AccessTransformer implements IClassTransformer
    041{
    042    private static final boolean DEBUG = false;
    043    private class Modifier
    044    {
    045        public String name = "";
    046        public String desc = "";
    047        public int oldAccess = 0;
    048        public int newAccess = 0;
    049        public int targetAccess = 0;
    050        public boolean changeFinal = false;
    051        public boolean markFinal = false;
    052        protected boolean modifyClassVisibility;
    053
    054        private void setTargetAccess(String name)
    055        {
    056            if (name.startsWith("public")) targetAccess = ACC_PUBLIC;
    057            else if (name.startsWith("private")) targetAccess = ACC_PRIVATE;
    058            else if (name.startsWith("protected")) targetAccess = ACC_PROTECTED;
    059
    060            if (name.endsWith("-f"))
    061            {
    062                changeFinal = true;
    063                markFinal = false;
    064            }
    065            else if (name.endsWith("+f"))
    066            {
    067                changeFinal = true;
    068                markFinal = true;
    069            }
    070        }
    071    }
    072
    073    private Multimap<String, Modifier> modifiers = ArrayListMultimap.create();
    074
    075    public AccessTransformer() throws IOException
    076    {
    077        this("fml_at.cfg");
    078    }
    079    protected AccessTransformer(String rulesFile) throws IOException
    080    {
    081        readMapFile(rulesFile);
    082    }
    083
    084    private void readMapFile(String rulesFile) throws IOException
    085    {
    086        File file = new File(rulesFile);
    087        URL rulesResource;
    088        if (file.exists())
    089        {
    090            rulesResource = file.toURI().toURL();
    091        }
    092        else
    093        {
    094            rulesResource = Resources.getResource(rulesFile);
    095        }
    096        Resources.readLines(rulesResource, Charsets.UTF_8, new LineProcessor<Void>()
    097        {
    098            @Override
    099            public Void getResult()
    100            {
    101                return null;
    102            }
    103
    104            @Override
    105            public boolean processLine(String input) throws IOException
    106            {
    107                String line = Iterables.getFirst(Splitter.on('#').limit(2).split(input), "").trim();
    108                if (line.length()==0)
    109                {
    110                    return true;
    111                }
    112                List<String> parts = Lists.newArrayList(Splitter.on(" ").trimResults().split(line));
    113                if (parts.size()>2)
    114                {
    115                    throw new RuntimeException("Invalid config file line "+ input);
    116                }
    117                Modifier m = new Modifier();
    118                m.setTargetAccess(parts.get(0));
    119                List<String> descriptor = Lists.newArrayList(Splitter.on(".").trimResults().split(parts.get(1)));
    120                if (descriptor.size() == 1)
    121                {
    122                    m.modifyClassVisibility = true;
    123                }
    124                else
    125                {
    126                    String nameReference = descriptor.get(1);
    127                    int parenIdx = nameReference.indexOf('(');
    128                    if (parenIdx>0)
    129                    {
    130                        m.desc = nameReference.substring(parenIdx);
    131                        m.name = nameReference.substring(0,parenIdx);
    132                    }
    133                    else
    134                    {
    135                        m.name = nameReference;
    136                    }
    137                }
    138                modifiers.put(descriptor.get(0).replace('/', '.'), m);
    139                return true;
    140            }
    141        });
    142    }
    143
    144    @SuppressWarnings("unchecked")
    145    @Override
    146    public byte[] transform(String name, byte[] bytes)
    147    {
    148        if (bytes == null) { return null; }
    149        if (!modifiers.containsKey(name)) { return bytes; }
    150
    151        ClassNode classNode = new ClassNode();
    152        ClassReader classReader = new ClassReader(bytes);
    153        classReader.accept(classNode, 0);
    154
    155        Collection<Modifier> mods = modifiers.get(name);
    156        for (Modifier m : mods)
    157        {
    158            if (m.modifyClassVisibility)
    159            {
    160                classNode.access = getFixedAccess(classNode.access, m);
    161                if (DEBUG)
    162                {
    163                    System.out.println(String.format("Class: %s %s -> %s", name, toBinary(m.oldAccess), toBinary(m.newAccess)));
    164                }
    165                continue;
    166            }
    167            if (m.desc.isEmpty())
    168            {
    169                for (FieldNode n : (List<FieldNode>) classNode.fields)
    170                {
    171                    if (n.name.equals(m.name) || m.name.equals("*"))
    172                    {
    173                        n.access = getFixedAccess(n.access, m);
    174                        if (DEBUG)
    175                        {
    176                            System.out.println(String.format("Field: %s.%s %s -> %s", name, n.name, toBinary(m.oldAccess), toBinary(m.newAccess)));
    177                        }
    178
    179                        if (!m.name.equals("*"))
    180                        {
    181                            break;
    182                        }
    183                    }
    184                }
    185            }
    186            else
    187            {
    188                for (MethodNode n : (List<MethodNode>) classNode.methods)
    189                {
    190                    if ((n.name.equals(m.name) && n.desc.equals(m.desc)) || m.name.equals("*"))
    191                    {
    192                        n.access = getFixedAccess(n.access, m);
    193                        if (DEBUG)
    194                        {
    195                            System.out.println(String.format("Method: %s.%s%s %s -> %s", name, n.name, n.desc, toBinary(m.oldAccess), toBinary(m.newAccess)));
    196                        }
    197
    198                        if (!m.name.equals("*"))
    199                        {
    200                            break;
    201                        }
    202                    }
    203                }
    204            }
    205        }
    206
    207        ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS);
    208        classNode.accept(writer);
    209        return writer.toByteArray();
    210    }
    211
    212    private String toBinary(int num)
    213    {
    214        return String.format("%16s", Integer.toBinaryString(num)).replace(' ', '0');
    215    }
    216
    217    private int getFixedAccess(int access, Modifier target)
    218    {
    219        target.oldAccess = access;
    220        int t = target.targetAccess;
    221        int ret = (access & ~7);
    222
    223        switch (access & 7)
    224        {
    225        case ACC_PRIVATE:
    226            ret |= t;
    227            break;
    228        case 0: // default
    229            ret |= (t != ACC_PRIVATE ? t : 0 /* default */);
    230            break;
    231        case ACC_PROTECTED:
    232            ret |= (t != ACC_PRIVATE && t != 0 /* default */? t : ACC_PROTECTED);
    233            break;
    234        case ACC_PUBLIC:
    235            ret |= (t != ACC_PRIVATE && t != 0 /* default */&& t != ACC_PROTECTED ? t : ACC_PUBLIC);
    236            break;
    237        default:
    238            throw new RuntimeException("The fuck?");
    239        }
    240
    241        // Clear the "final" marker on fields only if specified in control field
    242        if (target.changeFinal && target.desc == "")
    243        {
    244            if (target.markFinal)
    245            {
    246                ret |= ACC_FINAL;
    247            }
    248            else
    249            {
    250                ret &= ~ACC_FINAL;
    251            }
    252        }
    253        target.newAccess = ret;
    254        return ret;
    255    }
    256
    257    public static void main(String[] args)
    258    {
    259        if (args.length < 2)
    260        {
    261            System.out.println("Usage: AccessTransformer <JarPath> <MapFile> [MapFile2]... ");
    262            System.exit(1);
    263        }
    264
    265        boolean hasTransformer = false;
    266        AccessTransformer[] trans = new AccessTransformer[args.length - 1];
    267        for (int x = 1; x < args.length; x++)
    268        {
    269            try
    270            {
    271                trans[x - 1] = new AccessTransformer(args[x]);
    272                hasTransformer = true;
    273            }
    274            catch (IOException e)
    275            {
    276                System.out.println("Could not read Transformer Map: " + args[x]);
    277                e.printStackTrace();
    278            }
    279        }
    280
    281        if (!hasTransformer)
    282        {
    283            System.out.println("Culd not find a valid transformer to perform");
    284            System.exit(1);
    285        }
    286
    287        File orig = new File(args[0]);
    288        File temp = new File(args[0] + ".ATBack");
    289        if (!orig.exists() && !temp.exists())
    290        {
    291            System.out.println("Could not find target jar: " + orig);
    292            System.exit(1);
    293        }
    294
    295        if (!orig.renameTo(temp))
    296        {
    297            System.out.println("Could not rename file: " + orig + " -> " + temp);
    298            System.exit(1);
    299        }
    300
    301        try
    302        {
    303            processJar(temp, orig, trans);
    304        }
    305        catch (IOException e)
    306        {
    307            e.printStackTrace();
    308            System.exit(1);
    309        }
    310
    311        if (!temp.delete())
    312        {
    313            System.out.println("Could not delete temp file: " + temp);
    314        }
    315    }
    316
    317    private static void processJar(File inFile, File outFile, AccessTransformer[] transformers) throws IOException
    318    {
    319        ZipInputStream inJar = null;
    320        ZipOutputStream outJar = null;
    321
    322        try
    323        {
    324            try
    325            {
    326                inJar = new ZipInputStream(new BufferedInputStream(new FileInputStream(inFile)));
    327            }
    328            catch (FileNotFoundException e)
    329            {
    330                throw new FileNotFoundException("Could not open input file: " + e.getMessage());
    331            }
    332
    333            try
    334            {
    335                outJar = new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(outFile)));
    336            }
    337            catch (FileNotFoundException e)
    338            {
    339                throw new FileNotFoundException("Could not open output file: " + e.getMessage());
    340            }
    341
    342            ZipEntry entry;
    343            while ((entry = inJar.getNextEntry()) != null)
    344            {
    345                if (entry.isDirectory())
    346                {
    347                    outJar.putNextEntry(entry);
    348                    continue;
    349                }
    350
    351                byte[] data = new byte[4096];
    352                ByteArrayOutputStream entryBuffer = new ByteArrayOutputStream();
    353
    354                int len;
    355                do
    356                {
    357                    len = inJar.read(data);
    358                    if (len > 0)
    359                    {
    360                        entryBuffer.write(data, 0, len);
    361                    }
    362                }
    363                while (len != -1);
    364
    365                byte[] entryData = entryBuffer.toByteArray();
    366
    367                String entryName = entry.getName();
    368
    369                if (entryName.endsWith(".class") && !entryName.startsWith("."))
    370                {
    371                    ClassNode cls = new ClassNode();
    372                    ClassReader rdr = new ClassReader(entryData);
    373                    rdr.accept(cls, 0);
    374                    String name = cls.name.replace('/', '.').replace('\\', '.');
    375
    376                    for (AccessTransformer trans : transformers)
    377                    {
    378                        entryData = trans.transform(name, entryData);
    379                    }
    380                }
    381
    382                ZipEntry newEntry = new ZipEntry(entryName);
    383                outJar.putNextEntry(newEntry);
    384                outJar.write(entryData);
    385            }
    386        }
    387        finally
    388        {
    389            if (outJar != null)
    390            {
    391                try
    392                {
    393                    outJar.close();
    394                }
    395                catch (IOException e)
    396                {
    397                }
    398            }
    399
    400            if (inJar != null)
    401            {
    402                try
    403                {
    404                    inJar.close();
    405                }
    406                catch (IOException e)
    407                {
    408                }
    409            }
    410        }
    411    }
    412    public void ensurePublicAccessFor(String modClazzName)
    413    {
    414        Modifier m = new Modifier();
    415        m.setTargetAccess("public");
    416        m.modifyClassVisibility = true;
    417        modifiers.put(modClazzName, m);
    418    }
    419}