1 package org.hansel; 2 3 import java.io.IOException ; 4 import java.lang.instrument.ClassDefinition ; 5 import java.lang.instrument.ClassFileTransformer ; 6 import java.lang.instrument.IllegalClassFormatException ; 7 import java.lang.instrument.Instrumentation ; 8 import java.lang.instrument.UnmodifiableClassException ; 9 import java.lang.reflect.Method ; 10 import java.security.ProtectionDomain ; 11 import java.util.ArrayList ; 12 import java.util.HashMap ; 13 import java.util.HashSet ; 14 import java.util.List ; 15 import java.util.Map ; 16 17 public class Startup { 18 private static Instrumentation instrumentation; 19 20 private static Transformer transformer; 21 22 private static Map <ClassKey, byte[]> map = new HashMap <ClassKey, byte[]>(); 23 24 public static void premain(String options, Instrumentation instrumentation) { 25 Startup.instrumentation = instrumentation; 26 27 new ClassKey(null, null); 30 31 instrumentation.addTransformer(new ClassBufferCache()); 32 } 33 34 private static void delegate(HashSet <String > classNames) { 35 try { 36 Class clazz = ClassLoader.getSystemClassLoader().loadClass(Startup.class.getName()); 37 Method m = clazz.getMethod("init", new Class [] { HashSet .class }); 38 39 m.invoke(null, new Object [] { classNames }); 40 } catch (Exception e) { 41 throw new IllegalStateException (e); 42 } 43 } 44 45 public static void init(HashSet <String > classNames) 46 throws IOException , IllegalClassFormatException , ClassNotFoundException , 47 UnmodifiableClassException { 48 49 if (Startup.class.getClassLoader() != ClassLoader.getSystemClassLoader()) { 50 delegate(classNames); 51 return; 52 } 53 54 if (instrumentation == null) { 55 throw new IllegalStateException ("Instrumentation has not been initialized." + 56 " Please invoke VM with commandline option: " + 57 " -javaagent hansel.jar"); 58 } 59 60 Class [] allClasses = instrumentation.getAllLoadedClasses(); 61 62 transformer = new Transformer(classNames); 63 instrumentation.addTransformer(transformer); 64 65 List <ClassDefinition > redefine = new ArrayList <ClassDefinition >(); 66 67 for (int i=0; i<allClasses.length; i++) { 68 String className = allClasses[i].getName(); 69 70 if (classNames.contains(className)) { 71 74 byte[] classfileBuffer = map.get(new ClassKey(allClasses[i].getName().replace('.', '/'), 75 allClasses[i].getClassLoader())); 76 if (classfileBuffer == null) { 77 System.out.println("Skipping: " + allClasses[i].getName()); 78 new Exception ().printStackTrace(); 79 82 } else { 83 84 91 redefine.add(new ClassDefinition (allClasses[i], classfileBuffer)); 92 93 } 94 } 95 } 96 97 instrumentation.redefineClasses(redefine.toArray(new ClassDefinition [redefine.size()])); 98 99 } 100 101 117 118 public static void tearDown() 119 throws ClassNotFoundException , UnmodifiableClassException { 120 121 instrumentation.removeTransformer(transformer); 122 transformer.undo(instrumentation); 123 transformer = null; 124 } 125 126 private static class ClassKey { 127 private String classname; 128 private ClassLoader cl; 129 130 public ClassKey(String classname, ClassLoader cl) { 131 this.cl = cl; 132 this.classname = classname; 133 } 134 135 public boolean equals(Object obj) { 136 ClassKey ck = (ClassKey) obj; 137 return classname.equals(ck.classname) && cl.equals(ck.cl); 138 } 139 140 public int hashCode() { 141 int hashCode = classname.hashCode(); 142 if (cl != null) { 143 hashCode += cl.toString().hashCode(); 144 } 145 return hashCode; 146 } 147 148 public String toString() { 149 return classname + "-" + cl + "-" + hashCode(); 150 } 151 } 152 153 156 private static class ClassBufferCache implements ClassFileTransformer { 157 public byte[] transform(ClassLoader loader, 158 String className, 159 Class <?> classBeingRedefined, 160 ProtectionDomain protectionDomain, 161 byte[] classfileBuffer) { 162 map.put(new ClassKey(className, loader), classfileBuffer); 164 return null; 165 } 166 } 167 } 168 169 | Popular Tags |