001: package org.hansel;
002:
003: import java.io.IOException;
004: import java.lang.instrument.ClassDefinition;
005: import java.lang.instrument.ClassFileTransformer;
006: import java.lang.instrument.IllegalClassFormatException;
007: import java.lang.instrument.Instrumentation;
008: import java.lang.instrument.UnmodifiableClassException;
009: import java.lang.reflect.Method;
010: import java.security.ProtectionDomain;
011: import java.util.ArrayList;
012: import java.util.HashMap;
013: import java.util.HashSet;
014: import java.util.List;
015: import java.util.Map;
016:
017: public class Startup {
018: private static Instrumentation instrumentation;
019:
020: private static Transformer transformer;
021:
022: private static Map<ClassKey, byte[]> map = new HashMap<ClassKey, byte[]>();
023:
024: public static void premain(String options,
025: Instrumentation instrumentation) {
026: Startup.instrumentation = instrumentation;
027:
028: // load class, otherwise the class is loaded in transform(...) and might cause
029: // a deadlock.
030: new ClassKey(null, null);
031:
032: instrumentation.addTransformer(new ClassBufferCache());
033: }
034:
035: private static void delegate(HashSet<String> classNames) {
036: try {
037: Class clazz = ClassLoader.getSystemClassLoader().loadClass(
038: Startup.class.getName());
039: Method m = clazz.getMethod("init",
040: new Class[] { HashSet.class });
041:
042: m.invoke(null, new Object[] { classNames });
043: } catch (Exception e) {
044: throw new IllegalStateException(e);
045: }
046: }
047:
048: public static void init(HashSet<String> classNames)
049: throws IOException, IllegalClassFormatException,
050: ClassNotFoundException, UnmodifiableClassException {
051:
052: if (Startup.class.getClassLoader() != ClassLoader
053: .getSystemClassLoader()) {
054: delegate(classNames);
055: return;
056: }
057:
058: if (instrumentation == null) {
059: throw new IllegalStateException(
060: "Instrumentation has not been initialized."
061: + " Please invoke VM with commandline option: "
062: + " -javaagent hansel.jar");
063: }
064:
065: Class[] allClasses = instrumentation.getAllLoadedClasses();
066:
067: transformer = new Transformer(classNames);
068: instrumentation.addTransformer(transformer);
069:
070: List<ClassDefinition> redefine = new ArrayList<ClassDefinition>();
071:
072: for (int i = 0; i < allClasses.length; i++) {
073: String className = allClasses[i].getName();
074:
075: if (classNames.contains(className)) {
076: // Workaround for bug 5092850: Link class by calling getDeclaredFields.
077: //allClasses[i].getDeclaredFields();
078:
079: byte[] classfileBuffer = map.get(new ClassKey(
080: allClasses[i].getName().replace('.', '/'),
081: allClasses[i].getClassLoader()));
082: if (classfileBuffer == null) {
083: System.out.println("Skipping: "
084: + allClasses[i].getName());
085: new Exception().printStackTrace();
086: /* throw new IllegalStateException("Getting Classfilebuffer for class " + allClasses[i].getName() + "/"
087: + allClasses[i].getClassLoader() + "-" + new ClassKey(allClasses[i].getName().replace('.', '/'),
088: allClasses[i].getClassLoader()).hashCode() + ": " + map + " failed.");*/
089: } else {
090:
091: // getClassFileBuffer(className, allClasses[i].getClassLoader());
092: /*byte[] newClassDefinition =
093: transformer.transform(allClasses[i].getClassLoader(),
094: allClasses[i].getName(),
095: allClasses[i],
096: allClasses[i].getProtectionDomain(),
097: classfileBuffer);*/
098: redefine.add(new ClassDefinition(allClasses[i],
099: classfileBuffer/*newClassDefinition*/));
100:
101: }
102: }
103: }
104:
105: instrumentation.redefineClasses(redefine
106: .toArray(new ClassDefinition[redefine.size()]));
107:
108: }
109:
110: /* private static final byte[] getClassFileBuffer(String className, ClassLoader cl) throws IOException {
111: InputStream is = cl.getResourceAsStream(className.replace('.', '/') + ".class");
112: byte[] buf = new byte[1024];
113:
114: ByteArrayOutputStream baos = new ByteArrayOutputStream(is.available());
115: int read;
116: do {
117: read = is.read(buf);
118: if (read > -1) {
119: baos.write(buf, 0, read);
120: }
121: } while (read > -1);
122:
123: is.close();
124: return baos.toByteArray();
125: }*/
126:
127: public static void tearDown() throws ClassNotFoundException,
128: UnmodifiableClassException {
129:
130: instrumentation.removeTransformer(transformer);
131: transformer.undo(instrumentation);
132: transformer = null;
133: }
134:
135: private static class ClassKey {
136: private String classname;
137: private ClassLoader cl;
138:
139: public ClassKey(String classname, ClassLoader cl) {
140: this .cl = cl;
141: this .classname = classname;
142: }
143:
144: public boolean equals(Object obj) {
145: ClassKey ck = (ClassKey) obj;
146: return classname.equals(ck.classname) && cl.equals(ck.cl);
147: }
148:
149: public int hashCode() {
150: int hashCode = classname.hashCode();
151: if (cl != null) {
152: hashCode += cl.toString().hashCode();
153: }
154: return hashCode;
155: }
156:
157: public String toString() {
158: return classname + "-" + cl + "-" + hashCode();
159: }
160: }
161:
162: /**
163: * This is a horrible memory leak, but redefining classes is not possible otherwise :(
164: */
165: private static class ClassBufferCache implements
166: ClassFileTransformer {
167: public byte[] transform(ClassLoader loader, String className,
168: Class<?> classBeingRedefined,
169: ProtectionDomain protectionDomain,
170: byte[] classfileBuffer) {
171: //System.out.println("Transforming: " + new ClassKey(className, loader));
172: map.put(new ClassKey(className, loader), classfileBuffer);
173: return null;
174: }
175: }
176: }
|