001: /**
002: *
003: * Licensed to the Apache Software Foundation (ASF) under one or more
004: * contributor license agreements. See the NOTICE file distributed with
005: * this work for additional information regarding copyright ownership.
006: * The ASF licenses this file to You under the Apache License, Version 2.0
007: * (the "License"); you may not use this file except in compliance with
008: * the License. You may obtain a copy of the License at
009: *
010: * http://www.apache.org/licenses/LICENSE-2.0
011: *
012: * Unless required by applicable law or agreed to in writing, software
013: * distributed under the License is distributed on an "AS IS" BASIS,
014: * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015: * See the License for the specific language governing permissions and
016: * limitations under the License.
017: */package org.apache.openejb.javaagent;
018:
019: import java.io.File;
020: import java.io.FileOutputStream;
021: import java.io.IOException;
022: import java.io.InputStream;
023: import java.io.PrintWriter;
024: import java.io.OutputStreamWriter;
025: import java.io.Closeable;
026: import java.lang.instrument.Instrumentation;
027: import java.lang.management.ManagementFactory;
028: import java.lang.management.RuntimeMXBean;
029: import java.lang.reflect.Field;
030: import java.lang.reflect.Method;
031: import java.lang.reflect.ReflectPermission;
032: import java.net.JarURLConnection;
033: import java.net.URL;
034: import java.net.URLConnection;
035: import java.security.Permission;
036: import java.util.zip.ZipEntry;
037: import java.util.zip.ZipOutputStream;
038:
039: public class Agent {
040: private static final Permission ACCESS_PERMISSION = new ReflectPermission(
041: "suppressAccessChecks");
042: private static String agentArgs;
043: private static Instrumentation instrumentation;
044: private static boolean initialized = false;
045:
046: public static void premain(String agentArgs,
047: Instrumentation instrumentation) {
048: if (Agent.instrumentation != null)
049: return;
050:
051: Agent.agentArgs = agentArgs;
052: Agent.instrumentation = instrumentation;
053: initialized = true;
054: }
055:
056: public static void agentmain(String agentArgs,
057: Instrumentation instrumentation) {
058: if (Agent.instrumentation != null)
059: return;
060:
061: Agent.agentArgs = agentArgs;
062: Agent.instrumentation = instrumentation;
063: initialized = true;
064: }
065:
066: public static synchronized String getAgentArgs() {
067: SecurityManager sm = System.getSecurityManager();
068: if (sm != null)
069: sm.checkPermission(ACCESS_PERMISSION);
070: checkInitialization();
071: return agentArgs;
072: }
073:
074: /**
075: * Gets the instrumentation instance.
076: * You must have java.lang.ReflectPermission(suppressAccessChecks) to call this method
077: * @return the instrumentation instance
078: */
079: public static synchronized Instrumentation getInstrumentation() {
080: SecurityManager sm = System.getSecurityManager();
081: if (sm != null)
082: sm.checkPermission(ACCESS_PERMISSION);
083: checkInitialization();
084: return instrumentation;
085: }
086:
087: private static synchronized void checkInitialization() {
088: if (!initialized) {
089: try {
090: checkSystemClassPath();
091: dynamicLoadAgent();
092: } catch (Exception e) {
093: new IllegalStateException("Unable to initialize agent",
094: e).printStackTrace();
095: } finally {
096: initialized = true;
097: }
098: }
099: }
100:
101: private static void checkSystemClassPath()
102: throws NoSuchFieldException, IllegalAccessException {
103: if (instrumentation != null)
104: return;
105:
106: Class<?> systemAgentClass = null;
107: try {
108: ClassLoader systemCl = ClassLoader.getSystemClassLoader();
109: systemAgentClass = systemCl
110: .loadClass(Agent.class.getName());
111: } catch (ClassNotFoundException e) {
112: // java-agent jar was not on the system class path
113: return;
114: }
115:
116: Field instrumentationField = systemAgentClass
117: .getDeclaredField("instrumentation");
118: instrumentationField.setAccessible(true);
119: instrumentation = (Instrumentation) instrumentationField
120: .get(null);
121:
122: Field agentArgsField = systemAgentClass
123: .getDeclaredField("agentArgs");
124: agentArgsField.setAccessible(true);
125: agentArgs = (String) agentArgsField.get(null);
126: }
127:
128: private static void dynamicLoadAgent() throws Exception {
129: if (instrumentation != null)
130: return;
131:
132: try {
133: Class<?> vmClass = Class
134: .forName("com.sun.tools.attach.VirtualMachine");
135: Method attachMethod = vmClass.getMethod("attach",
136: String.class);
137: Method loadAgentMethod = vmClass.getMethod("loadAgent",
138: String.class);
139:
140: // find the agentJar
141: String agentPath = getAgentJar();
142:
143: // get the pid of the current process (for attach command)
144: String pid = getPid();
145:
146: // attach to the vm
147: Object vm = attachMethod.invoke(null, new String[] { pid });
148:
149: // load our agent
150: loadAgentMethod.invoke(vm, agentPath);
151:
152: // The AgentJar is loaded into the system classpath, and this class could
153: // be in a child classloader, so we need to double check the system classpath
154: checkSystemClassPath();
155: } catch (ClassNotFoundException e) {
156: // not a Sun VM
157: } catch (NoSuchMethodException e) {
158: // not a Sun VM
159: }
160: }
161:
162: private static String getPid() {
163: // This relies on the undocumented convention of the
164: // RuntimeMXBean's name starting with the PID, but
165: // there appears to be no other way to obtain the
166: // current process' id, which we need for the attach
167: // process
168: RuntimeMXBean bean = ManagementFactory.getRuntimeMXBean();
169: String pid = bean.getName();
170: if (pid.indexOf("@") != -1) {
171: pid = pid.substring(0, pid.indexOf("@"));
172: }
173: return pid;
174: }
175:
176: /**
177: * Try to find the openejb-javaagent jar, and if not found create a new jar
178: * file for the sole purpose of specifying an Agent-Class to load into the JVM.
179: */
180: private static String getAgentJar() throws IOException {
181: URL resource = Agent.class.getClassLoader().getResource(
182: Agent.class.getName().replace('.', '/') + ".class");
183: if (resource == null) {
184: throw new IllegalStateException(
185: "Could not find Agent class file in class path");
186: }
187:
188: URLConnection urlConnection = resource.openConnection();
189: if (urlConnection instanceof JarURLConnection) {
190: JarURLConnection jarURLConnection = (JarURLConnection) urlConnection;
191: return jarURLConnection.getJarFile().getName();
192: }
193:
194: InputStream in = urlConnection.getInputStream();
195: ZipOutputStream out = null;
196: File file = null;
197: try {
198: file = File.createTempFile(Agent.class.getName(), ".jar");
199: file.deleteOnExit();
200:
201: out = new ZipOutputStream(new FileOutputStream(file));
202:
203: // write manifest
204: out.putNextEntry(new ZipEntry("META-INF/MANIFEST.MF"));
205: try {
206: PrintWriter writer = new PrintWriter(
207: new OutputStreamWriter(out));
208: writer.println("Agent-Class: " + Agent.class.getName());
209: writer.println("Can-Redefine-Classes: true");
210: writer.println("Can-Retransform-Classes: true");
211: writer.flush();
212: } finally {
213: out.closeEntry();
214: }
215:
216: // write agent class
217: out.putNextEntry(new ZipEntry(Agent.class.getName()
218: .replace('.', '/')
219: + ".class"));
220: try {
221: byte[] buffer = new byte[4096];
222: for (int count = in.read(buffer); count >= 0; count = in
223: .read(buffer)) {
224: out.write(buffer, 0, count);
225: }
226: } finally {
227: out.closeEntry();
228: }
229:
230: return file.getAbsolutePath();
231: } catch (IOException e) {
232: if (file != null) {
233: file.delete();
234: }
235: throw e;
236: } finally {
237: close(in);
238: close(out);
239: }
240: }
241:
242: private static void close(Closeable closeable) {
243: if (closeable != null) {
244: try {
245: closeable.close();
246: } catch (IOException ignored) {
247: }
248: }
249: }
250: }
|