Search in sources :

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();
}
Also used : ClassWriter(org.objectweb.asm.ClassWriter) Manifest(java.util.jar.Manifest) Authorizer(io.cdap.cdap.security.spi.authorization.Authorizer) Arrays(java.util.Arrays) MethodVisitor(org.objectweb.asm.MethodVisitor) LoggerFactory(org.slf4j.LoggerFactory) TypeToken(com.google.common.reflect.TypeToken) Type(org.objectweb.asm.Type) HashSet(java.util.HashSet) Method(org.objectweb.asm.commons.Method) GeneratorAdapter(org.objectweb.asm.commons.GeneratorAdapter) ClassVisitor(org.objectweb.asm.ClassVisitor) Nullable(javax.annotation.Nullable) FinallyAdapter(io.cdap.cdap.internal.asm.FinallyAdapter) Opcodes(org.objectweb.asm.Opcodes) ImmutableSet(com.google.common.collect.ImmutableSet) Logger(org.slf4j.Logger) AccessController(io.cdap.cdap.security.spi.authorization.AccessController) Set(java.util.Set) IOException(java.io.IOException) Attributes(java.util.jar.Attributes) Collectors(java.util.stream.Collectors) File(java.io.File) FilterClassLoader(io.cdap.cdap.common.lang.FilterClassLoader) Stream(java.util.stream.Stream) ClassPathResources(io.cdap.cdap.common.lang.ClassPathResources) Signatures(io.cdap.cdap.internal.asm.Signatures) ClassReader(org.objectweb.asm.ClassReader) DirectoryClassLoader(io.cdap.cdap.common.lang.DirectoryClassLoader) BundleJarUtil(io.cdap.cdap.common.lang.jar.BundleJarUtil) VisibleForTesting(com.google.common.annotations.VisibleForTesting) DirUtils(io.cdap.cdap.common.utils.DirUtils) InputStream(java.io.InputStream) FinallyAdapter(io.cdap.cdap.internal.asm.FinallyAdapter) HashSet(java.util.HashSet) ImmutableSet(com.google.common.collect.ImmutableSet) Set(java.util.Set) Method(org.objectweb.asm.commons.Method) ClassVisitor(org.objectweb.asm.ClassVisitor) ClassWriter(org.objectweb.asm.ClassWriter) MethodVisitor(org.objectweb.asm.MethodVisitor) Type(org.objectweb.asm.Type) ClassReader(org.objectweb.asm.ClassReader) GeneratorAdapter(org.objectweb.asm.commons.GeneratorAdapter) HashSet(java.util.HashSet) Nullable(javax.annotation.Nullable)

Aggregations

VisibleForTesting (com.google.common.annotations.VisibleForTesting)1 ImmutableSet (com.google.common.collect.ImmutableSet)1 TypeToken (com.google.common.reflect.TypeToken)1 ClassPathResources (io.cdap.cdap.common.lang.ClassPathResources)1 DirectoryClassLoader (io.cdap.cdap.common.lang.DirectoryClassLoader)1 FilterClassLoader (io.cdap.cdap.common.lang.FilterClassLoader)1 BundleJarUtil (io.cdap.cdap.common.lang.jar.BundleJarUtil)1 DirUtils (io.cdap.cdap.common.utils.DirUtils)1 FinallyAdapter (io.cdap.cdap.internal.asm.FinallyAdapter)1 Signatures (io.cdap.cdap.internal.asm.Signatures)1 AccessController (io.cdap.cdap.security.spi.authorization.AccessController)1 Authorizer (io.cdap.cdap.security.spi.authorization.Authorizer)1 File (java.io.File)1 IOException (java.io.IOException)1 InputStream (java.io.InputStream)1 Arrays (java.util.Arrays)1 HashSet (java.util.HashSet)1 Set (java.util.Set)1 Attributes (java.util.jar.Attributes)1 Manifest (java.util.jar.Manifest)1