Java ASM bytecode manipulation - add code to constructor of a library class

283 views Asked by At

I have code to insert instructions into the constructor of a class from a third party library (okhttp3.OkHttpClient in this case). Disassembled class shows the added line. However, I am not sure how to make rest of the code to use this instrumented OkHttpClient class instead of using the original one. Once we have the modified class, do we need to install that in classpath or manipulate any class loaders? When the test method is run, I do not see "Constructor called" logged.

Here is the full code:

import java.io.FileOutputStream;
import java.io.IOException;

import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
public class ConstructorInjector {
    public static void main(String[] args) {
        String className = "okhttp3.OkHttpClient";

        byte[] modifiedBytecode = new byte[0];

        try {
            modifiedBytecode = ConstructorInjector.injectIntoConstructor(className);
        } catch (Exception e) {
            throw new IllegalArgumentException(e);
        }

        try (FileOutputStream fos = new FileOutputStream("OkHttpClientModified.class")) {
            fos.write(modifiedBytecode);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }

        try {
            test();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    // Example method to inject into constructor
    public static void logConstructorCall() {
        System.out.println("Constructor called");
    }

    // Example method to modify a class
    public static byte[] injectIntoConstructor(String className) throws Exception {
        ClassReader classReader = new ClassReader(className);
        ClassWriter classWriter = new ClassWriter(classReader, ClassWriter.COMPUTE_MAXS);
        ClassVisitor classVisitor = new ClassVisitor(Opcodes.ASM7, classWriter) {

            // Override visitMethod to modify the constructor
            @Override
            public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
                MethodVisitor methodVisitor = super.visitMethod(access, name, desc, signature, exceptions);
                if (name.equals("<init>")) {
                    // Override the constructor
                    return new MethodVisitor(Opcodes.ASM7, methodVisitor) {
                        @Override
                        public void visitInsn(int opcode) {
                            // Find the last instruction in the constructor
                            if (opcode == Opcodes.RETURN) {
                                // Insert code before the return instruction
                                visitMethodInsn(Opcodes.INVOKESTATIC, "ConstructorInjector", "logConstructorCall", "()V", false);
                            }
                            super.visitInsn(opcode);
                        }
                    };
                }
                return methodVisitor;
            }
        };

        // Modify the class using the classVisitor
        classReader.accept(classVisitor, ClassReader.EXPAND_FRAMES);

        // Return the modified bytecode
        return classWriter.toByteArray();
    }

    private static void test() throws IOException {
        OkHttpClient client = new OkHttpClient.Builder()
                .build();

        Request request = new Request.Builder()
                .url("http://www.publicobject.com/helloworld.txt")
                .header("User-Agent", "OkHttp Example")
                .build();

        Response response = client.newCall(request).execute();
        String body = response.body().string();
        System.out.println("++++ BODY: " + body);
        //response.body().close();
    }
}
0

There are 0 answers