1   /**
2    * 
3    */
4   package org.slf4j.instrumentation;
5   
6   import static org.slf4j.helpers.MessageFormatter.format;
7   
8   import java.io.ByteArrayInputStream;
9   import java.lang.instrument.ClassFileTransformer;
10  import java.security.ProtectionDomain;
11  
12  import javassist.CannotCompileException;
13  import javassist.ClassPool;
14  import javassist.CtBehavior;
15  import javassist.CtClass;
16  import javassist.CtField;
17  import javassist.NotFoundException;
18  
19  import org.slf4j.helpers.MessageFormatter;
20  
21  /**
22   * <p>
23   * LogTransformer does the work of analyzing each class, and if appropriate add
24   * log statements to each method to allow logging entry/exit.
25   * </p>
26   * <p>
27   * This class is based on the article <a href="http://today.java.net/pub/a/today/2008/04/24/add-logging-at-class-load-time-with-instrumentation.html"
28   * >Add Logging at Class Load Time with Java Instrumentation</a>.
29   * </p>
30   */
31  public class LogTransformer implements ClassFileTransformer {
32  
33  	/**
34  	 * Builder provides a flexible way of configuring some of many options on
35  	 * the parent class instead of providing many constructors.
36  	 * 
37  	 * {@link http 
38  	 * ://rwhansen.blogspot.com/2007/07/theres-builder-pattern-that-joshua.html}
39  	 * 
40  	 */
41  	public static class Builder {
42  
43  		/**
44  		 * Build and return the LogTransformer corresponding to the options set
45  		 * in this Builder.
46  		 * 
47  		 * @return
48  		 */
49  		public LogTransformer build() {
50  			if (verbose) {
51  				System.err.println("Creating LogTransformer");
52  			}
53  			return new LogTransformer(this);
54  		}
55  
56  		boolean addEntryExit;
57  
58  		/**
59  		 * Should each method log entry (with parameters) and exit (with
60  		 * parameters and returnvalue)?
61  		 * 
62  		 * @param b
63  		 *            value of flag
64  		 * @return
65  		 */
66  		public Builder addEntryExit(boolean b) {
67  			addEntryExit = b;
68  			return this;
69  		}
70  
71  		boolean addVariableAssignment;
72  
73  		// private Builder addVariableAssignment(boolean b) {
74  		// System.err.println("cannot currently log variable assignments.");
75  		// addVariableAssignment = b;
76  		// return this;
77  		// }
78  
79  		boolean verbose;
80  
81  		/**
82  		 * Should LogTransformer be verbose in what it does? This currently list
83  		 * the names of the classes being processed.
84  		 * 
85  		 * @param b
86  		 * @return
87  		 */
88  		public Builder verbose(boolean b) {
89  			verbose = b;
90  			return this;
91  		}
92  
93  		String[] ignore = { "org/slf4j/", "ch/qos/logback/",
94  				"org/apache/log4j/" };
95  
96  		public Builder ignore(String[] strings) {
97  			this.ignore = strings;
98  			return this;
99  		}
100 
101 		private String level = "info";
102 
103 		public Builder level(String level) {
104 			level = level.toLowerCase();
105 			if (level.equals("info") || level.equals("debug")
106 					|| level.equals("trace")) {
107 				this.level = level;
108 			} else {
109 				if (verbose) {
110 					System.err.println("level not info/debug/trace : " + level);
111 				}
112 			}
113 			return this;
114 		}
115 	}
116 
117 	private String level;
118 	private String levelEnabled;
119 
120 	private LogTransformer(Builder builder) {
121 		String s = "WARNING: javassist not available on classpath for javaagent, log statements will not be added";
122 		try {
123 			if (Class.forName("javassist.ClassPool") == null) {
124 				System.err.println(s);
125 			}
126 		} catch (ClassNotFoundException e) {
127 			System.err.println(s);
128 		}
129 
130 		this.addEntryExit = builder.addEntryExit;
131 		// this.addVariableAssignment = builder.addVariableAssignment;
132 		this.verbose = builder.verbose;
133 		this.ignore = builder.ignore;
134 		this.level = builder.level;
135 		this.levelEnabled = "is" + builder.level.substring(0, 1).toUpperCase()
136 				+ builder.level.substring(1) + "Enabled";
137 	}
138 
139 	private boolean addEntryExit;
140 	// private boolean addVariableAssignment;
141 	private boolean verbose;
142 	private String[] ignore;
143 
144 	public byte[] transform(ClassLoader loader, String className,
145 			Class<?> clazz, ProtectionDomain domain, byte[] bytes) {
146 
147 		try {
148 			return transform0(className, clazz, domain, bytes);
149 		} catch (Exception e) {
150 			System.err.println("Could not instrument " + className);
151 			e.printStackTrace();
152 			return bytes;
153 		}
154 	}
155 
156 	/**
157 	 * transform0 sees if the className starts with any of the namespaces to
158 	 * ignore, if so it is returned unchanged. Otherwise it is processed by
159 	 * doClass(...)
160 	 * 
161 	 * @param className
162 	 * @param clazz
163 	 * @param domain
164 	 * @param bytes
165 	 * @return
166 	 */
167 
168 	private byte[] transform0(String className, Class<?> clazz,
169 			ProtectionDomain domain, byte[] bytes) {
170 
171 		try {
172 			for (int i = 0; i < ignore.length; i++) {
173 				if (className.startsWith(ignore[i])) {
174 					return bytes;
175 				}
176 			}
177 			String slf4jName = "org.slf4j.LoggerFactory";
178 			try {
179 				if (domain != null && domain.getClassLoader() != null) {
180 					domain.getClassLoader().loadClass(slf4jName);
181 				} else {
182 					if (verbose) {
183 						System.err
184 								.println("Skipping "
185 										+ className
186 										+ " as it doesn't have a domain or a class loader.");
187 					}
188 					return bytes;
189 				}
190 			} catch (ClassNotFoundException e) {
191 				if (verbose) {
192 					System.err.println("Skipping " + className
193 							+ " as slf4j is not available to it");
194 				}
195 				return bytes;
196 			}
197 			if (verbose) {
198 				System.err.println("Processing " + className);
199 			}
200 			return doClass(className, clazz, bytes);
201 		} catch (Throwable e) {
202 			System.out.println("e = " + e);
203 			return bytes;
204 		}
205 	}
206 
207 	private String loggerName;
208 
209 	/**
210 	 * doClass() process a single class by first creates a class description
211 	 * from the byte codes. If it is a class (i.e. not an interface) the methods
212 	 * defined have bodies, and a static final logger object is added with the
213 	 * name of this class as an argument, and each method then gets processed
214 	 * with doMethod(...) to have logger calls added.
215 	 * 
216 	 * @param name
217 	 *            class name (slashes separate, not dots)
218 	 * @param clazz
219 	 * @param b
220 	 * @return
221 	 */
222 	private byte[] doClass(String name, Class<?> clazz, byte[] b) {
223 		ClassPool pool = ClassPool.getDefault();
224 		CtClass cl = null;
225 		try {
226 			cl = pool.makeClass(new ByteArrayInputStream(b));
227 			if (cl.isInterface() == false) {
228 
229 				loggerName = "_____log";
230 
231 				// We have to declare the log variable.
232 
233 				String pattern1 = "private static org.slf4j.Logger {};";
234 				String loggerDefinition = format(pattern1, loggerName);
235 				CtField field = CtField.make(loggerDefinition, cl);
236 
237 				// and assign it the appropriate value.
238 
239 				String pattern2 = "org.slf4j.LoggerFactory.getLogger({}.class);";
240 				String replace = name.replace('/', '.');
241 				String getLogger = format(pattern2, replace);
242 
243 				cl.addField(field, getLogger);
244 
245 				// then check every behaviour (which includes methods). We are
246 				// only
247 				// interested in non-empty ones, as they have code.
248 				// NOTE: This will be changed, as empty methods should be
249 				// instrumented too.
250 
251 				CtBehavior[] methods = cl.getDeclaredBehaviors();
252 				for (int i = 0; i < methods.length; i++) {
253 					if (methods[i].isEmpty() == false) {
254 						doMethod(methods[i]);
255 					}
256 				}
257 				b = cl.toBytecode();
258 			}
259 		} catch (Exception e) {
260 			System.err.println("Could not instrument " + name + ", " + e);
261 			e.printStackTrace(System.err);
262 		} finally {
263 			if (cl != null) {
264 				cl.detach();
265 			}
266 		}
267 		return b;
268 	}
269 
270 	/**
271 	 * process a single method - this means add entry/exit logging if requested.
272 	 * It is only called for methods with a body.
273 	 * 
274 	 * @param method
275 	 *            method to work on
276 	 * @throws NotFoundException
277 	 * @throws CannotCompileException
278 	 */
279 	private void doMethod(CtBehavior method) throws NotFoundException,
280 			CannotCompileException {
281 
282 		String signature = JavassistHelper.getSignature(method);
283 		String returnValue = JavassistHelper.returnValue(method);
284 
285 		if (addEntryExit) {
286 			String messagePattern = "if ({}.{}()) {}.{}(\">> {}\");";
287 			Object[] arg1 = new Object[] { loggerName, levelEnabled,
288 					loggerName, level, signature };
289 			String before = MessageFormatter.arrayFormat(messagePattern, arg1);
290 			// System.out.println(before);
291 			method.insertBefore(before);
292 
293 			String messagePattern2 = "if ({}.{}()) {}.{}(\"<< {}{}\");";
294 			Object[] arg2 = new Object[] { loggerName, levelEnabled,
295 					loggerName, level, signature, returnValue };
296 			String after = MessageFormatter.arrayFormat(messagePattern2, arg2);
297 			// System.out.println(after);
298 			method.insertAfter(after);
299 		}
300 	}
301 }