001    package ca.discotek.feenix.classloader;
002    
003    import java.io.InputStream;
004    import java.util.Iterator;
005    
006    import ca.discotek.rebundled.org.objectweb.asm.ClassReader;
007    import ca.discotek.rebundled.org.objectweb.asm.ClassVisitor;
008    import ca.discotek.rebundled.org.objectweb.asm.MethodVisitor;
009    import ca.discotek.rebundled.org.objectweb.asm.Opcodes;
010    import ca.discotek.rebundled.org.objectweb.asm.tree.ClassNode;
011    import ca.discotek.rebundled.org.objectweb.asm.tree.MethodNode;
012    
013    public class ClassLoaderClassVisitor extends ClassVisitor {
014    
015        static final String DEFINE_MY_CLASS_METHOD_NAME = "defineMyClass";
016        
017        static final String DEFINE_CLASS_METHOD_NAME = "defineClass";
018        static final String DEFINE_CLASS_METHOD_NEW_NAME = "_feenix_defineClass";
019        
020        static final String LOAD_CLASS_METHOD_NAME = "loadClass";
021        static final String LOAD_CLASS_METHOD_NEW_NAME = "_feenix_loadClass";
022        
023        static final String CURRENT_CLASS_NAMES[] = {ClassLoaderTargeted.class.getName().replace('.', '/')};
024        static final String NEW_CLASS_NAMES[] = {"java/lang/ClassLoader"};
025        
026        static final NameReplacerUtil classNameUtil = new NameReplacerUtil(CURRENT_CLASS_NAMES, NEW_CLASS_NAMES);
027        
028        static final String CURRENT_METHOD_NAMES[] = {LOAD_CLASS_METHOD_NAME, DEFINE_CLASS_METHOD_NAME};
029        static final String NEW_METHOD_NAMES[] = {LOAD_CLASS_METHOD_NEW_NAME, DEFINE_CLASS_METHOD_NEW_NAME};
030        
031        static final NameReplacerUtil METHOD_NAME_UTIL = new NameReplacerUtil(CURRENT_METHOD_NAMES, NEW_METHOD_NAMES);
032           
033        public ClassLoaderClassVisitor(ClassVisitor cv) {
034            super(Opcodes.ASM5, cv);
035        }
036    
037        public MethodVisitor visitMethod(int access,
038                String name,
039                String desc,
040                String signature,
041                String[] exceptions) {
042            
043            MethodVisitor mv = super.visitMethod(access, METHOD_NAME_UTIL.processName(name), desc, signature, exceptions);
044            if (name.equals(LOAD_CLASS_METHOD_NAME) && desc.equals("(Ljava/lang/String;)Ljava/lang/Class;"))
045                return new InvokeMethodNameReplacerMethodVisitor(mv, METHOD_NAME_UTIL);
046            else if (name.equals(DEFINE_CLASS_METHOD_NAME))
047                return new InvokeMethodNameReplacerMethodVisitor(mv, METHOD_NAME_UTIL);
048            else
049                return mv;
050        }
051    
052        static class NameReplacerUtil {
053            final String currentNames[];
054            final String newNames[];
055            
056            public NameReplacerUtil(String currentName, String newName) {
057                this(new String[]{currentName}, new String[]{newName});
058            }
059            
060            public NameReplacerUtil(String currentNames[], String newNames[]) {
061                this.currentNames = currentNames;
062                this.newNames = newNames;
063            }
064            
065            public String processName(String name) {
066                for (int i=0; i<currentNames.length; i++) {
067                    if (name.equals(currentNames[i])) 
068                        return newNames[i];
069                }
070                
071                return name;
072            }
073        }
074        
075        class InvokeMethodNameReplacerMethodVisitor extends MethodVisitor {
076    
077            final NameReplacerUtil util;
078            
079            public InvokeMethodNameReplacerMethodVisitor(MethodVisitor mv, NameReplacerUtil util) {
080                super(Opcodes.ASM5, mv);
081                this.util = util;
082            }
083            
084            public void visitMethodInsn(int opcode,
085                    String owner,
086                    String name,
087                    String desc,
088                    boolean itf) {
089    
090                super.visitMethodInsn(opcode, owner, util.processName(name), desc, itf);
091            }
092        }
093        
094        public void visitEnd() {
095            try {
096                InputStream is = 
097                    Thread.currentThread().getContextClassLoader().getResourceAsStream(ClassLoaderTargeted.class.getName().replace('.', '/') + ".class");
098                ClassReader cr = new ClassReader(is);
099                ClassNode node = new UpdateMethodInvocationsClassNode();
100                cr.accept(node, ClassReader.SKIP_FRAMES);
101                
102                Iterator<MethodNode> it = node.methods.listIterator();
103                MethodNode method;
104                String exceptions[];
105                while (it.hasNext()) {
106                    method = it.next();
107                    if (method.name.equals(DEFINE_CLASS_METHOD_NAME) || 
108                        method.name.equals(LOAD_CLASS_METHOD_NAME) || 
109                        method.name.equals(DEFINE_MY_CLASS_METHOD_NAME)) {
110                        
111                        exceptions = method.exceptions == null ? null : method.exceptions.toArray(new String[method.exceptions.size()]);
112                        MethodVisitor mv = super.visitMethod(method.access, method.name, method.desc, method.signature, exceptions);
113                        method.accept(mv);
114                    }
115                }
116            }
117            catch (Exception e) {
118                throw new Error("Unable to create classloader.", e);
119            }
120            
121            super.visitEnd();
122        }
123    
124        class UpdateMethodInvocationsClassNode extends ClassNode {
125            
126            public UpdateMethodInvocationsClassNode() {
127                super(Opcodes.ASM5);
128            }
129            
130            public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
131    
132                MethodNode mn = null;
133                
134                if (name.equals(DEFINE_CLASS_METHOD_NAME))
135                    mn = new UpdateMethodInvocationsMethodNode(access, name, desc, signature, exceptions, true);
136                else if (name.equals(LOAD_CLASS_METHOD_NAME))
137                    mn = new UpdateMethodInvocationsMethodNode(access, name, desc, signature, exceptions, !desc.equals("(Ljava/lang/String;)Ljava/lang/Class;"));
138                else if (name.equals(DEFINE_MY_CLASS_METHOD_NAME))
139                    mn = new UpdateMethodInvocationsMethodNode(access, name, desc, signature, exceptions, false);
140                else
141                    mn = new MethodNode(access, name, desc, signature, exceptions);
142                        
143                methods.add(mn);
144                
145                return mn;
146            }
147        }
148        
149        class UpdateMethodInvocationsMethodNode extends MethodNode {
150            
151            final boolean updateMethodNames;
152            
153            public UpdateMethodInvocationsMethodNode
154                (int access, String name, String desc, String signature, String exceptions[], boolean updateMethodNames) {
155                super(Opcodes.ASM5, access, name, desc, signature, exceptions);
156                this.updateMethodNames = updateMethodNames;
157            }
158            
159            public void visitMethodInsn(int opcode,
160                    String owner,
161                    String name,
162                    String desc,
163                    boolean itf) {
164                
165                String newMethodName = updateMethodNames ? METHOD_NAME_UTIL.processName(name) : name;
166                super.visitMethodInsn(opcode, classNameUtil.processName(owner), newMethodName, desc, itf);
167            }
168    
169        }
170        
171    }