lithium.classloadertest.common.TransformingURLClassLoader Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of multiverse-test Show documentation
The Multiverse Test Framework for Java
The newest version!
package lithium.classloadertest.common;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.net.URLClassLoader;
import java.security.AccessController;
import java.security.PrivilegedExceptionAction;

/**
 * @author jeff.collins
 */
public class TransformingURLClassLoader extends URLClassLoader {
	
	public interface ClassTransformer {
		byte[] transformClass(byte inputClass[]);
	}	
	
	final ClassTransformer transformer;

	public TransformingURLClassLoader(Class classesToFilter[], ClassTransformer transformer) {
		super(((URLClassLoader) TransformingURLClassLoader.class.getClassLoader()).getURLs(),
				TransformingURLClassLoader.class.getClassLoader().getParent());
		this.transformer = transformer;
		for (Class classToFilter : classesToFilter) {
			loadTransformedClass(classToFilter.getName());
		}
	}

	public TransformingURLClassLoader() {
		super(((URLClassLoader) TransformingURLClassLoader.class.getClassLoader()).getURLs(),
				TransformingURLClassLoader.class.getClassLoader().getParent());		
		this.transformer = null;
	}

		
	public Class loadTransformedClass(final String name)  {
		try {
		    return AccessController.doPrivileged(new PrivilegedExceptionAction>() {
			    @Override
				public Class run() throws ClassNotFoundException {
					String path = name.replace('.', '/').concat(".class");
					URL url = getResource(path);
					if (url != null) {
				    	byte b[] = transformer.transformClass(getBytes(url));
				    	return defineClass(name, b, 0, b.length);
					} else {
					    throw new ClassNotFoundException(name);
					}
			    }
			}, AccessController.getContext());
		} catch (java.security.PrivilegedActionException pae) {
		    throw new RuntimeException("Failed finding class", pae.getException());
		}
	}
	
	private byte[] getBytes(URL url) {
		ByteArrayOutputStream bais = new ByteArrayOutputStream();
		InputStream is = null;
		try {
			is = url.openStream();
			byte[] byteChunk = new byte[4096]; 
			int n;

			while ((n = is.read(byteChunk)) > 0) {
				bais.write(byteChunk, 0, n);
			}
		} catch (IOException e) {
			throw new RuntimeException("Failed while reading bytes from: " + url.toExternalForm(), e);
		} finally {
			if (is != null) {
				try {
					is.close();
				} catch (IOException e) {
				}
			}
		}
		return bais.toByteArray();
	}
	
}