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}