001/*
002 *  Copyright 2013 Christopher Pheby
003 *
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *
008 *      http://www.apache.org/licenses/LICENSE-2.0
009 *
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package org.jadira.reflection.access.classloader;
017
018import java.lang.reflect.Method;
019import java.util.HashMap;
020import java.util.Map;
021import java.util.concurrent.ConcurrentHashMap;
022
023/**
024 * A ClassLoader which can be used to load classes from arbitrary byte arrays.
025 * Jadira uses this to load classes generated using ASM.
026 */
027public class AccessClassLoader extends ClassLoader {
028
029        private static final Method DEFINE_METHOD;
030        
031        static {
032                Method defineMethod = null;
033                try {
034                        defineMethod = ClassLoader.class.getDeclaredMethod("defineClass", new Class[] { String.class, byte[].class, int.class, int.class });
035                        defineMethod.setAccessible(true);
036                } catch (NoSuchMethodException e) {
037                } catch (SecurityException e) {
038                }
039                DEFINE_METHOD = defineMethod;
040        }
041         
042        private static final ConcurrentHashMap<ClassLoader, AccessClassLoader> ASM_CLASS_LOADERS = new ConcurrentHashMap<ClassLoader, AccessClassLoader>();
043
044        private static final Map<String, byte[]> registeredClasses = new HashMap<String, byte[]>();
045        
046        /**
047         * Creates a new instance using a suitable ClassLoader for the specified class
048         * @param typeToBeExtended The class to use to obtain a ClassLoader
049         * @return A new instance, or an existing instance if one already exists.
050         */
051        public static final AccessClassLoader get(Class<?> typeToBeExtended) {
052
053                ClassLoader loader = typeToBeExtended.getClassLoader();
054                return get(loader == null ? ClassLoader.getSystemClassLoader() : loader);
055        }
056        
057        /**
058         * Creates an AccessClassLoader for the given parent
059         * @param parent The parent ClassLoader for this instance
060         * @return A new instance, or an existing instance if one already exists.
061         */
062        public synchronized static final AccessClassLoader get(ClassLoader parent) {
063                AccessClassLoader loader = (AccessClassLoader) ASM_CLASS_LOADERS.get(parent);
064                if (loader == null) {
065                        loader = new AccessClassLoader(parent);
066                        ASM_CLASS_LOADERS.put(parent, loader);
067                }
068                return loader;
069        }
070
071        private AccessClassLoader(ClassLoader parentClassLoader) {
072                super(parentClassLoader);
073        }
074
075        @Override
076        public Class<?> loadClass(String name) throws ClassNotFoundException {
077
078                Class<?> loadedClass = findLoadedClass(name);
079
080                if (loadedClass == null) {
081                        
082                        try {
083                                loadedClass = findClass(name);
084                        } catch (ClassNotFoundException e) {
085                                // Ignore
086                        }
087                        
088                        if (loadedClass == null) {
089                                loadedClass = super.loadClass(name);
090                        }
091                }
092
093                return loadedClass;
094        }
095        
096        /**
097         * Registers a class by its name
098         * @param name The name of the class to be registered
099         * @param bytes An array of bytes containing the class
100         */
101        public void registerClass(String name, byte[] bytes) {
102            
103                if (registeredClasses.containsKey(name)) {
104                        throw new IllegalStateException("Attempted to register a class that has been registered already: " + name);
105                }
106                registeredClasses.put(name, bytes);
107        }
108        
109        @Override
110    public Class<?> findClass(String name) throws ClassNotFoundException {
111            
112                byte[] bytes = registeredClasses.get(name);
113                if (bytes != null) {
114                        registeredClasses.remove(name);
115                        try {
116                                return (Class<?>) DEFINE_METHOD.invoke(getParent(), new Object[] { name, bytes, Integer.valueOf(0), Integer.valueOf(bytes.length) });
117                        } catch (Exception ignored) {
118                        }
119                        return defineClass(name, bytes, 0, bytes.length);
120                }
121                throw new ClassNotFoundException("Cannot find class: " + name);
122    }
123}