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