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