AgentInstaller.java

public class AgentInstaller {
  
	/** The created agent jar file name */
	protected static final AtomicReference<String> agentJar = new AtomicReference<String>(null); 
	
	/**
	 * Self installs the agent, then runs a person sayHello in a loop
	 * @param args None
	 */
	public static void main(String[] args) {
		try {
			// Get this JVM's PID
			String pid = ManagementFactory.getRuntimeMXBean().getName().split("@")[0];
			// Attach (to ourselves)
		    VirtualMachine vm = VirtualMachine.attach(pid);
		    // Create an agent jar (since we're in DEV mode)
	    	String fileName = createAgent();
	    	// Load the agent into this JVM
	    	vm.loadAgent(fileName);
	    	System.out.println("Agent Loaded");
			ObjectName on = new ObjectName("transformer:service=DemoTransformer");
			System.out.println("Instrumentation Deployed:" + ManagementFactory.getPlatformMBeanServer().isRegistered(on));
	    	// 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);
		}
	}
	
	
	/**
	 * 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");
						manifest.append("Premain-Class: " + AgentMain.class.getName() + "\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

public class DemoTransformer implements ClassFileTransformer {
  /** The normal form class name of the class to transform */
	protected String className;
	/** The class loader of the class */
	protected ClassLoader classLoader;
	/** The method name */
	protected String methodName;
	/** The method signature */
	protected String methodSignature;
	
	/**
	 * Creates a new DemoTransformer
	 * @param classLoader The classloader to match
	 * @param className The binary class name of the class to transform
	 * @param methodName The method name
	 * @param methodSignature A regular expression matching the method signature
	 */
	public DemoTransformer(ClassLoader classLoader, String className, String methodName, String methodSignature) {
		this.className = className.replace('.', '/');
		this.classLoader = classLoader;
		this.methodName = methodName;
		this.methodSignature = methodSignature;
	}

	/**
	 * {@inheritDoc}
	 * @see java.lang.instrument.ClassFileTransformer#transform(java.lang.ClassLoader, java.lang.String, java.lang.Class, java.security.ProtectionDomain, byte[])
	 */
	@Override
	public byte[] transform(ClassLoader loader, String className,
			Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
			byte[] classfileBuffer) throws IllegalClassFormatException {
		System.out.println("Examining class [" + className + "]");
		if(className.equals(this.className)  && loader.equals(classLoader)) {
			System.out.println("Instrumenting class [" + className + "]");
			return ModifyMethodTest.instrument(className, methodName, methodSignature, loader, classfileBuffer);
		}
		return classfileBuffer;
	}

}

ModifyMethodTest.java

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 + "]");
	}

}

TransformerService.java

public class TransformerService implements TransformerServiceMBean {
  /** The JVM's instrumentation instance */
	protected final Instrumentation instrumentation;
	
	/**
	 * Creates a new TransformerService
	 * @param instrumentation  The JVM's instrumentation instance 
	 */
	public TransformerService(Instrumentation instrumentation) {
		this.instrumentation = instrumentation;
	}

	/**
	 * {@inheritDoc}
	 * @see com.heliosapm.shorthandexamples.TransformerServiceMBean#transformClass(java.lang.String, java.lang.String, java.lang.String)
	 */
	@Override
	public void transformClass(String className, String methodName, String methodSignature) {
		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();
			transform(targetClazz, targetClassLoader, methodName, methodSignature);
			return;
		} catch (Exception ex) { /* Nope */ }
		// now try the hard/slow way
		for(Class<?> clazz: instrumentation.getAllLoadedClasses()) {
			if(clazz.getName().equals(className)) {
				targetClazz = clazz;
				targetClassLoader = targetClazz.getClassLoader();
				transform(targetClazz, targetClassLoader, methodName, methodSignature);
				return;				
			}
		}
		throw new RuntimeException("Failed to locate class [" + className + "]");
	}
	
	/**
	 * Registers a transformer and executes the transform
	 * @param clazz The class to transform
	 * @param classLoader The classloader the class was loaded from
	 * @param methodName The method name to instrument
	 * @param methodSignature The method signature to match
	 */
	protected void transform(Class<?> clazz, ClassLoader classLoader, String methodName, String methodSignature) {
		DemoTransformer dt = new DemoTransformer(classLoader, clazz.getName(), methodName, methodSignature);
		instrumentation.addTransformer(dt, true);
		try {
			instrumentation.retransformClasses(clazz);
		} catch (Exception ex) {
			throw new RuntimeException("Failed to transform [" + clazz.getName() + "]", ex);
		} finally {
			instrumentation.removeTransformer(dt);
		}		
	}
}

TransformerServiceMBean.java

public interface TransformerServiceMBean {
  /**
	 * Transforms the target class name
	 * @param className The binary name of the target class
	 * @param methodName  The name of the method to transform
	 * @param methodSignature A regular expression to match the method signature. (if null, matches ".*")
	 */
	public void transformClass(String className, String methodName, String methodSignature);
}