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 the
35     * 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 in
45       * 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 parameters
60       * 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 the
83       * 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/"};
94  
95      public Builder ignore(String[] strings) {
96        this.ignore = strings;
97        return this;
98      }
99  
100     private String level = "info";
101 
102     public Builder level(String level) {
103       level = level.toLowerCase();
104       if (level.equals("info") || level.equals("debug")
105           || level.equals("trace")) {
106         this.level = level;
107       } else {
108         if (verbose) {
109           System.err.println("level not info/debug/trace : " + level);
110         }
111       }
112       return this;
113     }
114   }
115 
116   private String level;
117   private String levelEnabled;
118 
119   private LogTransformer(Builder builder) {
120     String s = "WARNING: javassist not available on classpath for javaagent, log statements will not be added";
121     try {
122       if (Class.forName("javassist.ClassPool") == null) {
123         System.err.println(s);
124       }
125     } catch (ClassNotFoundException e) {
126       System.err.println(s);
127     }
128 
129     this.addEntryExit = builder.addEntryExit;
130     // this.addVariableAssignment = builder.addVariableAssignment;
131     this.verbose = builder.verbose;
132     this.ignore = builder.ignore;
133     this.level = builder.level;
134     this.levelEnabled = "is" + builder.level.substring(0, 1).toUpperCase()
135         + builder.level.substring(1) + "Enabled";
136   }
137 
138   private boolean addEntryExit;
139   // private boolean addVariableAssignment;
140   private boolean verbose;
141   private String[] ignore;
142 
143   public byte[] transform(ClassLoader loader, String className, Class<?> clazz,
144       ProtectionDomain domain, byte[] bytes) {
145 
146     try {
147       return transform0(className, clazz, domain, bytes);
148     } catch (Exception e) {
149       System.err.println("Could not instrument " + className);
150       e.printStackTrace();
151       return bytes;
152     }
153   }
154 
155   /**
156    * transform0 sees if the className starts with any of the namespaces to
157    * ignore, if so it is returned unchanged. Otherwise it is processed by
158    * doClass(...)
159    * 
160    * @param className
161    * @param clazz
162    * @param domain
163    * @param bytes
164    * @return
165    */
166 
167   private byte[] transform0(String className, Class<?> clazz,
168       ProtectionDomain domain, byte[] bytes) {
169 
170     try {
171       for (int i = 0; i < ignore.length; i++) {
172         if (className.startsWith(ignore[i])) {
173           return bytes;
174         }
175       }
176       String slf4jName = "org.slf4j.LoggerFactory";
177       try {
178         if (domain != null && domain.getClassLoader() != null) {
179           domain.getClassLoader().loadClass(slf4jName);
180         } else {
181           if (verbose) {
182             System.err.println("Skipping " + className
183                 + " as it doesn't have a domain or a class loader.");
184           }
185           return bytes;
186         }
187       } catch (ClassNotFoundException e) {
188         if (verbose) {
189           System.err.println("Skipping " + className
190               + " as slf4j is not available to it");
191         }
192         return bytes;
193       }
194       if (verbose) {
195         System.err.println("Processing " + className);
196       }
197       return doClass(className, clazz, bytes);
198     } catch (Throwable e) {
199       System.out.println("e = " + e);
200       return bytes;
201     }
202   }
203 
204   private String loggerName;
205 
206   /**
207    * doClass() process a single class by first creates a class description from
208    * the byte codes. If it is a class (i.e. not an interface) the methods
209    * defined have bodies, and a static final logger object is added with the
210    * name of this class as an argument, and each method then gets processed with
211    * doMethod(...) to have logger calls added.
212    * 
213    * @param name
214    *          class name (slashes separate, not dots)
215    * @param clazz
216    * @param b
217    * @return
218    */
219   private byte[] doClass(String name, Class<?> clazz, byte[] b) {
220     ClassPool pool = ClassPool.getDefault();
221     CtClass cl = null;
222     try {
223       cl = pool.makeClass(new ByteArrayInputStream(b));
224       if (cl.isInterface() == false) {
225 
226         loggerName = "_____log";
227 
228         // We have to declare the log variable.
229 
230         String pattern1 = "private static org.slf4j.Logger {};";
231         String loggerDefinition = format(pattern1, loggerName);
232         CtField field = CtField.make(loggerDefinition, cl);
233 
234         // and assign it the appropriate value.
235 
236         String pattern2 = "org.slf4j.LoggerFactory.getLogger({}.class);";
237         String replace = name.replace('/', '.');
238         String getLogger = format(pattern2, replace);
239 
240         cl.addField(field, getLogger);
241 
242         // then check every behaviour (which includes methods). We are only
243         // interested in non-empty ones, as they have code.
244         // NOTE: This will be changed, as empty methods should be
245         // instrumented too.
246 
247         CtBehavior[] methods = cl.getDeclaredBehaviors();
248         for (int i = 0; i < methods.length; i++) {
249           if (methods[i].isEmpty() == false) {
250             doMethod(methods[i]);
251           }
252         }
253         b = cl.toBytecode();
254       }
255     } catch (Exception e) {
256       System.err.println("Could not instrument " + name + ", " + e);
257       e.printStackTrace(System.err);
258     } finally {
259       if (cl != null) {
260         cl.detach();
261       }
262     }
263     return b;
264   }
265 
266   /**
267    * process a single method - this means add entry/exit logging if requested.
268    * It is only called for methods with a body.
269    * 
270    * @param method
271    *          method to work on
272    * @throws NotFoundException
273    * @throws CannotCompileException
274    */
275   private void doMethod(CtBehavior method) throws NotFoundException,
276       CannotCompileException {
277 
278     String signature = JavassistHelper.getSignature(method);
279     String returnValue = JavassistHelper.returnValue(method);
280 
281     if (addEntryExit) {
282       String messagePattern = "if ({}.{}()) {}.{}(\">> {}\");";
283       Object[] arg1 = new Object[] { loggerName, levelEnabled, loggerName,
284           level, signature };
285       String before = MessageFormatter.arrayFormat(messagePattern, arg1);
286       // System.out.println(before);
287       method.insertBefore(before);
288 
289       String messagePattern2 = "if ({}.{}()) {}.{}(\"<< {}{}\");";
290       Object[] arg2 = new Object[] { loggerName, levelEnabled, loggerName,
291           level, signature, returnValue };
292       String after = MessageFormatter.arrayFormat(messagePattern2, arg2);
293       // System.out.println(after);
294       method.insertAfter(after);
295     }
296   }
297 }