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    }