TransformableClassLoaderWhisperer - Add `ClassFileTransformer`s on the fly without Java agent

TransformableClassLoaderWhisperer.java

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.reflect.Field;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.Enumeration;
import java.util.List;
import java.util.Vector;
import java.util.concurrent.CopyOnWriteArrayList;

import sun.misc.Resource;
import sun.misc.URLClassPath;

/**
 * @author serkan
 */
public class TransformableClassLoaderWhisperer {

    private static final String SKIPPED_CLASS_PREFIX = 
            TransformableClassLoaderWhisperer.class.getName().replace(".", "/");
    private static final Field URL_CLASSPATH_FIELD;
    
    static {
        Field urlClasspathField = null;
        try {
            urlClasspathField = URLClassLoader.class.getDeclaredField("ucp");
            urlClasspathField.setAccessible(true);
        } catch (Throwable t) {
            System.err.println(
                    "Unable to get URL classpath field. " + 
                    "So TransformableClassLoaderWhisperer won't not available to use!");
            t.printStackTrace();
        }
        URL_CLASSPATH_FIELD = urlClasspathField;
    }

    private final List<ClassFileTransformer> transformers = new CopyOnWriteArrayList<ClassFileTransformer>();
    
    public static TransformableClassLoaderWhisperer get() {
        ClassLoader classLoader = ClassLoader.getSystemClassLoader();
        if (classLoader instanceof URLClassLoader) {
            return get((URLClassLoader) classLoader);
        } else {
            throw new IllegalStateException("Only URLClassLoader's are supported!");
        }
    }
    
    public static TransformableClassLoaderWhisperer get(URLClassLoader classLoader) {
        if (URL_CLASSPATH_FIELD == null) {
            throw new IllegalStateException("TransformableClassLoaderWhisperer is not available to use!");
        }
        synchronized (classLoader) {
            try {
                URLClassPath urlClassPath = (URLClassPath) URL_CLASSPATH_FIELD.get(classLoader);
                if (urlClassPath instanceof TransformerAwareClasspath) {
                    return ((TransformerAwareClasspath) urlClassPath).whisperer;
                } else {
                    return new TransformableClassLoaderWhisperer(classLoader);
                }
            } catch (IllegalArgumentException e) {
                throw new IllegalStateException("Unable to inject TransformableClassLoaderWhisperer!", e);
            } catch (IllegalAccessException e) {
                throw new IllegalStateException("Unable to inject TransformableClassLoaderWhisperer!", e);
            }
        }
    }
    
    private TransformableClassLoaderWhisperer(URLClassLoader urlClassLoader) {
        try {
            URL_CLASSPATH_FIELD.set(
                    urlClassLoader, 
                    new TransformerAwareClasspath(
                           urlClassLoader,
                           (URLClassPath) URL_CLASSPATH_FIELD.get(urlClassLoader)));
        } catch (Throwable t) {
            throw new IllegalStateException(t);
        } 
    }

    public void addTransformer(ClassFileTransformer transformer) {
        transformers.add(transformer);
    }
    
    public void removeTransformer(ClassFileTransformer transformer) {
        transformers.remove(transformer);
    }
   
    private class TransformerAwareClasspath extends URLClassPath {

        private final TransformableClassLoaderWhisperer whisperer = TransformableClassLoaderWhisperer.this;
        private final URLClassLoader urlClassLoader;
        private final URLClassPath urlClassPath;
       
        private TransformerAwareClasspath(URLClassLoader urlClassLoader, URLClassPath urlClassPath) {
           super(new URL[]{});
           this.urlClassLoader = urlClassLoader;
           this.urlClassPath = urlClassPath;
        }

        @Override
        public URL[] getURLs() {
            return urlClassPath.getURLs();
        }
        
        @Override
        public synchronized void addURL(URL url) {
           urlClassPath.addURL(url);
        }
    
        @Override
        public URL checkURL(URL url) {
            return urlClassPath.checkURL(url);
        }
    
        @Override
        public synchronized List<IOException> closeLoaders() {
            return urlClassPath.closeLoaders();
        }
    
        @Override
        public URL findResource(String name, boolean check) {
            return urlClassPath.findResource(name, check);
        }
    
        @Override
        public Enumeration<URL> findResources(String name, boolean check) {
            return urlClassPath.findResources(name, check);
        }
        
        @Override
        public Resource getResource(String name) {
            return processResource(urlClassPath.getResource(name));
        }
    
        @Override
        public Resource getResource(String name, boolean check) {
            return processResource(urlClassPath.getResource(name, check));
        }

        @Override
        public Enumeration<Resource> getResources(String name) {
            return processResources(urlClassPath.getResources(name));
        }
        
        @Override
        public Enumeration<Resource> getResources(String name, boolean check) {
            return processResources(urlClassPath.getResources(name, check));
        }

        private Resource processResource(final Resource resource) {
            if (resource == null || resource.getName().startsWith(SKIPPED_CLASS_PREFIX)) {
                return resource;
            }
            try {
                if (resource.getName().endsWith(".class")) {
                    String resourceName = resource.getName();
                    resourceName = resourceName.substring(0, resourceName.length() - ".class".length());
                    byte[] resourceData = resource.getBytes();
                    for (ClassFileTransformer transformer : transformers) {
                        byte[] transformedData;
                        try {
                            transformedData = transformer.transform(
                                    urlClassLoader, resourceName, null, null, resourceData);
                        } catch (IllegalClassFormatException e) {
                            e.printStackTrace();
                            continue;
                        }
                        if (transformedData != null) {
                            resourceData = transformedData;
                        }
                    }
                    final byte[] resourceDataArray = resourceData;
                    final int resourceDataLength = resourceData.length;
                    final ByteArrayInputStream resourceDataStream = new ByteArrayInputStream(resourceData);
                    return new Resource() {
                        @Override
                        public URL getURL() {
                            return resource.getURL();
                        }
                    
                        @Override
                        public String getName() {
                            return resource.getName();
                        }
                    
                        @Override
                        public InputStream getInputStream() throws IOException {
                            return resourceDataStream;
                        }
                        
                        @Override
                        public byte[] getBytes() throws IOException {
                            return resourceDataArray;
                        }
                    
                        @Override
                        public int getContentLength() throws IOException {
                            return resourceDataLength;
                        }
                    
                        @Override
                        public URL getCodeSourceURL() {
                            return resource.getCodeSourceURL();
                        }
                    };
                } 
                return resource;
            } catch (IOException e) {
                return resource;
            }
        }
        
        private Enumeration<Resource> processResources(Enumeration<Resource> resources) {
            Vector<Resource> processedResources = new Vector<Resource>();
            while (resources.hasMoreElements()) {
                processedResources.add(processResource(resources.nextElement()));
            }
            return processedResources.elements();
        }

   }

}