javaclassloaderjava-bytecode-asmdynamic-class-loaders

Java ASM byte code manipulation to inject code into a method not working


I want to inject some code into an existing class/method. But I am unable to get the classloader to "find" the class in order to use the modified byte code.

MyClassInjector.java

import org.objectweb.asm.*;
    
public class MyClassInjector {
    public static void main(String[] args) throws Exception {
        // Load the MyClass class
        ClassReader cr = new ClassReader("MyClass");
        ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS);
        MyClassVisitor cv = new MyClassVisitor(cw);
        cr.accept(cv, 0);

        // Inject code into the myMethod method
        MethodVisitor mv = cv.visitMethod(Opcodes.ACC_PUBLIC, "myMethod", "()V", null, null);
        mv.visitMethodInsn(Opcodes.INVOKESTATIC, "MyClassInjector", "newMethod", "()V", false);
        mv.visitInsn(Opcodes.RETURN);
        mv.visitMaxs(0, 0);
        mv.visitEnd();

        // Define the newMethod method
        MethodVisitor mv2 = cv.visitMethod(Opcodes.ACC_PUBLIC + Opcodes.ACC_STATIC, "newMethod", "()V", null, null);
        mv2.visitCode();
        mv2.visitFieldInsn(Opcodes.GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
        mv2.visitLdcInsn("Injected code");
        mv2.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V", false);
        mv2.visitInsn(Opcodes.RETURN);
        mv2.visitMaxs(0, 0);
        mv2.visitEnd();

        // Define the new byte array with the modified class bytecode
        byte[] modifiedClass = cw.toByteArray();

        // Define a new class loader to load the modified class
        ClassLoader cl = new ClassLoader() {
            @Override
            protected Class<?> findClass(String name) throws ClassNotFoundException {
                if (name.equals("MyClass")) {
                    return defineClass(name, modifiedClass, 0, modifiedClass.length);
                } else {
                    return super.findClass(name);
                }
            }
        };

        // Load the modified class and call myMethod
        Class<?> myClass = cl.loadClass("MyClass"); <----------------------- HERE
        Object myObject = myClass.newInstance();
        myClass.getMethod("myMethod").invoke(myObject);
    }
}

class MyClassVisitor extends ClassVisitor {
    public MyClassVisitor(ClassVisitor cv) {
        super(Opcodes.ASM5, cv);
    }
}

MyClass.java

public class MyClass {
    public void myMethod() {
        System.out.println("Hello, world!");
    }
}

When I call loadClass (marked by HERE above), it is not invoking "findClass" so method is not modified. From what I read, loadClass() is supposed to class findClass(). Any idea?


Solution

  • As already said by user16320675 in a comment, loadClass will attempt to load the class from the parent class loader first. So you can change new ClassLoader() to new ClassLoader(null) to set the bootstrap loader as its parent and it will not see the original definition.

    However, that will only work for very simple cases as then, the modified class can’t access other classes defined by the application class loader then.

    If the class has not been loaded yet, you can materialize the new definition like

    Class<?> myClass = MethodHandles.lookup().defineClass(modifiedClass);
    

    This will create the class in your current context, so even code using MyClass without Reflection will use the modified class.

    If you can’t preclude the environment from loading the class before this point, there is no way around writing a real Java Agent using the Instrumentation API or create a new environment with the custom class loader containing all classes the modified class might collaborate with. In other words, the environment of the class file transformator and the environment of the code to transform must be different then.


    In either case, after fixing this issue, you’ll get a java.lang.ClassFormatError: Duplicate method name "myMethod" with signature "()V" in class file MyClass, because your transformation code copies all artifacts of the original class file and adds another myMethod. To replace myMethod, you must intercept the class visitor when it encounters the original myMethod.

    import java.lang.invoke.MethodHandles;
    
    import org.objectweb.asm.*;
        
    public class MyClassInjector {
        public static void main(String[] args) throws Exception {
            // Load the MyClass class
            ClassReader cr = new ClassReader("MyClass");
            ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS);
            MyClassVisitor cv = new MyClassVisitor(cw);
            cr.accept(cv, 0);
    
            // Define the new byte array with the modified class bytecode
            byte[] modifiedClass = cw.toByteArray();
    
            Class<?> myClass = MethodHandles.lookup().defineClass(modifiedClass);
            Object myObject = myClass.getConstructor().newInstance(); // Class.newInstance() is deprecated
            myClass.getMethod("myMethod").invoke(myObject);
        }
    }
    
    class MyClassVisitor extends ClassVisitor {
        public MyClassVisitor(ClassVisitor cv) {
            super(Opcodes.ASM5, cv);
        }
    
        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
            MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
            if(name.equals("myMethod") && descriptor.equals("()V")) {
                instrument(mv);
                return null;
            }
            return mv;
        }
    
        private void instrument(MethodVisitor mv) {
          // change myMethod
          mv.visitCode();
          mv.visitMethodInsn(Opcodes.INVOKESTATIC, "MyClass", "newMethod", "()V", false);
          mv.visitInsn(Opcodes.RETURN);
          mv.visitMaxs(0, 0);
          mv.visitEnd();
        }
    
        @Override
        public void visitEnd() {
            // Define the newMethod method
            MethodVisitor mv = cv.visitMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "newMethod", "()V", null, null);
            mv.visitCode();
            mv.visitFieldInsn(Opcodes.GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
            mv.visitLdcInsn("Injected code");
            mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V", false);
            mv.visitInsn(Opcodes.RETURN);
            mv.visitMaxs(0, 0);
            mv.visitEnd();
            super.visitEnd();
        }
    }
    

    Note that there was another issue that you tried to invoke MyClassInjector.newMethod() in the modified code instead of MyClass.newMethod().