001    package cpw.mods.fml.common.asm.transformers;
002    
003    import java.io.BufferedInputStream;
004    import java.io.BufferedOutputStream;
005    import java.io.ByteArrayOutputStream;
006    import java.io.File;
007    import java.io.FileInputStream;
008    import java.io.FileNotFoundException;
009    import java.io.FileOutputStream;
010    import java.io.IOException;
011    import java.net.URL;
012    import java.util.List;
013    import java.util.zip.ZipEntry;
014    import java.util.zip.ZipInputStream;
015    import java.util.zip.ZipOutputStream;
016    
017    import org.objectweb.asm.ClassReader;
018    import org.objectweb.asm.ClassWriter;
019    import org.objectweb.asm.tree.ClassNode;
020    
021    import com.google.common.base.Charsets;
022    import com.google.common.base.Splitter;
023    import com.google.common.collect.ArrayListMultimap;
024    import com.google.common.collect.Iterables;
025    import com.google.common.collect.ListMultimap;
026    import com.google.common.collect.Lists;
027    import com.google.common.io.LineProcessor;
028    import com.google.common.io.Resources;
029    
030    import cpw.mods.fml.relauncher.IClassTransformer;
031    
032    public class MarkerTransformer implements IClassTransformer
033    {
034        private ListMultimap<String, String> markers = ArrayListMultimap.create();
035    
036        public MarkerTransformer() throws IOException
037        {
038            this("fml_marker.cfg");
039        }
040        protected MarkerTransformer(String rulesFile) throws IOException
041        {
042            readMapFile(rulesFile);
043        }
044    
045        private void readMapFile(String rulesFile) throws IOException
046        {
047            File file = new File(rulesFile);
048            URL rulesResource;
049            if (file.exists())
050            {
051                rulesResource = file.toURI().toURL();
052            }
053            else
054            {
055                rulesResource = Resources.getResource(rulesFile);
056            }
057            Resources.readLines(rulesResource, Charsets.UTF_8, new LineProcessor<Void>()
058            {
059                @Override
060                public Void getResult()
061                {
062                    return null;
063                }
064    
065                @Override
066                public boolean processLine(String input) throws IOException
067                {
068                    String line = Iterables.getFirst(Splitter.on('#').limit(2).split(input), "").trim();
069                    if (line.length()==0)
070                    {
071                        return true;
072                    }
073                    List<String> parts = Lists.newArrayList(Splitter.on(" ").trimResults().split(line));
074                    if (parts.size()!=2)
075                    {
076                        throw new RuntimeException("Invalid config file line "+ input);
077                    }
078                    List<String> markerInterfaces = Lists.newArrayList(Splitter.on(",").trimResults().split(parts.get(1)));
079                    for (String marker : markerInterfaces)
080                    {
081                        markers.put(parts.get(0), marker);
082                    }
083                    return true;
084                }
085            });
086        }
087    
088        @SuppressWarnings("unchecked")
089        @Override
090        public byte[] transform(String name, byte[] bytes)
091        {
092            if (bytes == null) { return null; }
093            if (!markers.containsKey(name)) { return bytes; }
094    
095            ClassNode classNode = new ClassNode();
096            ClassReader classReader = new ClassReader(bytes);
097            classReader.accept(classNode, 0);
098    
099            for (String marker : markers.get(name))
100            {
101                classNode.interfaces.add(marker);
102            }
103    
104            ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS);
105            classNode.accept(writer);
106            return writer.toByteArray();
107        }
108    
109        public static void main(String[] args)
110        {
111            if (args.length < 2)
112            {
113                System.out.println("Usage: MarkerTransformer <JarPath> <MapFile> [MapFile2]... ");
114                return;
115            }
116    
117            boolean hasTransformer = false;
118            MarkerTransformer[] trans = new MarkerTransformer[args.length - 1];
119            for (int x = 1; x < args.length; x++)
120            {
121                try
122                {
123                    trans[x - 1] = new MarkerTransformer(args[x]);
124                    hasTransformer = true;
125                }
126                catch (IOException e)
127                {
128                    System.out.println("Could not read Transformer Map: " + args[x]);
129                    e.printStackTrace();
130                }
131            }
132    
133            if (!hasTransformer)
134            {
135                System.out.println("Culd not find a valid transformer to perform");
136                return;
137            }
138    
139            File orig = new File(args[0]);
140            File temp = new File(args[0] + ".ATBack");
141            if (!orig.exists() && !temp.exists())
142            {
143                System.out.println("Could not find target jar: " + orig);
144                return;
145            }
146    /*
147            if (temp.exists())
148            {
149                if (orig.exists() && !orig.renameTo(new File(args[0] + (new SimpleDateFormat(".yyyy.MM.dd.HHmmss")).format(new Date()))))
150                {
151                    System.out.println("Could not backup existing file: " + orig);
152                    return;
153                }
154                if (!temp.renameTo(orig))
155                {
156                    System.out.println("Could not restore backup from previous run: " + temp);
157                    return;
158                }
159            }
160    */
161            if (!orig.renameTo(temp))
162            {
163                System.out.println("Could not rename file: " + orig + " -> " + temp);
164                return;
165            }
166    
167            try
168            {
169                processJar(temp, orig, trans);
170            }
171            catch (IOException e)
172            {
173                e.printStackTrace();
174            }
175    
176            if (!temp.delete())
177            {
178                System.out.println("Could not delete temp file: " + temp);
179            }
180        }
181    
182        private static void processJar(File inFile, File outFile, MarkerTransformer[] transformers) throws IOException
183        {
184            ZipInputStream inJar = null;
185            ZipOutputStream outJar = null;
186    
187            try
188            {
189                try
190                {
191                    inJar = new ZipInputStream(new BufferedInputStream(new FileInputStream(inFile)));
192                }
193                catch (FileNotFoundException e)
194                {
195                    throw new FileNotFoundException("Could not open input file: " + e.getMessage());
196                }
197    
198                try
199                {
200                    outJar = new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(outFile)));
201                }
202                catch (FileNotFoundException e)
203                {
204                    throw new FileNotFoundException("Could not open output file: " + e.getMessage());
205                }
206    
207                ZipEntry entry;
208                while ((entry = inJar.getNextEntry()) != null)
209                {
210                    if (entry.isDirectory())
211                    {
212                        outJar.putNextEntry(entry);
213                        continue;
214                    }
215    
216                    byte[] data = new byte[4096];
217                    ByteArrayOutputStream entryBuffer = new ByteArrayOutputStream();
218    
219                    int len;
220                    do
221                    {
222                        len = inJar.read(data);
223                        if (len > 0)
224                        {
225                            entryBuffer.write(data, 0, len);
226                        }
227                    }
228                    while (len != -1);
229    
230                    byte[] entryData = entryBuffer.toByteArray();
231    
232                    String entryName = entry.getName();
233    
234                    if (entryName.endsWith(".class") && !entryName.startsWith("."))
235                    {
236                        ClassNode cls = new ClassNode();
237                        ClassReader rdr = new ClassReader(entryData);
238                        rdr.accept(cls, 0);
239                        String name = cls.name.replace('/', '.').replace('\\', '.');
240    
241                        for (MarkerTransformer trans : transformers)
242                        {
243                            entryData = trans.transform(name, entryData);
244                        }
245                    }
246    
247                    ZipEntry newEntry = new ZipEntry(entryName);
248                    outJar.putNextEntry(newEntry);
249                    outJar.write(entryData);
250                }
251            }
252            finally
253            {
254                if (outJar != null)
255                {
256                    try
257                    {
258                        outJar.close();
259                    }
260                    catch (IOException e)
261                    {
262                    }
263                }
264    
265                if (inJar != null)
266                {
267                    try
268                    {
269                        inJar.close();
270                    }
271                    catch (IOException e)
272                    {
273                    }
274                }
275            }
276        }
277    }