AgentInstaller.java

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.concurrent.atomic.AtomicReference;
import java.util.jar.JarOutputStream;
import java.util.zip.ZipEntry;
import java.util.jar.Manifest;

import com.sun.tools.attach.VirtualMachine;


public class AgentInstaller {
  
	/** The created agent jar file name */
	protected static final AtomicReference<String> agentJar = new AtomicReference<String>(null); 
	
	/**
	 * Self installs the agent
	 * @param args are pid of target jvm, target class name, target method name. (e.g. 6332 Person sayHello)
	 */
	public static void main(String[] args) {
		try {
			// Attach (to ourselves)
		    VirtualMachine vm = VirtualMachine.attach(args[0]);
		    // Create an agent jar (since we're in DEV mode)
	    	String fileName = createAgent();
	    	// Load the agent into this JVM
	    	String agentArgs = args[1]+","+args[2];
	    	vm.loadAgent(fileName, agentArgs);
	    	System.out.println("Agent Loaded");
		    			
		} catch (Exception ex) {
			System.err.println("Agent Installation Failed. Stack trace follows...");
			ex.printStackTrace(System.err);
		}
	}
	
	
	/**
	 * Creates the temporary agent jar file if it has not been created
	 * @return The created agent file name
	 */
	public static String createAgent() {
		if(agentJar.get()==null) {
			synchronized(agentJar) {
				if(agentJar.get()==null) {
					FileOutputStream fos = null;
					JarOutputStream jos = null;
					try {
						File tmpFile = File.createTempFile(AgentMain.class.getName(), ".jar");
						System.out.println("Temp File:" + tmpFile.getAbsolutePath());
						tmpFile.deleteOnExit();		
						StringBuilder manifest = new StringBuilder();
						manifest.append("Manifest-Version: 1.0\nAgent-Class: " + AgentMain.class.getName() + "\n");
						manifest.append("Can-Redefine-Classes: true\n");
						manifest.append("Can-Retransform-Classes: true\n");
									
						ByteArrayInputStream bais = new ByteArrayInputStream(manifest.toString().getBytes());
						Manifest mf = new Manifest(bais);
						fos = new FileOutputStream(tmpFile, false);
						jos = new JarOutputStream(fos, mf);
						addClassesToJar(jos, AgentMain.class, DemoTransformer.class, ModifyMethodTest.class, TransformerService.class, TransformerServiceMBean.class);
						jos.flush();
						jos.close();
						fos.flush();
						fos.close();
						agentJar.set(tmpFile.getAbsolutePath());
					} catch (Exception e) {
						throw new RuntimeException("Failed to write Agent installer Jar", e);
					} finally {
						if(fos!=null) try { fos.close(); } catch (Exception e) {}
					}

				}
			}
		}
		return agentJar.get();
	}
	
	/**
	 * Writes the passed classes to the passed JarOutputStream
	 * @param jos the JarOutputStream
	 * @param clazzes The classes to write
	 * @throws IOException on an IOException
	 */
	protected static void addClassesToJar(JarOutputStream jos, Class<?>...clazzes) throws IOException {
		for(Class<?> clazz: clazzes) {
			jos.putNextEntry(new ZipEntry(clazz.getName().replace('.', '/') + ".class"));
			jos.write(getClassBytes(clazz));
			jos.flush();
			jos.closeEntry();
		}
	}
	
	/**
	 * Returns the bytecode bytes for the passed class
	 * @param clazz The class to get the bytecode for
	 * @return a byte array of bytecode for the passed class
	 */
	public static byte[] getClassBytes(Class<?> clazz) {
		InputStream is = null;
		try {
			is = clazz.getClassLoader().getResourceAsStream(clazz.getName().replace('.', '/') + ".class");
			ByteArrayOutputStream baos = new ByteArrayOutputStream(is.available());
			byte[] buffer = new byte[8092];
			int bytesRead = -1;
			while((bytesRead = is.read(buffer))!=-1) {
				baos.write(buffer, 0, bytesRead);
			}
			baos.flush();
			return baos.toByteArray();
		} catch (Exception e) {
			throw new RuntimeException("Failed to read class bytes for [" + clazz.getName() + "]", e);
		} finally {
			if(is!=null) { try { is.close(); } catch (Exception e) {} }
		}
	}	
}

AgentMain.java

public class AgentMain {
  
	/**
	 * Installs the transformation service
	 * @param agentArgs None supported
	 * @param inst The instrumentation instance
	 * @throws Exception thrown on any error
	 */
	public static void agentmain (String agentArgs, Instrumentation inst) throws Exception {
		System.out.println("Installing AgentMain...");
		TransformerService ts = new TransformerService(inst);
		ObjectName on = new ObjectName("transformer:service=DemoTransformer");
		// Could be a different MBeanServer. If so, pass a JMX Default Domain Name in agentArgs
		MBeanServer server = ManagementFactory.getPlatformMBeanServer();
		server.registerMBean(ts, on);
		// Set this property so the installer knows we're already here
		System.setProperty("demo.agent.installed", "true");		
		System.out.println("AgentMain Installed");
	}

}

DemoTransformer.java

import java.lang.instrument.Instrumentation;

public class AgentMain {
  
	/**
	 * Installs the transformation service
	 * @param agentArgs None supported
	 * @param inst The instrumentation instance
	 * @throws Exception thrown on any error
	 */
	public static void agentmain (String agentArgs, Instrumentation inst) throws Exception {
		System.out.println("Installing AgentMain..."+agentArgs);
		
		String[] tmpString = agentArgs.split(",");
		
		String className = tmpString[0];
		String methodName = tmpString[1];
		String methodSignature = null;
		Class<?> targetClazz = null;
		ClassLoader targetClassLoader = null;
		// first see if we can locate the class through normal means
		//try {
		//	targetClazz = Class.forName(className);
		//	targetClassLoader = targetClazz.getClassLoader();
		//	return;
		//} catch (Exception ex) { /* Nope */ }
		// now try the hard/slow way
		try {
			for(Class<?> clazz: inst.getAllLoadedClasses()) {
				if(clazz.getName().equals(className)) {
					targetClazz = clazz;
					targetClassLoader = targetClazz.getClassLoader();
		//			return;				
				}
			}
			// throw new RuntimeException("Failed to locate class [" + className + "]");
		} catch (Exception e) {
			throw new RuntimeException("Failed to locate class [" + className + "]");
		}
		DemoTransformer dt = new DemoTransformer(targetClassLoader, targetClazz.getName(), methodName, methodSignature);
		inst.addTransformer(dt, true);
		try {
			inst.retransformClasses(targetClazz);
		} catch (Exception ex) {
			throw new RuntimeException("Failed to transform [" + targetClazz.getName() + "]", ex);
		} finally {
			inst.removeTransformer(dt);
		}	
		
	}

}

ModifyMethodTest.java

import java.util.regex.Pattern;

import javassist.ByteArrayClassPath;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.LoaderClassPath;

public class ModifyMethodTest {
  
	/**
	 * Creates a new ModifyMethodTest
	 * @param className The internal form class name to modify
	 * @param methodName  The name of the method to transform
	 * @param methodSignature A regular expression to match the method signature. (if null, matches ".*")
	 * @param classLoader The intrumentation provided classloader
	 * @param byteCode The pre-transform byte code  
	 * @return  the modified byte code if successful, otherwise returns the original unmodified byte code
	 */
	public static byte[] instrument(String className, String methodName, String methodSignature, ClassLoader classLoader, byte[] byteCode) {
		String binName  = className.replace('/', '.');
		try {
			ClassPool cPool = new ClassPool(true);
			cPool.appendClassPath(new LoaderClassPath(classLoader));
			cPool.appendClassPath(new ByteArrayClassPath(binName, byteCode));
			CtClass ctClazz = cPool.get(binName);
			Pattern sigPattern = Pattern.compile((methodSignature==null || methodSignature.trim().isEmpty()) ? ".*" : methodSignature);
			int modifies = 0;
			for(CtMethod method: ctClazz.getDeclaredMethods()) {
				if(method.getName().equals(methodName)) {
					if(sigPattern.matcher(method.getSignature()).matches()) {
						ctClazz.removeMethod(method);
						String newCode = "System.out.println(\"\\n\\t-->Invoked method [" + binName + "." + method.getName() + "(" + method.getSignature() + ")]\");";
						System.out.println("[ModifyMethodTest] Adding [" + newCode + "]");
						method.insertBefore(newCode);
						ctClazz.addMethod(method);
						modifies++;
					}
				}
			}
			
			System.out.println("[ModifyMethodTest] Intrumented [" + modifies + "] methods");
			return ctClazz.toBytecode();
		} catch (Exception ex) {
			System.err.println("Failed to compile retransform class [" + binName + "] Stack trace follows...");
			ex.printStackTrace(System.err);
			return byteCode; 
		}
	}

}

Person.java

public class Person {
  public void sayHello(String name) {
		System.out.println("Hello [" + name + "]");
	}
	
	public void sayHello(int x) {
		System.out.println("Hello [" + x + "]");
	}

}

TestAI.java

public class TestAI {
  
  /**
	 * Runs a person sayHello in a loop
	 * @param args None
	 */
	public static void main(String[] args) {
		try {
	    	// Run sayHello in a loop
			Person person = new Person();
			for(int i = 0; i < 1000; i++) {
				person.sayHello(i);
				person.sayHello("" + (i*-1));
				Thread.currentThread().join(5000);
			}
		} catch (Exception ex) {
			System.err.println("Agent Installation Failed. Stack trace follows...");
			ex.printStackTrace(System.err);
		}
	}	
}