001/**
002 * Copyright (c) 2011, The University of Southampton and the individual contributors.
003 * All rights reserved.
004 *
005 * Redistribution and use in source and binary forms, with or without modification,
006 * are permitted provided that the following conditions are met:
007 *
008 *   *  Redistributions of source code must retain the above copyright notice,
009 *      this list of conditions and the following disclaimer.
010 *
011 *   *  Redistributions in binary form must reproduce the above copyright notice,
012 *      this list of conditions and the following disclaimer in the documentation
013 *      and/or other materials provided with the distribution.
014 *
015 *   *  Neither the name of the University of Southampton nor the names of its
016 *      contributors may be used to endorse or promote products derived from this
017 *      software without specific prior written permission.
018 *
019 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
020 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
021 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
022 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
023 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
024 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
025 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
026 * ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
027 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
028 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
029 */
030package org.openimaj.aop.classloader;
031
032import java.io.File;
033import java.io.IOException;
034import java.net.URL;
035import java.net.URLClassLoader;
036import java.util.Arrays;
037import java.util.jar.Attributes;
038import java.util.jar.JarFile;
039import java.util.jar.Manifest;
040
041import javassist.ClassPool;
042import javassist.Loader;
043
044import org.openimaj.aop.ClassTransformer;
045import org.openimaj.aop.MultiTransformClassFileTransformer;
046
047/**
048 * {@link ClassLoaderTransform} provides an alternative to using a java agent to
049 * perform byte-code manipulation by providing a classloader that will
050 * automatically transform classes as they are loaded.
051 * 
052 * @author Jonathon Hare (jsh2@ecs.soton.ac.uk)
053 * 
054 */
055public class ClassLoaderTransform {
056        /**
057         * Manifest property key for specifying the actual main class when using the
058         * loading technique provided by {@link #main(String[])}.
059         */
060        public static final String MAIN_CLASS = "OpenIMAJ-Main-Class";
061
062        /**
063         * Manifest property key for specifying the {@link ClassTransformer} classes
064         * when using the loading technique provided by {@link #main(String[])}.
065         */
066        public static final String TRANSFORMERS = "OpenIMAJ-Transformers";
067
068        /**
069         * Separator character for separating multiple {@link ClassTransformer}s
070         * when using the loading technique provided by {@link #main(String[])}.
071         */
072        public static final String TRANSFORMERS_SEPARATOR = ",";
073
074        /**
075         * Run the main method of the given jar file, passing all classes as they
076         * are loaded through the transformer to modify the bytecode. The main
077         * method must be specified by the "Main-Class" manifest entry.
078         * 
079         * @param pool
080         *            the classpool to hold classes as they are loaded
081         * @param tf
082         *            the transform(s) to apply
083         * @param jarFile
084         *            the jar file to run
085         * @param args
086         *            the arguments to pass to the main method
087         * @return the classloader used for loading the classes in the jar file
088         * @throws Throwable
089         *             if an error occurs
090         */
091        public static Loader run(ClassPool pool, MultiTransformClassFileTransformer tf, File jarFile, String[] args)
092                        throws Throwable
093        {
094                final String mainClass = getMainClass(jarFile);
095
096                if (mainClass == null)
097                        throw new Exception("Failed to load Main-Class manifest attribute from " + jarFile);
098
099                pool.appendClassPath(jarFile.getAbsolutePath());
100
101                final ClassLoader parent = new URLClassLoader(new URL[] { jarFile.toURI().toURL() });
102
103                return run(parent, pool, tf, mainClass, args);
104        }
105
106        /**
107         * Get the main class from the manifest of the given jar file
108         * 
109         * @param jarFile
110         *            the jar file
111         * @return the main class or null if not found
112         * @throws IOException
113         *             if an error occurs when reading
114         */
115        public static String getMainClass(File jarFile) throws IOException {
116                JarFile jar = null;
117
118                try {
119                        jar = new JarFile(jarFile);
120                        final Manifest manifest = jar.getManifest();
121                        final Attributes attr = manifest.getMainAttributes();
122                        final String mainClass = attr.getValue("Main-Class");
123
124                        return mainClass;
125                } finally {
126                        try {
127                                if (jar != null)
128                                        jar.close();
129                        } catch (final IOException e) {
130                        }
131                }
132        }
133
134        /**
135         * Run the main method of the given class, transforming any classes found on
136         * the given classpath as they are loaded.
137         * 
138         * @param pool
139         *            the classpool to hold classes as they are loaded
140         * @param tf
141         *            the transform(s) to apply
142         * @param classpath
143         *            the classpath on which to search for classes
144         * @param mainClass
145         *            the main class with the <code>main(String[])</code> method
146         * @param args
147         *            the arguments to pass to the main method
148         * @return the classloader used for loading the classes in the jar file
149         * @throws Throwable
150         *             if an error occurs
151         */
152        public static Loader run(ClassPool pool, MultiTransformClassFileTransformer tf, String classpath, String mainClass,
153                        String[] args)
154                        throws Throwable
155        {
156                final String[] paths = classpath.split(File.pathSeparator);
157                final URL[] urls = new URL[paths.length];
158                for (int i = 0; i < paths.length; i++)
159                        urls[i] = new File(paths[i]).toURI().toURL();
160
161                final ClassLoader parent = new URLClassLoader(urls);
162
163                pool.appendPathList(classpath);
164
165                return run(parent, pool, tf, mainClass, args);
166        }
167
168        /**
169         * Re-load the given class in a newly created {@link Loader} that is
170         * configured to apply the given transform(s), and then run the main method.
171         * 
172         * @param transformer
173         *            the transformer
174         * @param clz
175         *            the class
176         * @param args
177         *            the arguments
178         * @throws Throwable
179         *             if an error occurs
180         */
181        public static void run(MultiTransformClassFileTransformer transformer, Class<?> clz, String[] args) throws Throwable {
182                final ClassPool pool = new ClassPool();
183                pool.appendSystemPath();
184
185                final Loader loader = new Loader(pool);
186
187                // skip args4j
188                loader.delegateLoadingOf("org.apache.log4j.");
189
190                loader.addTranslator(pool, transformer);
191
192                loader.run(clz.getName(), args);
193        }
194
195        protected static Loader run(final ClassLoader parent, ClassPool pool, MultiTransformClassFileTransformer tf,
196                        String mainClass, String[] args)
197                        throws Throwable
198        {
199                final Loader cl = new Loader(parent, pool);
200
201                // Set the correct app name on OSX. Are there similar controls for other
202                // platforms?
203                System.setProperty("com.apple.mrj.application.apple.menu.about.name", mainClass);
204
205                // skip args4j
206                cl.delegateLoadingOf("org.apache.log4j.");
207
208                cl.addTranslator(pool, tf);
209
210                cl.run(mainClass, args);
211
212                return cl;
213        }
214
215        /**
216         * If the the given class has not already been loaded in by {@link Loader},
217         * load it in a newly created {@link Loader} that is configured to apply the
218         * given transform(s), run the main method and return true. If the class's
219         * classload is a {@link Loader} then return false.
220         * 
221         * @param clz
222         *            the class to load, transform and run
223         * @param args
224         *            the arguments to the main method
225         * @param transformer
226         *            the first transformer
227         * @param transformers
228         *            any additional transformers
229         * @return true if the class has been reloaded; false otherwise
230         * @throws Throwable
231         *             if an error occurs.
232         */
233        public static boolean
234                        run(Class<?> clz, String[] args, ClassTransformer transformer, ClassTransformer... transformers)
235                                        throws Throwable
236        {
237                if (!(clz.getClassLoader() instanceof Loader)) {
238                        final MultiTransformClassFileTransformer mtf =
239                                        new MultiTransformClassFileTransformer(transformer, transformers);
240
241                        ClassLoaderTransform.run(mtf, clz, args);
242
243                        return true;
244                }
245
246                return false;
247        }
248
249        /**
250         * Main method for an alternative in-place loading strategy. Allows a
251         * self-contained program jar to automatically run {@link ClassTransformer}s
252         * on itself.
253         * <p>
254         * Modify the assembled jar file so that the Main-Class is set as this class
255         * (<code>org.openimaj.aop.classloader.ClassLoaderTransform</code>) and add
256         * additional keys to the manifest to point at the actual class main (
257         * {@link #MAIN_CLASS}) and the transformer classes ({@link #TRANSFORMERS},
258         * separated by {@link #TRANSFORMERS_SEPARATOR}).
259         * <p>
260         * The specified transformer classes must all be instances of
261         * {@link ClassTransformer} and also have a no-args constructor.
262         * 
263         * @param args
264         *            the arguments to pass to the program being run
265         * @throws Throwable
266         *             if an error occurs
267         */
268        public static void main(String[] args) throws Throwable {
269                final Class<ClassLoaderTransform> clazz = ClassLoaderTransform.class;
270                final String className = clazz.getSimpleName() + ".class";
271                final String classPath = clazz.getResource(className).toString();
272                if (!classPath.startsWith("jar")) {
273                        // Class not from JAR
274                        return;
275                }
276                final String manifestPath = classPath.substring(0,
277                                classPath.lastIndexOf("!") + 1) + "/META-INF/MANIFEST.MF";
278                final Manifest manifest = new Manifest(new URL(manifestPath).openStream());
279                final Attributes attr = manifest.getMainAttributes();
280
281                final String mainClass = attr.getValue(MAIN_CLASS);
282                final String[] tfs = attr.getValue(TRANSFORMERS).split(TRANSFORMERS_SEPARATOR);
283                final ClassTransformer[] cts = new ClassTransformer[tfs.length];
284                for (int i = 0; i < tfs.length; i++)
285                        cts[i] = (ClassTransformer) Class.forName(tfs[i]).newInstance();
286
287                final ClassPool cp = ClassPool.getDefault();
288
289                run(null,
290                                cp,
291                                new MultiTransformClassFileTransformer(cts[0], Arrays.copyOfRange(cts, 1, cts.length)),
292                                mainClass,
293                                args);
294        }
295}