Search code examples
javaclassmetaprogrammingclassloaderbytecode

Can Java Classloader's rewrite the bytecode of (only their copy of) System classes?


So I have a classloader (MyClassLoader) that maintains a set of "special" classes in memory. These special classes are dynamically compiled and stored in a byte array inside MyClassLoader. When the MyClassLoader is asked for a class, it first checks if its specialClasses dictionary contains it, before delegating to the System classloader. It looks something like this:

class MyClassLoader extends ClassLoader {
    Map<String, byte[]> specialClasses;

    public MyClassLoader(Map<String, byte[]> sb) {
        this.specialClasses = sb;
    }

    @Override
    public Class<?> loadClass(String name) throws ClassNotFoundException {
        if (specialClasses.containsKey(name)) return findClass(name);
        else return super.loadClass(name);
    }

    @Override
    public Class findClass(String name) {
        byte[] b = specialClasses.get(name);
        return defineClass(name, b, 0, b.length);
    }    
}

If I want to perform transformations (e.g. instrumentation) on the specialClasses, I can do so simply by modifying the byte[] before i call defineClass() on it.

I would also like to transform the classes which are provided by the System classloader, but the System classloader doesn't seem to provide any way of accessing the raw byte[] of the classes it provides, and gives me the Class objects directly.

I could use a -javaagent instrument all classes loaded into the JVM, but that would add overhead to the classes which I do not want to instrument; I only really want the classes loaded by MyClassLoader to be instrumented.

  • Is there any way of retrieving the raw byte[] of the classes provided by the parent classloader, so I can instrument them before defining my own copy?
  • Alternately, is there any way of emulating the functionality of the System classloader, in terms of where it grabs it's byte[]s from, so that MyClassLoader can instrument and define its own copy of all the System classes (Object, String, etc.)?

EDIT:

So I tried another approach:

  • Using a -javaagent, Capture the byte[] of every class that is loaded and store it in a hashtable, keyed by the name of the class.
  • MyClassLoader, instead of delegating the system classes to its parent classloader, would instead load their bytecode from this hashtable using the class name and define it

In theory this would let MyClassLoader define its own version of the system classes with instrumentation. However, it fails with a

java.lang.SecurityException: Prohibited package name: java.lang

Clearly the JVM does not like me defining the java.lang classes myself, even though it should (in theory) be from the same byte[] source that the bootstrap-loaded classes should come from. The search for a solution continues.

EDIT2:

I found a (really sketchy) solution for this problem, but if someone who knows more than me about the intricacies of the Java classloading/instrumentation can come up with something less sketchy, that would be awesome.


Solution

  • So I have found a solution to this. It's not a very elegant solution, and it would cause a lot of angry emails at code-review time, but it seems to work. The basic points are:

    JavaAgent

    Use a java.lang.instrumentation and a -javaagent to store the Instrumentation object to use later

    class JavaAgent {
        private JavaAgent() {}
    
        public static void premain(String agentArgs, Instrumentation inst) {
            System.out.println("Agent Premain Start");
            Transformer.instrumentation = inst;
            inst.addTransformer(new Transformer(), inst.isRetransformClassesSupported());
        }    
    }
    

    ClassFileTransformer

    Add a Transformer to the Instrumentation that only works on the marked classes. Something like

    public class Transformer implements ClassFileTransformer {
        public static Set<Class<?>> transformMe = new Set<>()
        public static Instrumentation instrumentation = null; // set during premain()
        @Override
        public byte[] transform(ClassLoader loader,
                                String className,
                                Class<?> classBeingRedefined,
                                ProtectionDomain protectionDomain,
                                byte[] origBytes) {
    
    
            if (transformMe.contains(classBeingRedefined)) {
                return instrument(origBytes, loader);
            } else {
                return null;
            }
        }
        public byte[] instrument(byte[] origBytes) {
            // magic happens here
        }
    }
    

    ClassLoader

    In the classloader, explicitly mark each loaded class (even the classes whose loading is delegated to the parent) by placing it in transformMe before asking the Instrumentation to transform it

    public class MyClassLoader extends ClassLoader{
        public Class<?> instrument(Class<?> in){
            try{
                Transformer.transformMe.add(in);
                Transformer.instrumentation.retransformClasses(in);
                Transformer.transformMe.remove(in);
                return in;
            }catch(Exception e){ return null; }
        }
        @Override
        public Class<?> loadClass(String name) throws ClassNotFoundException {
            return instrument(super.loadClass(name));
        }
    }
    

    ... and voila! Every class which is loaded by MyClassLoader gets transformed by the instrument() method, including all the system classes like java.lang.Object and friends, while all the classes which are loaded by the default ClassLoader are left untouched.

    I have tried this using a memory-profiling instrument() method, which inserts callback hooks to track memory allocations in the instrumented bytecode, and can confirm that the MyClassLoad classes are firing the callbacks when their methods run (even system classes) while the "normal" classes are not.

    Victory!

    This is, of course, terrible code. Shared mutable state everywhere, non-local side-effects, globals, everything that you can possibly imagine. Probably isn't threadsafe either. But it shows that such a thing is possible, you can indeed selectively instrument the bytecode of classes, even system classes, as part of a custom ClassLoader's operation, while leaving the "rest" of the program untouched.

    Open Problems

    If someone else has any ideas how to make this code less terrible, I would be glad to hear it. I could not figure out a way of:

    • Making the Instrumentation only instrument classes on demand via retransformClasses() and not instrument classes loaded otherwise
    • Store some metadata in each Class<?> object which would allow the Transformer to tell whether it should be transformed or not, without the global-mutable-hashtable lookup.
    • Transform a system class without using the Instrumentation.retransformClass() method. As mentioned, any attempts to dynamically defineClass a byte[] into a java.lang.* class fails due to hardcoded checks in ClassLoader.java.

    If anyone can find a way around any of these problems, it would make this much less sketchy. Regardless, I'm guessing being able to instrument (e.g. for profiling) some sub-system (i.e. the one you're interested in) while leaving the rest of the JVM untouched (with no instrumentation overhead) will be useful to someone else besides me, so here it is.