001    package org.shiftone.jrat.inject.bytecode.asm;
002    
003    
004    import org.objectweb.asm.*;
005    import org.objectweb.asm.commons.GeneratorAdapter;
006    import org.objectweb.asm.commons.Method;
007    import org.shiftone.jrat.core.JRatException;
008    import org.shiftone.jrat.inject.bytecode.InjectorStrategy;
009    import org.shiftone.jrat.inject.bytecode.Modifier;
010    import org.shiftone.jrat.util.VersionUtil;
011    import org.shiftone.jrat.util.log.Logger;
012    
013    import java.util.Date;
014    
015    
016    public class InjectClassVisitor extends ClassAdapter implements Constants, Opcodes {
017    
018        private static final Logger LOG = Logger.getLogger(InjectClassVisitor.class);
019        private int handlerCount;
020        private Type classType;
021        private GeneratorAdapter initializer;
022    
023        public InjectClassVisitor(ClassVisitor visitor) {
024            super(visitor);
025        }
026    
027    
028        public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
029    
030            handlerCount = 0;
031            classType = Type.getType("L" + name.replace('.', '/') + ";");
032    
033            super.visit(version, access, name, signature, superName, interfaces);
034    
035            initializer = addInitializer();
036        }
037    
038    
039        private GeneratorAdapter addInitializer() {
040    
041            int access = Modifier.PRIVATE_STATIC;
042            String descriptor = "()V";
043            MethodVisitor initMethodVisitor = super.visitMethod(access, initializeName, descriptor, null, null);
044            GeneratorAdapter initializer = new GeneratorAdapter(initMethodVisitor, access, initializeName,
045                    descriptor);
046    
047            initMethodVisitor.visitCode();
048    
049            return initializer;
050        }
051    
052    
053        public void visitEnd() {
054    
055            initializer.returnValue();
056            initializer.endMethod();
057            addCommentField();
058            super.visitEnd();
059        }
060    
061    
062        private void addCommentField() {
063    
064            FieldVisitor commentField = super.visitField(Modifier.PRIVATE_STATIC_FINAL,
065                    InjectorStrategy.COMMENT_FIELD_NAME, "Ljava/lang/String;", null,
066                    "Class enhanced on " + new Date() + " w/ version JRat v"
067                            + VersionUtil.getBuiltOn() + " built on " + VersionUtil.getBuiltOn());
068    
069            commentField.visitEnd();
070        }
071    
072    
073        public FieldVisitor visitField(final int access, final String name, final String desc, final String signature,
074                                       final Object value) {
075    
076            if (name.equals(InjectorStrategy.COMMENT_FIELD_NAME)) {
077                throw new JRatException("this class was previously injected by JRat");
078            }
079    
080            return super.visitField(access, name, desc, signature, value);
081        }
082    
083    
084        private void addMethodHandlerField(String fieldName, String methodName, String descriptor) {
085    
086            FieldVisitor handler = super.visitField(Modifier.PRIVATE_STATIC_FINAL, fieldName,
087                    MethodHandler.TYPE.getDescriptor(), null, null);
088    
089            handler.visitEnd();
090            initializer.push(classType.getClassName());
091            initializer.push(methodName);
092            initializer.push(descriptor);
093            initializer.invokeStatic(HandlerFactory.TYPE, HandlerFactory.getMethodHandler);
094            initializer.putStatic(classType, fieldName, MethodHandler.TYPE);
095        }
096    
097    
098        public void pushThis(GeneratorAdapter adapter, boolean isStatic) {
099    
100            if (isStatic) {
101                adapter.push("test");
102            } else {
103                adapter.loadThis();
104            }
105        }
106    
107    
108        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature,
109                                         String[] exceptions) {
110    
111            if (name.equals("<clinit>") || name.equals("<init>") || Modifier.isAbstract(access)
112                    || Modifier.isNative(access)) {
113    
114                // LOG.debug("skipping " + name);
115                return super.visitMethod(access, name, descriptor, signature, exceptions);
116            }
117    
118            int index = (handlerCount++);
119            String handlerFieldName = InjectorStrategy.HANDLER_PREFIX + index;
120            String targetMethodName = name + InjectorStrategy.METHOD_POSTFIX;
121    
122            addMethodHandlerField(handlerFieldName, name, descriptor);
123    
124            // -- [ Proxy Method ] --
125            {
126                Method method = new Method(name, descriptor);
127                MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
128                ProxyMethodVisitor visitor = new ProxyMethodVisitor(access, method, mv, classType, targetMethodName,
129                        handlerFieldName);
130    
131                visitor.visitCode();
132                visitor.visitEnd();
133            }
134    
135            // -- [ Target Method ] --
136            return super.visitMethod(Modifier.makePrivate(access), targetMethodName, descriptor, signature, exceptions);
137        }
138    }