Javassist: How to add agent to classpath?

Question

I have an agent that I'm dynamically loading into a running Java application which opens a simple Swing JFrame when it's attached. It also allows appending new lines into a TextArea inside that JFrame.

My goal is to change how some methods work inside the application the agent is loaded into.

public class MyAgent {
    public static void agentmain(String args, Instrumentation instrumentation) {
        UI.openWindow();
        UI.addMessage("Agent loaded: %s", args);
        instrumentation.addTransformer(new MyTransformer());
        instrumentation.redefineClasses(new ClassDefinition(Class.forName("app.TargetClass"), ...));
    }
}

The UI window is managed in another class accessible from the agent. It successfully opens a window and appends a text message when the agent is loaded.

public class UI {
    private static SwingWindow swingWindow;
    
    public static void addMessage(String format, Object... args) {
        System.out.println("UI: " + String.format(format, args));
        swingWindow.appendToTextArea(format, args);
    }
    
    public static void openWindow() {
        try {
            SwingUtilities.invokeAndWait(() -> swingWindow = new SwingWindow());
        }
        catch (Exception e) {}
    }
}

I'm using Javassist to generate bytecode inside my transformer.

public class MyTransformer implements ClassFileTransformer {
    @Override
    public byte[] transform(ClassLoader loader, String className, ..., byte[] classBuffer) {
        if (className.equals("app/TargetClass")) {
            UI.addMessage("Now transforming the class that I need!");

            try {
                ClassPool classPool = ClassPool.getDefault();
                CtClass targetClass = classPool.get("app.TargetClass");
                CtMethod targetMethod = targetClass.getDeclaredMethod("importantMethod");
                targetMethod.insertBefore("me.domain.agent.ui.UI.addMessage(\"Hello from Javassist!\")");
                byte[] byteCode = targetClass.toBytecode();
                targetClass.detach();
                return byteCode;
            }
            catch (Exception e) {
                UI.addMessage("Couldn't transform the class I needed.");
            }
        }

        return classBuffer;
    }
}

The target class gets found, but the bytecode does not compile:

UI: Failed transforming class app.TargetClass: [source error] no such class: me$domain.agent.ui.UI

However, the UI class is inside the agent:

agent.jar
├── META-INF
└── me.domain.agent
    ├── ui
    │   └── UI.class
    ├── MyTransformer.class
    └── Agent.class

I've tried adding the agent's ClassLoader to Javassist's ClassPool:

classPool.insertClassPath(new LoaderClassPath(Agent.class.getClassLoader()));

But it doesn't work. How can I add a call to my agent UI into the bytecode?

Answer 1

I've decided to use ASM to call my agent UI from bytecode. There are no issues with finding the class.

Here's how the ASM-based class transformer looks like:

public class MyTransformer implements ClassFileTransformer {
    public void transformClass(ClassNode classNode) {
        MethodNode methodNode = findMethodNodeOfClass(classNode, "importantMethod", "()V");
        if (methodNode == null) {
            throw new TransformerException("app.TargetClass#importantMethod not found");
        }

        AbstractInsnNode firstInsn = findFirstInstruction(methodNode);
        if (firstInsn == null) {
            throw new TransformerException("No instructions in app.TargetClass#importantMethod");
        }

        InsnList insnList = new InsnList();
        insnList.add(new LdcInsnNode("Hello from ASM!"));
        insnList.add(new MethodInsnNode(INVOKESTATIC, Type.getInternalName(UI.class), "addMessage", "(Ljava/lang/String;)V"));
        methodNode.instructions.insertBefore(firstInsn, insnList);
    }

    @Override
    public byte[] transform(ClassLoader classLoader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] byteCode) {
        if (className.equals("app/TargetClass")) {
            try {
                ClassNode classNode = new ClassNode();
                ClassReader classReader = new ClassReader(byteCode);
                classReader.accept(classNode, 0);
                this.transformClass(classNode);
                ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES);
                classNode.accept(classWriter);
                return classWriter.toByteArray();
            }
            catch (Exception e) {
                e.printStackTrace();
            }
        }

        return byteCode;
    }
}

This will insert the static call UI.addMessage("Hello from ASM!") before the first non-label instruction of the app.TargetClass#importantMethod .