001    package edu.rice.cs.cunit;
002    
003    import edu.rice.cs.cunit.classFile.ClassFile;
004    import edu.rice.cs.cunit.classFile.ClassFileTools;
005    import edu.rice.cs.cunit.classFile.attributes.AAnnotationsAttributeInfo;
006    import edu.rice.cs.cunit.classFile.attributes.AAttributeInfo;
007    import edu.rice.cs.cunit.classFile.attributes.InstrumentationAttributeInfo;
008    import edu.rice.cs.cunit.classFile.attributes.RuntimeInvisibleAnnotationsAttributeInfo;
009    import edu.rice.cs.cunit.classFile.constantPool.APoolInfo;
010    import edu.rice.cs.cunit.classFile.constantPool.ClassPoolInfo;
011    import edu.rice.cs.cunit.classFile.constantPool.AUTFPoolInfo;
012    import edu.rice.cs.cunit.classFile.constantPool.visitors.ADefaultPoolInfoVisitor;
013    import edu.rice.cs.cunit.classFile.constantPool.visitors.CheckUTFVisitor;
014    import edu.rice.cs.cunit.instrumentors.DoNotInstrument;
015    import edu.rice.cs.cunit.instrumentors.IInstrumentationStrategy;
016    import edu.rice.cs.cunit.util.Debug;
017    import edu.rice.cs.cunit.util.FileOps;
018    
019    import java.io.*;
020    import java.util.ArrayList;
021    import java.util.HashMap;
022    import java.util.List;
023    
024    /**
025     * Class loader that instruments loaded classes by applying IInstrumentors to them.
026     *
027     * @author Mathias Ricken
028     */
029    public class InstrumentingClassLoader extends ClassLoader {
030        /**
031         * Should the instrumented classes be saved?
032         */
033        protected boolean _saveInstrumented = true;
034    
035        /**
036         * Make backups of original classes?
037         */
038        protected boolean _makeBackups = false;
039    
040        /**
041         * List of instrumentors that should be applied to a class.
042         */
043        protected ArrayList<IInstrumentationStrategy> _instrumentors = new ArrayList<IInstrumentationStrategy>();
044    
045        /**
046         * Class path to consider when loading classes.
047         */
048        protected ArrayList<String> _classPath = new ArrayList<String>();
049        
050        /**
051         * Hash map of classes already processed. Once processing of a class has begun, it is added with a value of null.
052         * When processing is done, this value is updated to the actual Class object.If a class has not been processed, it
053         * is not present.
054         */
055        protected HashMap<String, Class<?>> _processedClasses = new HashMap<String, Class<?>>();
056    
057        /**
058         * Creates a new instrumenting class loader using specified instrumentors.
059         *
060         * @param instrumentors array of instrumentors
061         */
062        public InstrumentingClassLoader(IInstrumentationStrategy[] instrumentors) {
063            for(IInstrumentationStrategy itor : instrumentors) {
064                _instrumentors.add(itor);
065            }
066            String[] paths = System.getProperty("java.class.path").split(System.getProperty("path.separator"));
067            for(String path: paths) { _classPath.add(path); }
068        }
069    
070        /**
071         * Accessor for the save instrumented flag
072         *
073         * @return flag
074         */
075        public boolean isSaveInstrumented() {
076            return _saveInstrumented;
077        }
078    
079        /**
080         * Mutator for the save instrumented flag
081         *
082         * @param saveInstrumented new flag
083         */
084        public void setSaveInstrumented(boolean saveInstrumented) {
085            _saveInstrumented = saveInstrumented;
086        }
087    
088        /**
089         * Mutator for the make backups flag
090         *
091         * @param makeBackups new flag
092         */
093        public void setMakeBackups(boolean makeBackups) {
094            _makeBackups = makeBackups;
095        }
096    
097        /**
098         * Load and instrument the class with the specified name.
099         *
100         * @param name    class name
101         * @param resolve true if class should be resolved
102         *
103         * @return Class instance for the loaded class
104         *
105         * @throws ClassNotFoundException
106         */
107        public Class loadClass(String name, boolean resolve)
108            throws ClassNotFoundException {
109            SyncPointBuffer.setRecording(true);
110            try {
111                ClassFile cf = null;
112                Class c;
113    
114                Debug.out.print(name);
115                c = findLoadedClass(name);
116                if ((c == null) && (!_processedClasses.containsKey(name))) {
117                    try {
118                        _processedClasses.put(name, null);
119    
120                        // find class on class path
121                        ClassFileTools.ClassLocation cl = ClassFileTools.findClassFile(name, _classPath);
122                        if (cl!=null) {
123                            boolean fromJAR = (cl.getJarFile()!=null);
124                            File origFile = cl.getFile();
125                            byte[] barr = null;
126                            cf = cl.getClassFile();
127                            
128                            try {
129                                if (fromJAR) {
130                                    Debug.out.print(" from jar file "+cl.getJarFile().getName());
131                                }
132                                else {
133                                    Debug.out.print(" from file "+cl.getFile().getPath());
134                                }
135                                
136                                // go ahead and close the file now
137                                try {
138                                    cl.close();
139                                }
140                                catch(IOException e) { /* ignore */ }
141                                cl = null;
142                                    
143                                // find out if we should instrument this class, or if it has already been instrumented
144                                boolean doInstrument = cf.getAttribute(InstrumentationAttributeInfo.NAME) == null;
145                                if (doInstrument) {
146                                    String[] doNotInstrumentPatterns = new String[0];
147                                    AAttributeInfo attr = cf.getAttribute(
148                                        RuntimeInvisibleAnnotationsAttributeInfo.getAttributeName());
149                                    if (null != attr) {
150                                        RuntimeInvisibleAnnotationsAttributeInfo ann
151                                            = (RuntimeInvisibleAnnotationsAttributeInfo)attr;
152                                        for(AAnnotationsAttributeInfo.Annotation a : ann.getAnnotations()) {
153                                            String typeString = ClassFileTools.getTypeString(a.getType(), "");
154                                            if (typeString.substring(0, typeString.length() - 1).equals(
155                                                DoNotInstrument.class.getName())) {
156                                                for(AAnnotationsAttributeInfo.Annotation.NameValuePair nvp : a.getPairs()) {
157                                                    if (nvp.getName().toString().equals("instrumentors")) {
158                                                        AAnnotationsAttributeInfo.Annotation.ConstantMemberValue cmv = nvp.getValue().execute(
159                                                            AAnnotationsAttributeInfo.Annotation.CheckConstantMemberVisitor.singleton(),null);
160                                                        AUTFPoolInfo patternString = cmv.getConstValue().execute(
161                                                            CheckUTFVisitor.singleton(), null);
162                                                        doNotInstrumentPatterns = patternString.toString().split(";");
163                                                        break;
164                                                    }
165                                                }
166                                                break;
167                                            }
168                                        }
169                                    }
170                                    List<IInstrumentationStrategy> appliedInstrumentors
171                                        = new ArrayList<IInstrumentationStrategy>();
172                                    // apply all instrumentors
173                                    for(IInstrumentationStrategy itor : _instrumentors) {
174                                        if (!ClassFileTools.classNameMatches(itor.getClass().getName(),
175                                                                             doNotInstrumentPatterns)) {
176                                            itor.instrument(cf);
177                                            appliedInstrumentors.add(itor);
178                                        }
179                                        else {
180                                            Debug.out.println("Class " + cf.getThisClassName() + " not instrumented with "
181                                                                  + itor.getClass().getName()
182                                                                  + " because it matched @DoNotInstrument instrumentors");
183                                        }
184                                    }
185                                    InstrumentationAttributeInfo.addInstrumentationAttributeInfo(cf, appliedInstrumentors);
186                                    ByteArrayOutputStream baos = new ByteArrayOutputStream();
187                                    cf.write(baos);
188                                    barr = baos.toByteArray();
189    
190                                    if (_saveInstrumented && !fromJAR) {
191                                        // we can only save and make backups if this is coming from a file, not a jar
192                                        if (_makeBackups) {
193                                            File backupFile = new File(origFile.getPath().replace(".class", ".class~"));
194                                            FileOps.copyFile(origFile, backupFile);
195                                        }
196                                        FileOutputStream fos = new FileOutputStream(origFile);
197                                        fos.write(barr);
198                                        fos.close();
199                                    }
200                                    Debug.out.println(" <instrumented>");
201                                }
202                                else {
203                                    Debug.out.println(" <skipped: attribute/annotation>");
204                                    ByteArrayOutputStream baos = new ByteArrayOutputStream();
205                                    cf.write(baos);
206                                    barr = baos.toByteArray();
207                                }
208                                
209                                c = defineClass(name, barr, 0, barr.length);
210                            }
211                            catch(IOException e) {
212                                e.printStackTrace();
213                            }     
214                        }
215                        else {
216                            // not found on classpath
217                            // try to load the class from the system
218                            c = findSystemClass(name);
219                            Debug.out.println(" from system");
220                        }
221                    }
222                    catch(ClassFormatError e) {
223                        System.err.println(e.getMessage());
224                        throw new ClassNotFoundException(name + ":" + e.getMessage(), e);
225                    }
226                }
227                else {
228                    Debug.out.println(" already processed");
229                }
230                if (resolve) {
231                    resolveClass(c);
232                }
233    
234                _processedClasses.put(name, c);
235                
236                return c;
237            }
238            finally {
239                SyncPointBuffer.setRecording(false);
240            }
241        }
242    
243        /**
244         * Private exception used to transport an exception out of the inner class.
245         */
246        private static class RecurIntoClassesException extends RuntimeException {
247            public RecurIntoClassesException(Throwable cause) {
248                super(cause);
249            }
250        }
251    
252        /**
253         * Recur into all the classes defined in the class.
254         *
255         * @param cf      class file
256         * @param resolve resolve classes?
257         *
258         * @throws ClassNotFoundException
259         */
260        protected void recurIntoClasses(ClassFile cf, final boolean resolve) throws ClassNotFoundException {
261            for(APoolInfo cpi : cf.getConstantPool()) {
262                try {
263                    cpi.execute(new ADefaultPoolInfoVisitor<Object, Object>() {
264                        public Object defaultCase(APoolInfo host, Object o) {
265                            // do nothing
266                            return null;
267                        }
268    
269                        public Object classCase(ClassPoolInfo host, Object o) {
270                            String className = ClassFileTools.getClassName(host.getName().toString().replace('/', '.'));
271                            int space = className.indexOf(' ');
272                            if (space >= 0) {
273                                className = className.substring(0, space);
274                            }
275    
276                            if ((!ClassFileTools.isPrimitive(host.getName().toString())) &&
277                                (!className.startsWith("edu.rice.cs.cunit.")) &&
278                                (!_processedClasses.containsKey(className))) {
279                                // load this class
280                                //Debug.out.println("Recur into class "+className);
281                                try {
282                                    Class c = loadClass(className, resolve);
283                                    c = null;
284                                }
285                                catch(ClassNotFoundException e) {
286                                    throw new RecurIntoClassesException(e);
287                                }
288                            }
289                            return null;
290                        }
291                    }, null);
292                }
293                catch(RecurIntoClassesException e) {
294                    throw (ClassNotFoundException)(e.getCause());
295                }
296            }
297        }
298    }
299