-
-
Notifications
You must be signed in to change notification settings - Fork 44
/
ClassRegistry.java
389 lines (367 loc) · 14.4 KB
/
ClassRegistry.java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
package me.nallar.tickthreading.patcher;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.zip.ZipEntry;
import java.util.zip.ZipException;
import java.util.zip.ZipFile;
import java.util.zip.ZipInputStream;
import java.util.zip.ZipOutputStream;
import com.google.common.io.ByteStreams;
import com.google.common.io.Files;
import javassist.ClassPath;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.NotFoundException;
import javassist.bytecode.MethodInfo;
import me.nallar.tickthreading.Log;
import me.nallar.tickthreading.util.CollectionsUtil;
import me.nallar.tickthreading.util.IterableEnumerationWrapper;
import me.nallar.tickthreading.util.LocationUtil;
import me.nallar.unsafe.UnsafeUtil;
import net.minecraft.server.MinecraftServer;
// Unchecked - Enumeration cast due to old java APIs.
// ResultOfMethodCallIgnored - No, I do not care if mkdir fails to make the directory.
// FieldRepeatedlyAccessedInMethod - I don't care about minor optimisations in patcher code, can't find a way to change inspections per package.
@SuppressWarnings ({"unchecked", "ResultOfMethodCallIgnored", "FieldRepeatedlyAccessedInMethod"})
public class ClassRegistry {
private final String hashFileName = "patcher.hash";
private final String patchedModsFolderName = "patchedMods";
private final String modsFolderName = "mods";
private final Map<String, File> classNameToLocation = new HashMap<String, File>();
private final Map<File, Integer> locationToPatchHash = new HashMap<File, Integer>();
private final Map<File, Integer> expectedPatchHashes = new HashMap<File, Integer>();
private final Set<File> loadedFiles = new HashSet<File>();
private final Set<File> updatedFiles = new HashSet<File>();
private final Map<String, Set<File>> duplicateClassNamesToLocations = new HashMap<String, Set<File>>();
private final Set<ClassPath> classPathSet = new HashSet<ClassPath>();
private final Map<String, byte[]> replacementFiles = new HashMap<String, byte[]>();
public File serverFile;
private File patchedModsFolder;
public final ClassPool classes = new ClassPool(false);
public boolean disableJavassistLoading = false;
public boolean forcePatching = false;
public boolean writeAllClasses = false;
{
MethodInfo.doPreverify = true;
classes.appendSystemPath();
classes.importPackage("java.util");
classes.importPackage("java.io");
classes.importPackage("me.nallar.tickthreading");
}
public void clearClassInfo() {
closeClassPath();
classNameToLocation.clear();
updatedFiles.clear();
duplicateClassNamesToLocations.clear();
replacementFiles.clear();
loadedFiles.clear();
}
public void closeClassPath() {
for (ClassPath classPath : classPathSet) {
classes.removeClassPath(classPath);
}
classPathSet.clear();
}
public void loadFiles(Iterable<File> filesToLoad) throws IOException {
for (File file : filesToLoad) {
String extension = file.getName().toLowerCase();
extension = extension.substring(extension.lastIndexOf('.') + 1);
try {
File[] files = file.listFiles();
if (files != null) {
if (!".disabled".equals(file.getName()) && !patchedModsFolderName.equalsIgnoreCase(file.getName())) {
loadFiles(Arrays.asList(files));
}
} else if ("jar".equals(extension) || "zip".equals(extension) || "litemod".equals(extension)) {
ZipFile zipFile = new ZipFile(file);
try {
loadZip(zipFile);
} finally {
zipFile.close();
}
loadHashes(file);
}
} catch (ZipException e) {
throw new ZipException(e.getMessage() + " file: " + file);
}
}
}
public void update(String className, byte[] replacement) {
Collection<File> duplicates = duplicateClassNamesToLocations.get(className);
if (duplicates != null) {
Log.warning(className + " is in multiple jars: " + CollectionsUtil.join(duplicates, ", "));
updatedFiles.addAll(duplicates);
}
updatedFiles.add(classNameToLocation.get(className));
replacementFiles.put(className.replace('.', '/') + ".class", replacement);
}
@SuppressWarnings ("IOResourceOpenedButNotSafelyClosed")
public void save(File backupDirectory) throws IOException {
closeClassPath();
File tempFile = null, renameFile = null;
File tempDirectory = new File(backupDirectory.getParentFile(), "TTTemp");
tempDirectory.mkdir();
ZipInputStream zin = null;
ZipOutputStream zout = null;
backupDirectory.mkdir();
patchedModsFolder.mkdir();
File modsFolder = new File(patchedModsFolder.getParent(), "mods");
int patchedClasses = 0;
try {
for (File zipFile : updatedFiles) {
if (zipFile == serverFile || (!zipFile.equals(LocationUtil.locationOf(PatchMain.class).getAbsoluteFile()) && !modsFolderName.equals(zipFile.getParentFile().getName()))) {
File backupFile = new File(backupDirectory, zipFile.getName());
if (backupFile.exists() && !backupFile.delete()) {
Log.warning("Failed to remove old backup");
}
Files.copy(zipFile, backupFile);
tempFile = makeTempFile(tempDirectory, zipFile);
if (zipFile.renameTo(tempFile)) {
renameFile = zipFile;
tempFile.deleteOnExit();
} else {
throw new IOException("Couldn't rename " + zipFile + " -> " + tempFile);
}
zin = new ZipInputStream(new FileInputStream(tempFile));
zout = new ZipOutputStream(new FileOutputStream(zipFile));
writeChanges(zipFile, zin, zout, false);
if (!tempFile.delete()) {
Log.warning("Failed to delete temporary patching file " + tempFile + " after patching " + zipFile);
}
renameFile = null;
} else {
zin = new ZipInputStream(new FileInputStream(zipFile));
File patchedModFile = new File(patchedModsFolder, zipFile.getName());
if (patchedModFile.exists() && !patchedModFile.delete()) {
Log.severe("Failed to write patches for " + zipFile + ", could not delete old patchedMods file.");
}
zout = new ZipOutputStream(new FileOutputStream(patchedModFile));
patchedClasses += writeChanges(zipFile, zin, zout, true);
}
zin = null;
zout = null;
}
} catch (Exception e) {
if (zin != null) {
zin.close();
}
if (zout != null) {
zout.close();
}
if (renameFile != null && !tempFile.renameTo(renameFile)) {
Log.warning("Failed to restore " + renameFile + " after patching, you will need to get a new copy of it.");
}
throw UnsafeUtil.throwIgnoreChecked(e);
} finally {
delete(tempDirectory);
}
Log.info("Patched " + patchedClasses + " mod classes.");
for (File file : patchedModsFolder.listFiles()) {
if (!new File(modsFolder, file.getName()).exists() && file.delete()) {
Log.info("Deleted old patched mod file " + file.getName());
}
}
}
public CtClass getClass(String className) throws NotFoundException {
return classes.get(className);
}
public void loadPatchHashes(PatchManager patchManager) {
Map<String, Integer> patchHashes = patchManager.getHashes();
for (Map.Entry<String, Integer> stringIntegerEntry : patchHashes.entrySet()) {
File location = classNameToLocation.get(stringIntegerEntry.getKey());
if (location == null) {
continue;
}
int hash = stringIntegerEntry.getValue();
Integer currentHash = expectedPatchHashes.get(location);
currentHash = (currentHash == null) ? hash : currentHash * 13 + hash;
expectedPatchHashes.put(location, currentHash);
}
}
public boolean shouldPatch(String className) {
return shouldPatch(classNameToLocation.get(className));
}
boolean shouldPatch(File file) {
return forcePatching || file == null || !(expectedPatchHashes.get(file).equals(locationToPatchHash.get(file)));
}
public boolean shouldPatch() {
boolean shouldPatch = false;
for (Map.Entry<File, Integer> fileIntegerEntry : expectedPatchHashes.entrySet()) {
if (shouldPatch(fileIntegerEntry.getKey())) {
Integer expectedPatchHash = fileIntegerEntry.getValue();
Log.warning("Patching required for " + fileIntegerEntry.getKey() + " because " + (expectedPatchHash == null ? "it has not been patched" : "it is out of date." +
"\nExpected " + expectedPatchHash + ", got " + locationToPatchHash.get(fileIntegerEntry.getKey())));
shouldPatch = true;
}
}
return shouldPatch;
}
public void restoreBackups(File backupDirectory) {
for (Map.Entry<File, Integer> fileIntegerEntry : locationToPatchHash.entrySet()) {
Integer expectedHash = expectedPatchHashes.get(fileIntegerEntry.getKey());
Integer actualHash = fileIntegerEntry.getValue();
if (actualHash == null) {
continue;
}
if (forcePatching || !actualHash.equals(expectedHash)) {
File backupFile = new File(backupDirectory, fileIntegerEntry.getKey().getName());
if (new File(patchedModsFolder, backupFile.getName()).exists()) {
continue;
}
boolean restoreFailed = false;
if (!backupFile.exists()) {
restoreFailed = true;
} else {
fileIntegerEntry.getKey().delete();
try {
Files.move(backupFile, fileIntegerEntry.getKey());
} catch (IOException e) {
restoreFailed = true;
Log.severe("Failed to restore unpatched backup before patching", e);
}
}
if (restoreFailed) {
Log.severe("Can't patch - no backup for " + fileIntegerEntry.getKey().getName() + " exists, and a patched copy is already in the mods directory." +
"\nYou will need to replace this file with a new unpatched copy.");
throw new Error("Missing backup for patched file");
}
}
}
}
void appendClassPath(String path) throws NotFoundException {
if (!disableJavassistLoading) {
classPathSet.add(classes.appendClassPath(path));
}
}
void loadHashes(File zipFile) throws IOException {
if (!modsFolderName.equalsIgnoreCase(zipFile.getParentFile().getName())) {
return;
}
if (patchedModsFolder == null) {
patchedModsFolder = new File(zipFile.getParentFile().getParentFile(), patchedModsFolderName);
}
File file = new File(patchedModsFolder, zipFile.getName());
if (!file.exists()) {
return;
}
try {
if (MinecraftServer.getServer() == null) {
throw new NullPointerException();
}
} catch (Throwable t) {
if (!file.delete()) {
Log.warning("Unable to delete old patchedMods file " + file);
}
return;
}
ZipFile zip = new ZipFile(file);
try {
for (ZipEntry zipEntry : new IterableEnumerationWrapper<ZipEntry>((Enumeration<ZipEntry>) zip.entries())) {
String name = zipEntry.getName();
if (name.equals(hashFileName)) {
ByteArrayOutputStream output = new ByteArrayOutputStream();
ByteStreams.copy(zip.getInputStream(zipEntry), output);
int hash = Integer.valueOf(new String(output.toByteArray(), "UTF-8"));
locationToPatchHash.put(zipFile, hash);
}
}
} finally {
zip.close();
}
}
void loadZip(ZipFile zip) throws IOException {
File file = new File(zip.getName());
if (!loadedFiles.add(file)) {
return;
}
try {
appendClassPath(file.getAbsolutePath());
} catch (Exception e) {
Log.severe("Javassist could not load " + file, e);
}
for (ZipEntry zipEntry : new IterableEnumerationWrapper<ZipEntry>((Enumeration<ZipEntry>) zip.entries())) {
String name = zipEntry.getName();
if (name.endsWith(".class")) {
String className = name.replace('/', '.').substring(0, name.lastIndexOf('.'));
if (classNameToLocation.containsKey(className)) {
Set<File> locations = duplicateClassNamesToLocations.get(className);
if (locations == null) {
locations = new HashSet<File>();
locations.add(classNameToLocation.get(className));
duplicateClassNamesToLocations.put(className, locations);
}
locations.add(file);
} else {
classNameToLocation.put(className, file);
}
} else if (name.equals(hashFileName)) {
ByteArrayOutputStream output = new ByteArrayOutputStream();
ByteStreams.copy(zip.getInputStream(zipEntry), output);
int hash = Integer.valueOf(new String(output.toByteArray(), "UTF-8"));
locationToPatchHash.put(file, hash);
}
}
}
private static File makeTempFile(File tempLocation, File file) {
File tempFile = new File(tempLocation, file.getName() + ".tmp");
if (tempFile.exists() && !tempFile.delete()) {
throw new Error("Failed to delete old temp file " + tempFile);
}
return tempFile;
}
private static void delete(File f) {
File[] files = f.listFiles();
if (files != null) {
for (File c : files) {
delete(c);
}
}
f.delete();
}
@SuppressWarnings ("StatementWithEmptyBody")
private int writeChanges(File zipFile, ZipInputStream zin, ZipOutputStream zout, boolean onlyClasses) throws Exception {
int patchedClasses = 0;
Set<String> replacements = new HashSet<String>();
ZipEntry zipEntry;
while ((zipEntry = zin.getNextEntry()) != null) {
String entryName = zipEntry.getName();
if (entryName.equals(hashFileName) || (entryName.startsWith("META-INF/") && !(!entryName.isEmpty() && entryName.charAt(entryName.length() - 1) == '/') && !entryName.toUpperCase().endsWith("MANIFEST.MF")) && (entryName.length() - entryName.replace("/", "").length() == 1)) {
// Skip
} else if (onlyClasses && !entryName.toLowerCase().endsWith(".class")) {
// Skip
} else if (replacementFiles.containsKey(entryName)) {
replacements.add(entryName);
patchedClasses++;
} else if (!onlyClasses || writeAllClasses) {
zout.putNextEntry(new ZipEntry(entryName));
ByteStreams.copy(zin, zout);
patchedClasses++;
}
}
for (String name : replacements) {
zout.putNextEntry(new ZipEntry(name));
zout.write(replacementFiles.get(name));
zout.closeEntry();
}
boolean hasPatchHash = expectedPatchHashes.containsKey(zipFile);
zout.putNextEntry(new ZipEntry(hashFileName));
String patchHash = hasPatchHash ? String.valueOf(expectedPatchHashes.get(zipFile)) : "-1";
zout.write(patchHash.getBytes("UTF-8"));
if (hasPatchHash) {
Log.info("Patched " + replacements.size() + " classes in " + zipFile + ", patchHash: " + patchHash + ", " + (onlyClasses ? "mod" : "server jar"));
}
zin.close();
zout.close();
return patchedClasses;
}
}