Implementing a selective ClassLoader

I want to instrument the bytecode of some classes on the classpath at loading time. Since these are 3rd party libraries, I know exactly when they are loaded. The problem is that I need to do the instrumentation selectively, i.e. instrument only some classes. Now if I do not load a class with my classloader but with its parent, this parent gets set as the classes classloader and all succinct classes are loaded by that parent, effectively putting my classloader out of use. So I need to implement a parent-last classloader (see How to put custom ClassLoader to use?).

So I need to load classes myself. If those classes are system classes (starting with "java" or "sun") I delegate to the parent. Otherwise I read the bytecode and call defineClass(name, byteBuffer, 0, byteBuffer.length);. But now a java.lang.ClassNotFoundException: java.lang.Object is thrown.

Here is the code, any comment highly appreciated:

public class InstrumentingClassLoader extends ClassLoader {
private final BytecodeInstrumentation instrumentation = new BytecodeInstrumentation();

@Override
public Class<?> loadClass(String name) throws ClassNotFoundException {
    Class<?> result = defineClass(name);
    if (result != null) {
        return result;
    }
    result = findLoadedClass(name);
    if(result != null){
        return result;
    }
    result = super.findClass(name);
    return result;
}

private Class<?> defineClass(String name) throws ClassFormatError {
    byte[] byteBuffer = null;
    if (instrumentation.willInstrument(name)) {
        byteBuffer = instrumentByteCode(name);
    }
    else {
        byteBuffer = getRegularByteCode(name);
    }
    if (byteBuffer == null) {
        return null;
    }
    Class<?> result = defineClass(name, byteBuffer, 0, byteBuffer.length);
    return result;
}

private byte[] getRegularByteCode(String name) {
    if (name.startsWith("java") || name.startsWith("sun")) {
        return null;
    }
    try {
        InputStream is = ClassLoader.getSystemResourceAsStream(name.replace('.', '/') + ".class");
        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        int nRead;
        byte[] data = new byte[16384];
        while ((nRead = is.read(data, 0, data.length)) != -1) {
            buffer.write(data, 0, nRead);
        }
        buffer.flush();
        return buffer.toByteArray();
    } catch (IOException exc) {
        return null;
    }
}

private byte[] instrumentByteCode(String fullyQualifiedTargetClass) {
    try {
        String className = fullyQualifiedTargetClass.replace('.', '/');
        return instrumentation.transformBytes(className, new ClassReader(fullyQualifiedTargetClass));
    } catch (Exception e) {
        throw new RuntimeException(e);
    }
}
}

The code can be executed e.g. with:

    InstrumentingClassLoader instrumentingClassLoader = new InstrumentingClassLoader();
    Class<?> changedClass = instrumentingClassLoader.loadClass(ClassLoaderTestSubject.class.getName());

The ClassLoaderTestSubject should call some other classes, where the 开发者_JAVA百科called classes are target of instrumentation, but the ClassLoaderTestSubject itself is not...

---

I'd recommend you to use regular class loader strategy, i.e. parent first. But put all classes that you want to instrument into separate jar file and do not add it to the classpath of the application. Instantiate these classes using your class loader that extends URL class loader and knows to search jars in other location. In this case all JDK classes will be known automatically and your code will be simpler. You do not have to "think" whether to instrument the class: if it is not loaded by parent class loader it is your class that has to be instrumented.

---

Stupid mistake. The parent classloader is not the parent as in the inheritance hierarchy. It is the parent as given to the constructor. So the correct code looks like this:

public InstrumentingClassLoader() {
    super(InstrumentingClassLoader.class.getClassLoader());
    this.classLoader = InstrumentingClassLoader.class.getClassLoader();
}

@Override
public Class<?> loadClass(String name) throws ClassNotFoundException {
    [... as above ...]
    result = classLoader.loadClass(name);
    return result;
}