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.agent;
031
032import java.io.File;
033import java.io.FileOutputStream;
034import java.io.IOException;
035import java.io.InputStream;
036import java.io.OutputStream;
037import java.io.UnsupportedEncodingException;
038import java.lang.instrument.Instrumentation;
039import java.lang.management.ManagementFactory;
040import java.lang.reflect.Method;
041import java.net.URL;
042import java.net.URLClassLoader;
043import java.security.AccessController;
044import java.security.PrivilegedActionException;
045import java.security.PrivilegedExceptionAction;
046import java.util.ArrayList;
047import java.util.List;
048import java.util.jar.JarEntry;
049import java.util.jar.JarOutputStream;
050
051import org.apache.logging.log4j.Logger;
052import org.apache.logging.log4j.LogManager;
053
054/**
055 * Dynamic agent loader. Provides methods to extract an agent jar with a
056 * specific agent class, and to attempt to dynamically load agents on Oracle
057 * JVMs (including OpenJDK). Dynamic loading won't work on all JVMs (i.e. IBMs),
058 * but the standard "-javaagent" commandline option can be used instead. If
059 * dynamic loading fails, instructions on using the command line flag will be
060 * printed to stderr.
061 *
062 * @author Jonathon Hare (jsh2@ecs.soton.ac.uk)
063 */
064public class AgentLoader {
065        private static final Logger logger = LogManager.getLogger(AgentLoader.class);
066        private static final String VMCLASS = "com.sun.tools.attach.VirtualMachine";
067
068        private static long copy(InputStream input, OutputStream output) throws IOException {
069                long count = 0;
070                int n = 0;
071                final byte[] buffer = new byte[4096];
072
073                while (-1 != (n = input.read(buffer))) {
074                        output.write(buffer, 0, n);
075                        count += n;
076                }
077
078                return count;
079        }
080
081        private static byte[] createManifest(Class<?> agentClass) {
082                final StringBuffer sb = new StringBuffer();
083
084                try {
085                        agentClass.getDeclaredMethod("premain", String.class, Instrumentation.class);
086                        sb.append("Premain-Class: " + agentClass.getName() + "\n");
087                } catch (final NoSuchMethodException e) {
088                        // IGNORE//
089                }
090
091                try {
092                        agentClass.getDeclaredMethod("agentmain", String.class, Instrumentation.class);
093                        sb.append("Agent-Class: " + agentClass.getName() + "\n");
094                } catch (final NoSuchMethodException e) {
095                        // IGNORE//
096                }
097
098                sb.append("Can-Redefine-Classes: true\n");
099                sb.append("Can-Retransform-Classes: true\n");
100
101                try {
102                        return sb.toString().getBytes("US-ASCII");
103                } catch (final UnsupportedEncodingException e) {
104                        throw new RuntimeException("Charset US-ASCII isn't supported!! This should never happen.");
105                }
106        }
107
108        /**
109         * Create an agent jar file with the required manifest entries.
110         *
111         * @param file
112         *            the location to create the jar
113         * @param agentClass
114         *            the agent class
115         * @throws IOException
116         *             if an error occurs
117         */
118        public static void createAgentJar(File file, Class<?> agentClass) throws IOException {
119                final JarOutputStream jos = new JarOutputStream(new FileOutputStream(file));
120
121                String classEntryPath = agentClass.getName().replace(".", "/") + ".class";
122                final InputStream classStream = agentClass.getClassLoader().getResourceAsStream(classEntryPath);
123
124                if (classEntryPath.startsWith("/"))
125                        classEntryPath = classEntryPath.substring(1);
126
127                JarEntry entry = new JarEntry(classEntryPath);
128                jos.putNextEntry(entry);
129                copy(classStream, jos);
130                jos.closeEntry();
131
132                entry = new JarEntry("META-INF/MANIFEST.MF");
133                jos.putNextEntry(entry);
134                jos.write(createManifest(agentClass));
135                jos.closeEntry();
136
137                jos.close();
138        }
139
140        /**
141         * Attempt to locate potential "tools.jar" jars
142         */
143        private static List<File> getPotentialToolsJars() {
144                final List<File> jars = new ArrayList<File>();
145
146                final File javaHome = new File(System.getProperty("java.home"));
147
148                final File jreSourced = new File(javaHome, "lib/tools.jar");
149                if (jreSourced.exists()) {
150                        jars.add(jreSourced);
151                }
152
153                if ("jre".equals(javaHome.getName())) {
154                        final File jdkHome = new File(javaHome, "../");
155                        final File jdkSourced = new File(jdkHome, "lib/tools.jar");
156                        if (jdkSourced.exists()) {
157                                jars.add(jdkSourced);
158                        }
159                }
160
161                return jars;
162        }
163
164        /**
165         * Try and get the VirtualMachine class
166         */
167        private static Class<?> tryGetVMClass() {
168                try {
169                        return AccessController.doPrivileged(new PrivilegedExceptionAction<Class<?>>() {
170                                @Override
171                                public Class<?> run() throws Exception {
172                                        try {
173                                                return ClassLoader.getSystemClassLoader().loadClass(VMCLASS);
174                                        } catch (final ClassNotFoundException e) {
175                                                for (final File jar : getPotentialToolsJars()) {
176                                                        try {
177                                                                return new URLClassLoader(new URL[] { jar.toURI().toURL() }).loadClass(VMCLASS);
178                                                        } catch (final Throwable t) {
179                                                                logger.trace("Exception while loading tools.jar from " + jar, t);
180                                                        }
181                                                }
182                                        }
183                                        return null;
184                                }
185                        });
186                } catch (final PrivilegedActionException pae) {
187                        final Throwable actual = pae.getCause();
188
189                        if (actual instanceof ClassNotFoundException) {
190                                logger.trace("No VirtualMachine found");
191                                return null;
192                        }
193
194                        throw new RuntimeException("Unexpected checked exception : " + actual);
195                }
196        }
197
198        private static void loadFailed() {
199                System.err.println("Unable to load the java agent dynamically");
200                // FIXME: instructions
201        }
202
203        /**
204         * Attempt to dynamically load the given agent class
205         *
206         * @param agentClass
207         *            the agent class
208         * @throws IOException
209         *             if an error occurs creating the agent jar
210         */
211        public static void loadAgent(Class<?> agentClass) throws IOException {
212                final File tmp = File.createTempFile("agent", ".jar");
213                tmp.deleteOnExit();
214                createAgentJar(tmp, agentClass);
215
216                final String nameOfRunningVM = ManagementFactory.getRuntimeMXBean().getName();
217                final int p = nameOfRunningVM.indexOf('@');
218                final String pid = nameOfRunningVM.substring(0, p);
219
220                final Class<?> vmClass = tryGetVMClass();
221
222                if (vmClass == null) {
223                        loadFailed();
224                } else {
225                        try {
226                                final Method attach = vmClass.getMethod("attach", String.class);
227                                final Method loadAgent = vmClass.getMethod("loadAgent", String.class);
228                                final Method detach = vmClass.getMethod("detach");
229
230                                final Object vm = attach.invoke(null, pid);
231                                try {
232                                        loadAgent.invoke(vm, tmp.getAbsolutePath());
233                                } finally {
234                                        detach.invoke(vm);
235                                }
236                        } catch (final Exception e) {
237                                logger.warn("Loading the agent failed", e);
238                                loadFailed();
239                        }
240                }
241        }
242}