Examples with FinallyAdapter io.cdap.cdap.internal.asm.FinallyAdapter used on opensource projects

Example 1 with FinallyAdapter

use of io.cdap.cdap.internal.asm.FinallyAdapter in project cdap by caskdata.

the class AccessControllerClassLoader method rewriteClass.

@Nullable
@Override
public byte[] rewriteClass(String className, InputStream input) throws IOException {
    if (!accessControllerClassName.equals(className)) {
        return null;
    }
    // Rewrite the AccessController class to wrap every methods call with context classloader change
    Set<java.lang.reflect.Method> accessControlMethods = Stream.of(Authorizer.class, AccessController.class).flatMap(c -> Stream.of(c.getMethods())).collect(Collectors.toSet());
    ClassReader cr = new ClassReader(input);
    ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
    cr.accept(new ClassVisitor(Opcodes.ASM7, cw) {

        private String superName;

        @Override
        public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
            super.visit(version, access, name, signature, superName, interfaces);
            this.superName = superName;
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
            MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
            Method method = new Method(name, descriptor);
            // Only rewrite methods defined in the Authorizer or AccessController interface.
            if (accessControlMethods.removeIf(m -> method.equals(Method.getMethod(m)))) {
                return rewriteMethod(access, name, descriptor, mv);
            }
            return mv;
        }

        @Override
        public void visitEnd() {
            if (Type.getType(Object.class).getInternalName().equals(superName)) {
                // the one we need and it's not trivial as those interfaces can be in the current jar we are trying to load.
                return;
            }
            // Generates all the missing methods on the Authorizer interface so that we can wrap them with the
            // context classloader switch.
            Set<Method> generatedMethods = new HashSet<>();
            new HashSet<>(accessControlMethods).forEach(m -> {
                Method method = Method.getMethod(m);
                // Guard against same method signature that comes from different parent interfaces
                if (!generatedMethods.add(method)) {
                    return;
                }
                // Generate the method by calling super.[method]
                String signature = Signatures.getMethodSignature(method, TypeToken.of(m.getGenericReturnType()), Arrays.stream(m.getGenericParameterTypes()).map(TypeToken::of).toArray(TypeToken[]::new));
                String[] exceptions = Arrays.stream(m.getExceptionTypes()).map(Type::getInternalName).toArray(String[]::new);
                MethodVisitor mv = rewriteMethod(Opcodes.ACC_PUBLIC, method.getName(), method.getDescriptor(), visitMethod(Opcodes.ACC_PUBLIC, method.getName(), method.getDescriptor(), signature, exceptions));
                GeneratorAdapter generator = new GeneratorAdapter(Opcodes.ACC_PUBLIC, method, mv);
                generator.visitCode();
                generator.loadThis();
                generator.loadArgs();
                generator.visitMethodInsn(Opcodes.INVOKESPECIAL, superName, method.getName(), method.getDescriptor(), false);
                generator.returnValue();
                generator.endMethod();
            });
            super.visitEnd();
        }

        /**
         * Rewrites the method by wrapping the whole method call with the context classloader switch to the
         * AccessControllerClassLoader.
         */
        private MethodVisitor rewriteMethod(int access, String name, String descriptor, MethodVisitor mv) {
            return new FinallyAdapter(Opcodes.ASM7, mv, access, name, descriptor) {

                int currentThread;

                int oldClassLoader;

                @Override
                protected void onMethodEnter() {
                    // Thread currentThread = Thread.currentThread();
                    invokeStatic(THREAD_TYPE, new Method("currentThread", THREAD_TYPE, new Type[0]));
                    currentThread = newLocal(THREAD_TYPE);
                    storeLocal(currentThread, THREAD_TYPE);
                    // ClassLoader oldClassLoader = currentThread.getContextClassLoader();
                    loadLocal(currentThread, THREAD_TYPE);
                    invokeVirtual(THREAD_TYPE, new Method("getContextClassLoader", CLASSLOADER_TYPE, new Type[0]));
                    oldClassLoader = newLocal(CLASSLOADER_TYPE);
                    storeLocal(oldClassLoader, CLASSLOADER_TYPE);
                    // currentThread.setContextClassLoader(getClass().getClassLoader());
                    loadLocal(currentThread, THREAD_TYPE);
                    loadThis();
                    invokeVirtual(Type.getType(Object.class), new Method("getClass", Type.getType(Class.class), new Type[0]));
                    invokeVirtual(Type.getType(Class.class), new Method("getClassLoader", CLASSLOADER_TYPE, new Type[0]));
                    invokeVirtual(THREAD_TYPE, new Method("setContextClassLoader", Type.VOID_TYPE, new Type[] { CLASSLOADER_TYPE }));
                    beginTry();
                }

                @Override
                protected void onFinally(int opcode) {
                    // currentThread.setContextClassLoader(oldClassLoader);
                    loadLocal(currentThread, THREAD_TYPE);
                    loadLocal(oldClassLoader, CLASSLOADER_TYPE);
                    invokeVirtual(THREAD_TYPE, new Method("setContextClassLoader", Type.VOID_TYPE, new Type[] { CLASSLOADER_TYPE }));
                }
            };
        }
    }, ClassReader.EXPAND_FRAMES);
    return cw.toByteArray();
}