Java Byte Code Instrumentation Example

Agent.java

package kr.rvs.instrumentation;

import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.Instrumentation;

/**
 * Created by Junhyeong Lim on 2017-02-02.
 */
public class Agent {
    public static void agentmain(String agentArg, Instrumentation inst) throws Exception {
        inst.addTransformer(new ClassTransformer());

        InputStream inStream = ClassLoader.getSystemResourceAsStream("java/lang/Exception.class");
        ByteArrayOutputStream outStream = new ByteArrayOutputStream();
        int read;
        byte[] data = new byte[65536];
        while ((read = inStream.read(data, 0, data.length)) != -1) {
            outStream.write(data, 0, read);
        }

        inst.redefineClasses(new ClassDefinition(Exception.class, outStream.toByteArray()));
    }
}

ClassTransformer.java

package kr.rvs.instrumentation;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;

import java.io.FileOutputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

import static org.objectweb.asm.Opcodes.ASM5;

/**
 * Created by Junhyeong Lim on 2017-02-02.
 */
public class ClassTransformer implements ClassFileTransformer {
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        byte[] ret = classfileBuffer;
        if (className.equals("java/lang/Exception")) {
            try {
                ClassReader reader = new ClassReader("java.lang.Exception");
                ClassWriter writer = new ClassWriter(0);
                ClassVisitor visitor = new TransformClassVisitor(ASM5, writer);

                reader.accept(visitor, 0);
                ret = writer.toByteArray();

                FileOutputStream out = new FileOutputStream("Exception.class");
                out.write(ret);
                out.close();

            } catch (Exception ex) {
                ex.printStackTrace();
            }
        }

        return ret;
    }
}

Main.java

package kr.rvs.instrumentation;

import com.sun.tools.attach.VirtualMachine;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.lang.management.ManagementFactory;
import java.util.jar.Attributes;
import java.util.jar.JarEntry;
import java.util.jar.JarOutputStream;
import java.util.jar.Manifest;

/**
 * Created by Junhyeong Lim on 2017-02-02.
 */
public class Main {
    public static void main(String[] args) {
        try {
            String jvm = ManagementFactory.getRuntimeMXBean().getName();
            String pid = jvm.substring(0, jvm.indexOf('@'));
            VirtualMachine vm = VirtualMachine.attach(pid);
            vm.loadAgent(generateJar(Agent.class, Utils.class, ClassTransformer.class, TransformClassVisitor.class, TransformMethodVisitor.class).getAbsolutePath());
            throw new Exception();
        } catch (Exception ex) {
            // Ignore
        }
    }

    public static File generateJar(Class agent, Class... resources) throws IOException {
        File jar = new File("agent.jar");
        jar.deleteOnExit();

        Manifest manifest = new Manifest();
        Attributes attr = manifest.getMainAttributes();

        attr.put(Attributes.Name.MANIFEST_VERSION, "1.0");
        attr.put(new Attributes.Name("Agent-Class"), "kr.rvs.instrumentation.Agent");
        attr.put(new Attributes.Name("Can-Retransform-Classes"), "true");
        attr.put(new Attributes.Name("Can-Redefine-Classes"), "true");

        JarOutputStream out = new JarOutputStream(new FileOutputStream(jar), manifest);
        out.putNextEntry(new JarEntry(Utils.getClassName(Agent.class)));
        out.write(Utils.getBytesAsClass(agent));
        out.closeEntry();

        for (Class cls : resources) {
            String name = Utils.getClassName(cls);
            out.putNextEntry(new JarEntry(name));
            out.write(Utils.getBytesAsClass(cls));
            out.closeEntry();
        }

        out.close();
        return jar;
    }
}

TransformClassVisitor.java

package kr.rvs.instrumentation;

import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.MethodVisitor;

/**
 * Created by Junhyeong Lim on 2017-02-02.
 */
public class TransformClassVisitor extends ClassVisitor {
    public TransformClassVisitor(int i, ClassVisitor classVisitor) {
        super(i, classVisitor);
    }

    @Override
    public MethodVisitor visitMethod(int i, String s, String s1, String s2, String[] strings) {
        MethodVisitor visitor = super.visitMethod(i, s, s1, s2, strings);
        if (s.equals("<init>")) {
            visitor = new TransformMethodVisitor(api, visitor);
        }
        return visitor;
    }
}

TransformMethodVisitor.java

package kr.rvs.instrumentation;

import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;

import static org.objectweb.asm.Opcodes.ARETURN;
import static org.objectweb.asm.Opcodes.ATHROW;
import static org.objectweb.asm.Opcodes.DRETURN;
import static org.objectweb.asm.Opcodes.FRETURN;
import static org.objectweb.asm.Opcodes.IRETURN;
import static org.objectweb.asm.Opcodes.LRETURN;
import static org.objectweb.asm.Opcodes.RETURN;

public class TransformMethodVisitor extends MethodVisitor {
    public TransformMethodVisitor(int i, MethodVisitor methodVisitor) {
        super(i, methodVisitor);
    }

    @Override
    public void visitInsn(int var1) {
        switch (var1) {
            case ARETURN:
            case DRETURN:
            case FRETURN:
            case IRETURN:
            case LRETURN:
            case RETURN:
            case ATHROW:
                visitFieldInsn(Opcodes.GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
                visitLdcInsn("Exception detected");
                visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V", false);
            default:
                break;
        }
        super.visitInsn(var1);
    }

    @Override
    public void visitEnd() {
        super.visitEnd();
    }
}

Utils.java

package kr.rvs.instrumentation;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;

/**
 * Created by Junhyeong Lim on 2017-02-02.
 */
public class Utils {
    public static String getClassName(Class cls) {
        return cls.getName().replace(".", "/") + ".class";
    }

    public static InputStream getStreamAsClass(Class cls) {
        return ClassLoader.getSystemResourceAsStream(getClassName(cls));
    }

    public static byte[] getBytesAsStream(InputStream in) throws IOException {
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int read;
        byte[] data = new byte[65536];
        while ((read = in.read(data, 0, data.length)) != -1) {
            out.write(data, 0, read);
        }
        return out.toByteArray();
    }

    public static byte[] getBytesAsClass(Class cls) throws IOException {
        return getBytesAsStream(getStreamAsClass(cls));
    }
}