Skip to content

Instantly share code, notes, and snippets.

@ppdouble
Forked from nickman/AgentInstaller.java
Created September 11, 2013 08:17
Show Gist options
  • Save ppdouble/6520694 to your computer and use it in GitHub Desktop.
Save ppdouble/6520694 to your computer and use it in GitHub Desktop.
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) {} }
}
}
}
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");
}
}
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;
}
}
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;
}
}
}
public class Person {
public void sayHello(String name) {
System.out.println("Hello [" + name + "]");
}
public void sayHello(int x) {
System.out.println("Hello [" + x + "]");
}
}
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);
}
}
}
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);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment