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 }