Byte code manipulation using Java instrumentation and jboss Javassist

In previous post we discussed how to use Java instrumentation to find size of object. In this post we will discuss how to modify Java byte codes using Transformer class (An agent implementing ClassFileTransformer interface) and jboss Javassist byte code modification library.
Lets understand importance of transformer class, how to register transformer class using instrumentation instance and finally how does it help in byte code instrumentation.

Transformer class:- A class which implements ClassFileTransformer interface is called transformer class.Transformer class implements transform() method. Transformer class is registered with instrumentation instance so that any further class loading(after Agent class by system class loader) invokes transform method of this transformer class. Signature of transform method is as follows :

public byte[] transform(ClassLoader loader, String className, Class classBeingRedefined,
ProtectionDomain protectionDomain, byte[] classfileBuffer){
           // returns  modified bytecode of class  or null, if not modified.
 }

Classloader invokes transform method for each class being loaded and passes information of classloader name which loaded this class or null if loaded by bootstrap classloader, class name and bytecode of class file.
Note:- Byte code of class file passed by class loader in the form of  "byte[] classfileBuffer" is not modified directly. We make a copy of classfileBuffer, modify it and return to class loader and class loader loads this modified byte code of class file.

What is uses of javassist.jar :- This library provides high level API to modify byte code of class file so that behaviour of class file can be changed dynamically(either by adding new methods or modifying existing one). We have another library in the form of Apache BCEL, ASM, etc which serves same purpose but in different way.
Bytecode modification framework can be broadly classified in two categories - One which provides high level API, that allows us to get away from learning low level opcodes and JVM internals (e.g, javaassist and CGLIB) and another one low level frameworks when we need to understand JVM or use some bytecode generation tools (ASM and BCEL).

Byte code injection demonstration:- 
In order to show how byte code injection works with javassist library we are creating a java project with following classes and manifest file as created in previous post(Fundamental of Instrumentation).

Create a class "InstrumentationAgent.java" and copy following code lines in it.(Adjust package name accordingly)

package com.devinline.instrumentation;

import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;

public class InstrumentationAgent {
 /*
  * System classloader after loading the Agent Class, then invokes the
  * premain (premain is roughly equal to main method for normal Java classes)
  */
 public static void premain(String agentArgs, Instrumentation inst) {
   inst.addTransformer(new Transformer());
 }

 public static void agentmain(String agentArgs, Instrumentation inst)
   throws ClassNotFoundException, UnmodifiableClassException {

 }
}

Create another class "Transformer.java" which implements ClassFileTransformer interface and its transform() method. Use following code lines.

package com.devinline.instrumentation;

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

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

public class Transformer implements ClassFileTransformer {

 @Override
 public byte[] transform(ClassLoader classLoader, String className,
   Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
   byte[] classfileBuffer) throws IllegalClassFormatException {
  byte[] modifiedClassByteCode = classfileBuffer;
  if (className.equals("com/devinline/client/SampleClass")) {
   try {
    ClassPool classPool = ClassPool.getDefault();
    CtClass ctClass = classPool.makeClass(new ByteArrayInputStream(
      classfileBuffer));
    CtMethod[] methods = ctClass.getDeclaredMethods();
    for (CtMethod method : methods) {
     if (method.getName().equalsIgnoreCase("method2")) {
      System.out.println("Start Instrumentation in method " + method.getName());
      method.addLocalVariable("elapsedTime", CtClass.longType);
      method.insertBefore("elapsedTime = System.currentTimeMillis();");
      method.insertAfter("{elapsedTime = System.currentTimeMillis() - elapsedTime;"
        + "System.out.println(\"Method Executed in ms: \" + elapsedTime);}");
     }

    }
    modifiedClassByteCode = ctClass.toBytecode();
    ctClass.detach();
    System.out.println("Instrumentation complete.");
   } catch (Throwable ex) {
    System.out.println("Exception: " + ex);
    ex.printStackTrace();
   }
  }
  return modifiedClassByteCode;
 }

}

Now create "SampleClass.java" whose methods are modified by transform method, when this class is loaded by classloader and transform method is invoked.

package com.devinline.client;

public class SampleClass {
 public void method1() throws InterruptedException {
  // randomly sleeps between 1000ms and 30000ms
  long randomSleepDuration = (long) (1000 + Math.random() * 2000);
  System.out.printf("Sleeping for %d ms ..\n", randomSleepDuration);
  Thread.sleep(randomSleepDuration);
 }

 public void method2(String input) throws InterruptedException {
  // randomly sleeps between 1000ms and 30000ms
  long randomSleepDuration = (long) (1000 + Math.random() * 2000);
  System.out.printf("Sleeping for %d ms ..\n", randomSleepDuration);
  Thread.sleep(randomSleepDuration);
 }

}

Finally, create a main class from where class loading starts and triggers method executions of sample class and verify that method has been instrumented of not.

package com.devinline.client;

public class TestInstrumentation {

 public static void main(String args[]) throws InterruptedException {
  SampleClass l = new SampleClass();
  System.out.println("Executing method1(Not instrumented) ");
  l.method1();
  System.out.println("Executing method2(Instrumented) ");
  l.method2("Hello");

 }
}

Update MANIFEST.Mf with property "premain-class" and "Boot-Class-Path". premain-class indicates the Agent class name to JVM and Boot-Class-Path is used to provide access of external jar (e.g: javassist.jar). Create a file manifest.txt and add following two entries in it:
premain-class: com.devinline.instrumentation.InstrumentationAgent
Boot-Class-Path: ../lib/javassist.jar
Note:- we have created lib folder in java project directory and placed javassist.jar inside it.

Create agent jar file:-  Execute following command from bin directory and create agent jar which will be passed to JVM instance with -javaagent:<agent.jar>.
> jar -cmf manifest.txt agent.jar com
It will create a jar executable jar file in bin directory with name agant.jar.

Execute main method with agent.jar  
Use follwoing command to run main method of TestInstrumentaion with agent.jar
> java -javaagent:agent.jar -cp . com.devinline.client.TestInstrumentation

Referring above diagram of program execution, we can see that method2 is instrumented and execution time of method 2 is printed and it is more than sleeping time J