Skip to content

Instantly share code, notes, and snippets.

@entrypointkr
Last active May 24, 2021 02:23
Show Gist options
  • Save entrypointkr/152f089f6f3884047abcd19d39297c9e to your computer and use it in GitHub Desktop.
Save entrypointkr/152f089f6f3884047abcd19d39297c9e to your computer and use it in GitHub Desktop.
Java Byte Code Instrumentation Example
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()));
}
}
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;
}
}
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;
}
}
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;
}
}
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();
}
}
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));
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment