Question

For my instrumentation tool, I want to provide a wrapping ClassLoader that is used to start a main method after specific classes have been instrumented. My ClassLoader should load instrumented versions of certain classes. But for Jetty and JUnit, this approach is severly limited because they build their own classloading hierarchy.

I don’t want to pass VM arguments, so I can’t change the SystemClassLoader. But I can force-feed it with my classes by using reflection to make ClassLoader.defineClass(String, byte[], int, int) public.

ClassLoader scl = ClassLoader.getSystemClassLoader();
Method defineClass = ClassLoader.class.getDeclaredMethod(
        "defineClass", String.class, byte[].class, int.class, int.class);
defineClass.setAccessible(true);
for (String binaryName : classNamesToLoad) {
    byte[] bytecode = this.declaredClasses.get(binaryName);
    defineClass.invoke(scl, binaryName, bytecode, 0, bytecode.length);
}
defineClass.setAccessible(false);

This is just great – but there’s one problem left: If some of my classes inherit from or contain other classes, they have to be loaded in the right order because the SystemClassLoader loads all classes the current one depends on – and would load the uninstrumented version.

Here is an example with some (poorly named) classes and the order they would have to be loaded in:

A
A.A extends B.A
B
B.A extends B.C
B.C

would have to be loaded in order

B
B.C
B.A
A
A.A

if I want to load only the instrumented version.

Is there an easy way out – e.g. a “setSystemClassLoader” method I didn’t spot yet?

A workaround by which I wouldn’t need to manipulate the SystemClassLoader?

Or do I really have to do a full transitive dependency analysis starting on the classes I want to load to determine the right order (and in this case: is there any “prior art” I can work with)?

Thanks!

Answer 1

Looks like there’s no way around the transitive dependency analysis.

I solved it this way, and I really hope someone can profit from this implementation:

import java.io.IOException;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.tree.ClassNode;

public class DependencyDetector {

    private static class Node implements Comparable<Node> {
        private final String binaryName;
        private final Node[] imports;
        private final int score;

        private Node(String binaryName, Node...imports) {
            this.binaryName = binaryName;
            this.imports = imports;
            this.score = calculateScore(imports);
        }

        public int compareTo(Node o) {
            return score - o.score;
        }

        private int calculateScore(Node...imports) {
            int newScore = 0;
            for (Node n : imports) {
                if (n.score >= newScore) {
                    newScore = n.score + 1;
                }
            }
            return newScore;
        }
    }

    private Map<String, Node> nodes = new HashMap<String, Node>();

    public DependencyDetector add(ClassNode node) {
        Node n = nodes.get(node.name);
        if (n == null) {
            n = createNode(node);
        }
        return this;
    }

    private Node createNode(ClassNode node) {
        String binaryName = node.name;
        String[] importNames = extractImportedBinaryNames(node);
        Node[] imports = new Node[importNames.length];
        for (int i = 0; i < imports.length; i++) {
            String importName = importNames[i];
            Node imp = nodes.get(importName);
            if (imp == null) {
                ClassNode cn = new ClassNode();
                String path = importName.replace('.', '/') + ".class";
                try {
                    new ClassReader(
                            ClassLoader.getSystemResourceAsStream(path)
                        ).accept(cn, ClassReader.SKIP_CODE);
                } catch (IOException e) {
                    throw new RuntimeException(
                        "could not read class " + importName);
                }
                imp = createNode(cn);
                nodes.put(importName, imp);
            }
            imports[i] = imp;
        }
        Node result = new Node(binaryName, imports);
        nodes.put(binaryName, result);
        return result;
    }

    private String[] extractImportedBinaryNames(ClassNode node) {
        String binaryName = node.name;
        ArrayList<String> nodesToAdd = new ArrayList<String>();
        int endOfOuter = binaryName.lastIndexOf('$');
        if (endOfOuter >= 0) {
            nodesToAdd.add(binaryName.substring(0, endOfOuter));
        }
        if (node.superName != null) {
            nodesToAdd.add(node.superName);
        }
        if (node.interfaces != null) {
            for (String interf : (List<String>) node.interfaces) {
                if (interf != null) {
                    nodesToAdd.add(interf);
                }
            }
        }
        return nodesToAdd.toArray(new String[nodesToAdd.size()]);
    }

    public String[] getClassesToLoad(String...binaryNames) {
        String[] classNames = binaryNames != null && binaryNames.length > 0
                ? binaryNames.clone()
                : nodes.keySet().toArray(new String[nodes.size()]);
        ArrayDeque<Node> dependencyQueue = new ArrayDeque<Node>();
        for (String className : classNames) {
            Node node = nodes.get(className.replace('.', '/'));
            dependencyQueue.add(node);
            if (node == null) {
                throw new RuntimeException(
                    "Class " + className + " was not registered");
            }
        }
        HashMap<String, Node> dependencyMap = new HashMap<String, Node>();
        while (!dependencyQueue.isEmpty()) {
            Node node = dependencyQueue.removeFirst();
            dependencyMap.put(node.binaryName, node);
            for (Node i : node.imports) {
                dependencyQueue.addLast(i);
            }
        }
        ArrayList<Node> usedNodes =
            new ArrayList<Node>(dependencyMap.values());
        Collections.sort(usedNodes);
        String[] result = new String[usedNodes.size()];
        int i = 0;
        for (Node n : usedNodes) {
            result[i++] = n.binaryName.replace('/', '.');
        }
        return result;
    }

    public boolean contains(String binaryName) {
        return nodes.containsKey(binaryName.replace('.', '/'));
    }
}

It’s used like this: on a DependencyDetector, you call add(ClassNode) to add a ClassNode and all its dependencies (all classes it extends or implements or is contained by). When you are done building the dependency tree, you call getClassesToLoad() to retrieve all dependencies as a String[] containing the binary names in the necessary order. You can also just ask for a subset of all added classes and their dependencies by specifying the binary names as a parameter of getClassesToLoad(...).

Now, when I instrument classes, I also add the ClassNode to the DependencyDetector and can retrieve everything I need to pass it into a method like this:

/**
 * load the specified classes (or all instrumented classes)
 * and all their dependencies with the specified ClassLoader.
 * @param loader
 * @param binaryNames binary names of all classes you want to load
 *        - none loads all instrumented classes
 */
public void loadIntoClassLoader(ClassLoader loader, String...binaryNames) {
    final String[] classNamesToLoad =
        dependencies.getClassesToLoad(binaryNames);
    Method defineClass = null;
    Method findLoadedClass = null;
    try {
        // crack ClassLoader wide open and force-feed it with our classes
        defineClass = ClassLoader.class.getDeclaredMethod(
                "defineClass", String.class, byte[].class,
                int.class, int.class);
        defineClass.setAccessible(true);
        findLoadedClass = ClassLoader.class.getDeclaredMethod(
                "findLoadedClass", String.class);
        findLoadedClass.setAccessible(true);
        for (String binaryName : classNamesToLoad) {
            if (!binaryName.startsWith("java.")) {
                if (findLoadedClass.invoke(loader, binaryName) == null) {
                    byte[] bytecode = getBytecode(binaryName);
                    defineClass.invoke(loader, binaryName, bytecode,
                        0, bytecode.length);
                } else if (declaredClasses.containsKey(binaryName)) {
                    throw new RuntimeException(
                        "Class " + binaryName + " was already loaded, " +
                        "it must not be redeclared");
                }
            }
        }
    } catch (Exception e) {
        e.printStackTrace();
        throw new RuntimeException(
            "could not load classes into ClassLoader", e);
    } finally {
        rehideMethod(findLoadedClass);
        rehideMethod(defineClass);
    }
}

private void rehideMethod(Method m) {
    if (m != null) {
        try {
            m.setAccessible(false);
        } catch (Exception e) {
        }
    }
}

which relies on

private final DependencyDetector dependencies = new DependencyDetector();
private final Map<String, byte[]> declaredClasses = new HashMap<String, byte[]>();

private byte[] getBytecode(String binaryName) {
    byte[] bytecode = declaredClasses.get(binaryName);
    if (bytecode == null) {
        // asBytes loads the class as byte[]
        bytecode =
            asBytes(binaryName.replace('.', '/') + ".class");
    }
    return bytecode;
}

That’s pretty much it and it works great in every situation I encountered so far.