LaunchClassLoader.java

package net.minecraft.launchwrapper;

import cpw.mods.fml.relauncher.FMLRelaunchLog;
import nallar.tickthreading.patcher.PatchHook;
import nallar.unsafe.UnsafeUtil;

import java.io.*;
import java.net.*;
import java.security.*;
import java.util.*;
import java.util.jar.*;
import java.util.jar.Attributes.*;
import java.util.logging.*;

public class LaunchClassLoader extends URLClassLoader {
	private IClassTransformer prePatchTransformer;
	public static final int BUFFER_SIZE = 1 << 12;
	private List<URL> sources;
	private ClassLoader parent = getClass().getClassLoader();

	private List<IClassTransformer> transformers = new ArrayList<IClassTransformer>(2);
	private Map<String, Class<?>> cachedClasses = new HashMap<String, Class<?>>(1000);
	private Set<String> invalidClasses = new HashSet<String>(1000);

	private Set<String> classLoaderExceptions = new HashSet<String>();
	private Set<String> transformerExceptions = new HashSet<String>();
	private Map<String, byte[]> resourceCache = new HashMap<String, byte[]>(1000);
	private Set<String> negativeResourceCache = new HashSet<String>();

	private IClassNameTransformer renameTransformer;

	private final ThreadLocal<byte[]> loadBuffer = new ThreadLocal<byte[]>();

	private static final String[] RESERVED_NAMES = {"CON", "PRN", "AUX", "NUL", "COM1", "COM2", "COM3", "COM4", "COM5", "COM6", "COM7", "COM8", "COM9", "LPT1", "LPT2", "LPT3", "LPT4", "LPT5", "LPT6", "LPT7", "LPT8", "LPT9"};

	private static final boolean DEBUG = Boolean.parseBoolean(System.getProperty("legacy.debugClassLoading", "false"));
	private static final boolean DEBUG_FINER = DEBUG && Boolean.parseBoolean(System.getProperty("legacy.debugClassLoadingFiner", "false"));
	private static final boolean DEBUG_SAVE = DEBUG && Boolean.parseBoolean(System.getProperty("legacy.debugClassLoadingSave", "false"));
	private static File tempFolder = null;
	public static LaunchClassLoader instance;

	public LaunchClassLoader(URL[] sources) {
		super(sources, null);
		instance = this;
		this.sources = new ArrayList<URL>(Arrays.asList(sources));
		Thread.currentThread().setContextClassLoader(this);

		// classloader exclusions
		addClassLoaderExclusion("java.");
		addClassLoaderExclusion("sun.");
		addClassLoaderExclusion("org.lwjgl.");
		addClassLoaderExclusion("org.apache.logging.");
		addClassLoaderExclusion("net.minecraft.launchwrapper.");

		// transformer exclusions
		addTransformerExclusion("javax.");
		addTransformerExclusion("argo.");
		addTransformerExclusion("org.objectweb.asm.");
		addTransformerExclusion("com.google.common.");
		addTransformerExclusion("org.bouncycastle.");
		addTransformerExclusion("net.minecraft.launchwrapper.injector.");

		if (DEBUG_SAVE) {
			int x = 1;
			tempFolder = new File(Launch.minecraftHome, "CLASSLOADER_TEMP");
			while (tempFolder.exists() && x <= 10) {
				tempFolder = new File(Launch.minecraftHome, "CLASSLOADER_TEMP" + x++);
			}

			if (tempFolder.exists()) {
				LogWrapper.info("DEBUG_SAVE enabled, but 10 temp directories already exist, clean them and try again.");
				tempFolder = null;
			} else {
				LogWrapper.info("DEBUG_SAVE Enabled, saving all classes to \"%s\"", tempFolder.getAbsolutePath().replace('\\', '/'));
				tempFolder.mkdirs();
			}
		}
	}

	public void registerTransformer(String transformerClassName) {
		try {
			IClassTransformer transformer = (IClassTransformer) loadClass(transformerClassName).newInstance();
			transformers.add(transformer);
			if (transformer instanceof IClassNameTransformer && renameTransformer == null) {
				renameTransformer = (IClassNameTransformer) transformer;
			}
			if (transformerClassName.equals("cpw.mods.fml.common.asm.transformers.DeobfuscationTransformer")) {
				prePatchTransformer = transformer;
			}
		} catch (Exception e) {
			LogWrapper.log(Level.SEVERE, "Critical problem occurred registering the ASM transformer class %s", transformerClassName);
		}
	}

	public boolean excluded(String name) {
		for (final String exception : classLoaderExceptions) {
			if (name.startsWith(exception)) {
				return true;
			}
		}
		return false;
	}

	@Override
	public Class<?> findClass(final String name) throws ClassNotFoundException {
		if (invalidClasses.contains(name)) {
			throw new ClassNotFoundException(name);
		}

		if (excluded(name)) {
			return parent.loadClass(name);
		}

		if (cachedClasses.containsKey(name)) {
			return cachedClasses.get(name);
		}

		for (final String exception : transformerExceptions) {
			if (name.startsWith(exception)) {
				try {
					final Class<?> clazz = super.findClass(name);
					cachedClasses.put(name, clazz);
					return clazz;
				} catch (ClassNotFoundException e) {
					invalidClasses.add(name);
					throw e;
				}
			}
		}

		try {
			final String transformedName = transformName(name);
			if (cachedClasses.containsKey(transformedName)) {
				return cachedClasses.get(transformedName);
			}

			final String untransformedName = untransformName(name);

			final int lastDot = untransformedName.lastIndexOf('.');
			final String packageName = lastDot == -1 ? "" : untransformedName.substring(0, lastDot);
			final String fileName = untransformedName.replace('.', '/').concat(".class");
			URLConnection urlConnection = findCodeSourceConnectionFor(fileName);

			CodeSigner[] signers = null;

			if (lastDot > -1 && !untransformedName.startsWith("net.minecraft.")) {
				if (urlConnection instanceof JarURLConnection) {
					final JarURLConnection jarURLConnection = (JarURLConnection) urlConnection;
					final JarFile jarFile = jarURLConnection.getJarFile();

					if (jarFile != null && jarFile.getManifest() != null) {
						final Manifest manifest = jarFile.getManifest();
						final JarEntry entry = jarFile.getJarEntry(fileName);

						Package pkg = getPackage(packageName);
						getClassBytes(untransformedName);
						signers = entry.getCodeSigners();
						if (pkg == null) {
							definePackage(packageName, manifest, jarURLConnection.getJarFileURL());
						} else {
							if (pkg.isSealed() && !pkg.isSealed(jarURLConnection.getJarFileURL())) {
								LogWrapper.severe("The jar file %s is trying to seal already secured path %s", jarFile.getName(), packageName);
							} else if (isSealed(packageName, manifest)) {
								LogWrapper.severe("The jar file %s has a security seal for path %s, but that path is defined and not secure", jarFile.getName(), packageName);
							}
						}
					}
				} else {
					Package pkg = getPackage(packageName);
					if (pkg == null) {
						definePackage(packageName, null, null, null, null, null, null, null);
					} else if (pkg.isSealed()) {
						LogWrapper.severe("The URL %s is defining elements for sealed path %s", urlConnection.getURL(), packageName);
					}
				}
			}

			final byte[] transformedClass = runTransformers(untransformedName, transformedName, getClassBytes(untransformedName));
			if (DEBUG_SAVE) {
				saveTransformedClass(transformedClass, transformedName);
			}

			final CodeSource codeSource = urlConnection == null ? null : new CodeSource(urlConnection.getURL(), signers);
			final Class<?> clazz = defineClass(transformedName, transformedClass, 0, transformedClass.length, codeSource);
			cachedClasses.put(transformedName, clazz);
			return clazz;
		} catch (Throwable e) {
			invalidClasses.add(name);
			if (DEBUG) {
				LogWrapper.log(Level.FINE, e, "Exception encountered attempting classloading of %s", name);
			}
			throw new ClassNotFoundException(name, e);
		}
	}

	private void saveTransformedClass(final byte[] data, final String transformedName) {
		if (tempFolder == null) {
			return;
		}

		final File outFile = new File(tempFolder, transformedName.replace('.', File.separatorChar) + ".class");
		final File outDir = outFile.getParentFile();

		if (!outDir.exists()) {
			outDir.mkdirs();
		}

		if (outFile.exists()) {
			outFile.delete();
		}

		try {
			LogWrapper.fine("Saving transformed class \"%s\" to \"%s\"", transformedName, outFile.getAbsolutePath().replace('\\', '/'));

			final OutputStream output = new FileOutputStream(outFile);
			output.write(data);
			output.close();
		} catch (IOException ex) {
			LogWrapper.log(Level.WARNING, ex, "Could not save transformed class \"%s\"", transformedName);
		}
	}

	private String untransformName(final String name) {
		if (renameTransformer != null) {
			return renameTransformer.unmapClassName(name);
		}

		return name;
	}

	private String transformName(final String name) {
		if (renameTransformer != null) {
			return renameTransformer.remapClassName(name);
		}

		return name;
	}

	private boolean isSealed(final String path, final Manifest manifest) {
		Attributes attributes = manifest.getAttributes(path);
		String sealed = null;
		if (attributes != null) {
			sealed = attributes.getValue(Name.SEALED);
		}

		if (sealed == null) {
			attributes = manifest.getMainAttributes();
			if (attributes != null) {
				sealed = attributes.getValue(Name.SEALED);
			}
		}
		return "true".equalsIgnoreCase(sealed);
	}

	private URLConnection findCodeSourceConnectionFor(final String name) {
		final URL resource = findResource(name);
		if (resource != null) {
			try {
				return resource.openConnection();
			} catch (IOException e) {
				throw new RuntimeException(e);
			}
		}

		return null;
	}

	private byte[] runTransformer(final String name, final String transformedName, byte[] basicClass, final IClassTransformer transformer) {
		try {
			return transformer.transform(name, transformedName, basicClass);
		} catch (Throwable t) {
			String message = t.getMessage();
			if (message != null && message.contains("for invalid side")) {
				throw UnsafeUtil.throwIgnoreChecked(t);
			} else if (basicClass != null || DEBUG_FINER) {
				FMLRelaunchLog.log((DEBUG_FINER && basicClass != null) ? Level.WARNING : Level.FINE, t, "Failed to transform " + name);
			}
			return basicClass;
		}
	}

	private byte[] runTransformers(final String name, final String transformedName, byte[] basicClass) {
		basicClass = transformUpToSrg(name, transformedName, basicClass);
		basicClass = PatchHook.hook(name, transformedName, basicClass);
		basicClass = transformAfterSrg(name, transformedName, basicClass);
		return basicClass;
	}

	private byte[] transformUpToSrg(final String name, final String transformedName, byte[] basicClass) {
		for (final IClassTransformer transformer : transformers) {
			basicClass = runTransformer(name, transformedName, basicClass, transformer);
			if (transformer == prePatchTransformer) {
				return basicClass;
			}
		}
		throw new RuntimeException("No SRG transformer!");
	}

	private byte[] transformAfterSrg(final String name, final String transformedName, byte[] basicClass) {
		boolean pastSrg = false;
		for (final IClassTransformer transformer : transformers) {
			if (pastSrg) {
				basicClass = runTransformer(name, transformedName, basicClass, transformer);
			}
			if (transformer == prePatchTransformer) {
				pastSrg = true;
			}
		}
		if (!pastSrg) {
			throw new RuntimeException("No SRG transformer!");
		}
		return basicClass;
	}

	public byte[] getSrgBytes(String name) {
		final String transformedName = transformName(name);
		name = untransformName(name);
		try {
			byte[] bytes = getClassBytes(name);
			return transformUpToSrg(name, transformedName, bytes);
		} catch (Throwable t) {
			throw new RuntimeException(t);
		}
	}

	@Override
	public void addURL(final URL url) {
		super.addURL(url);
		sources.add(url);
	}

	public List<URL> getSources() {
		return sources;
	}

	private byte[] readFully(InputStream stream) {
		try {
			byte[] buffer = getOrCreateBuffer();

			int read;
			int totalLength = 0;
			while ((read = stream.read(buffer, totalLength, buffer.length - totalLength)) != -1) {
				totalLength += read;

				// Extend our buffer
				if (totalLength >= buffer.length - 1) {
					byte[] newBuffer = new byte[buffer.length + BUFFER_SIZE];
					System.arraycopy(buffer, 0, newBuffer, 0, buffer.length);
					buffer = newBuffer;
				}
			}

			final byte[] result = new byte[totalLength];
			System.arraycopy(buffer, 0, result, 0, totalLength);
			return result;
		} catch (Throwable t) {
			LogWrapper.log(Level.WARNING, t, "Problem loading class");
			return new byte[0];
		}
	}

	private byte[] getOrCreateBuffer() {
		byte[] buffer = loadBuffer.get();
		if (buffer == null) {
			loadBuffer.set(new byte[BUFFER_SIZE]);
			buffer = loadBuffer.get();
		}
		return buffer;
	}

	public List<IClassTransformer> getTransformers() {
		return Collections.unmodifiableList(transformers);
	}

	public void addClassLoaderExclusion(String toExclude) {
		classLoaderExceptions.add(toExclude);
	}

	public void addTransformerExclusion(String toExclude) {
		transformerExceptions.add(toExclude);
	}

	public byte[] getClassBytes(String name) throws IOException {
		if (negativeResourceCache.contains(name)) {
			return null;
		} else if (resourceCache.containsKey(name)) {
			return resourceCache.get(name);
		}
		if (name.indexOf('.') == -1) {
			for (final String reservedName : RESERVED_NAMES) {
				if (name.toUpperCase(Locale.ENGLISH).startsWith(reservedName)) {
					final byte[] data = getClassBytes("_" + name);
					if (data != null) {
						resourceCache.put(name, data);
						return data;
					}
				}
			}
		}

		InputStream classStream = null;
		try {
			final String resourcePath = name.replace('.', '/').concat(".class");
			final URL classResource = findResource(resourcePath);

			if (classResource == null) {
				if (DEBUG) LogWrapper.finest("Failed to find class resource %s", resourcePath);
				negativeResourceCache.add(name);
				return null;
			}
			classStream = classResource.openStream();

			if (DEBUG) LogWrapper.finest("Loading class %s from resource %s", name, classResource.toString());
			final byte[] data = readFully(classStream);
			resourceCache.put(name, data);
			return data;
		} finally {
			closeSilently(classStream);
		}
	}

	private static void closeSilently(Closeable closeable) {
		if (closeable != null) {
			try {
				closeable.close();
			} catch (IOException ignored) {
			}
		}
	}

	public void clearNegativeEntries(Set<String> entriesToClear) {
		negativeResourceCache.removeAll(entriesToClear);
	}
}