How to write a Simple Agent

Posted on November 8, 2015 Updated on November 8, 2015

After two long introduction chapters, finally we are writing our simple java agent. Main target of our agent is to add an extra print statement to one of the methods in our main class (SimpleMain). In order to do that, we need to create following three files, which are the content of our agent.

• SimpleAgent.class
• SimpleClassFileTransformer.class
• MANIFEST.MF

We can generate these class files once we compile our .java files. But, before we start writing, check whether you have got maven installed with you, because we are going to need it when we generate the manifest file. Otherwise, generating the manifest file going to be a real pain when you start improving your agent. Therefore, better get use to it from the first lesson itself.

So, assuming you have already created a maven project let’s add dependency of Javassist to our pom.xml file. If you don’t include it, you won’t be able to use the methods of Javassist. (hope those experienced viewers won’t get annoyed with these details, because um trying to provide answers to all the questions i came across when i was writing my simple agent) You can obtain the dependency and jar from here based on your preference.

pom.xml 

<dependency>
 <groupId>org.javassist</groupId>
 <artifactId>javassist</artifactId>
 <version>3.20.0-GA</version>
</dependency>

SimpleClassFileTransformer.java

Next, lets write our SimpleClassFileTransformer class. This is where we include all the operations we need to do during the instrumentation process.

 
package com.testProject.javaagent;

import javassist.*;

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

public class SimpleClassFileTransformer implements 
ClassFileTransformer {
 public byte[] transform(ClassLoader loader, 
 String className, Class classBeingRedefined,
 ProtectionDomain protectionDomain, byte[] classfileBuffer)
 throws IllegalClassFormatException {

   byte[] byteCode = classfileBuffer;
   if(className.equals("com/testMain/SimpleMain")){
      ClassPool classPool = ClassPool.getDefault(); 
      try { 
         CtClass ctClass = classPool.makeClass(
                new ByteArrayInputStream(classfileBuffer)); 
         CtMethod[] methods = ctClass.getDeclaredMethods(); 
         for (CtMethod method : methods) {  
            if (method.getName().equals("print")) {
                System.out.println("Instrumenting method : " + method.getName());
                method.insertAt(1,
                  "System.out.println(\"This is the injected line\");"); 
            } 
         } 
         byteCode = ctClass.toBytecode(); 
         ctClass.detach(); 
       } catch (Throwable e) { 
            e.printStackTrace(); 
       }
   } 
 return byteCode;
 }
}

This class implements the ClassFileTransformer interface and override the transform method which can read and write class files as byte arrays. Parameters available in this method can be used in different situations based on our requirement. ClassfileBuffer contains the content of the currently instrumenting class as a byte array. When this method is called, it will obtain the current byte array and return a byte array including all the modifications done using the byte code instrumentation library. In this example, we try to obtain the class which match the fully qualified class name of our main class (com/testMain/SimpleMain) and instrument the print(). insertAt method of javassist allows us to insert our print statement to the specified line of the method. (Let’s explore other main methods of javassist later) Once the instrumentation is complete, it will recompile and verify whether insertions are according to the standards and load the classes to JVM.

SimpleAgent.java

This is where we create a transformer object and instruct to attach instruction to class (using addTransformer()). As it is mentioned earlier, invocation method of the agent is decided based on the main method used. In our example, we will be using the premain method which pass instrumentation object and other agent parameters and invoke agent before classes are load to VM.

 
package com.testProject.javaagent;

import java.lang.instrument.Instrumentation;

public class SimpleAgent {

 public static void premain(String agentArgs, 
                   Instrumentation instrumentation){

 System.out.println("Starting Agent");
 SimpleClassFileTransformer transformer = 
                  new SimpleClassFileTransformer();
 instrumentation.addTransformer(transformer);
 }
}

MANIFEST.MF

As it is mentioned in previous post, there are few attributes that need to include in the manifest file. But, for our simple agent, we will need a simple manifest file as follows. (I’ll post the code segment needed for the generation of manifest as a separate post). Class path is provided assuming all the jars needed by the agent are included in the ‘lib’ folder

Manifest-Version: 1.0
Premain-Class: org.wso2.javaagent.JDBCAgent
Class-Path: lib/javassist-3.20-GA.jar

These are the main files of our agent. Next step would be bundling of the .class files into our Agent.jar. This can also be automated using small piece of code added to our pom file. You can generate the jar using the command line as well. But i won’t go into detail about it, because it was bit annoying. But if you dig around a bit you would be able to figure it out. But, if you can wait a little, I am sure you’ll be happy with the automated way. Anyhow, assuming we have bundled our files into our Agent.jar let’s write the main class. Because, we need a main class to instrument with our agent.

SimpleMain.java

 
package com.testMain;

public class SimpleMain {
 public static void main(String[] args){
 SimpleMain sm = new SimpleMain();
 System.out.println("Simple Main");
 sm.print();
 }
 
 public void print(){
 System.out.println("Print method of Simple Main class");
 }
}

That’s it, we have everything to test our agent. Let’s run our agent using the command line. First compile the SimpleMain class. Then run the following command after going to the location of the SimpleMain.class.

> java -javaagent:path/to/agent/Agent.jar com.testMain.SimpleMain

There is a special keyword of java called as ‘-javaagent’ which is used to pass our java agent as a parameter to JVM. Value given after jar is the fully qualified name of our main class. You can also provide the class path of the main class using ‘-cp’ attribute of java.

> java -cp path/to/main -javaagent:path/to/agent/Agent.jar com.testMain.SimpleMain

Once we run our agent using above command our output will be as follows,

Starting Agent
Instrumenting method : print
Simple Main
This is the injected line
Print method of Simple Main class

Well, now you know how to write a simple java agent. Adding more features can be done modifying the the transform method. Next post will be on ‘How to generate a MANIFEST using Maven’. Future, ill get back with another simple example which uses more features of javassist.

Hope, you guys enjoy writing your first agent. Have a nice time…!!!