001    /*
002     * Cobertura - http://cobertura.sourceforge.net/
003     *
004     * Copyright (C) 2003 jcoverage ltd.
005     * Copyright (C) 2005 Mark Doliner
006     * Copyright (C) 2005 Joakim Erdfelt
007     * Copyright (C) 2005 Grzegorz Lukasik
008     * Copyright (C) 2006 John Lewis
009     * Copyright (C) 2006 Jiri Mares 
010     * Contact information for the above is given in the COPYRIGHT file.
011     *
012     * Cobertura is free software; you can redistribute it and/or modify
013     * it under the terms of the GNU General Public License as published
014     * by the Free Software Foundation; either version 2 of the License,
015     * or (at your option) any later version.
016     *
017     * Cobertura is distributed in the hope that it will be useful, but
018     * WITHOUT ANY WARRANTY; without even the implied warranty of
019     * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
020     * General Public License for more details.
021     *
022     * You should have received a copy of the GNU General Public License
023     * along with Cobertura; if not, write to the Free Software
024     * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
025     * USA
026     */
027    
028    package net.sourceforge.cobertura.instrument;
029    
030    import java.io.ByteArrayOutputStream;
031    import java.io.File;
032    import java.io.FileInputStream;
033    import java.io.FileNotFoundException;
034    import java.io.FileOutputStream;
035    import java.io.IOException;
036    import java.io.InputStream;
037    import java.io.OutputStream;
038    import java.util.ArrayList;
039    import java.util.Collection;
040    import java.util.Iterator;
041    import java.util.List;
042    import java.util.Vector;
043    import java.util.zip.ZipEntry;
044    import java.util.zip.ZipInputStream;
045    import java.util.zip.ZipOutputStream;
046    
047    import net.sourceforge.cobertura.coveragedata.CoverageDataFileHandler;
048    import net.sourceforge.cobertura.coveragedata.ProjectData;
049    import net.sourceforge.cobertura.util.ArchiveUtil;
050    import net.sourceforge.cobertura.util.CommandLineBuilder;
051    import net.sourceforge.cobertura.util.Header;
052    import net.sourceforge.cobertura.util.IOUtil;
053    import net.sourceforge.cobertura.util.RegexUtil;
054    
055    import org.apache.log4j.Logger;
056    import org.objectweb.asm.ClassReader;
057    import org.objectweb.asm.ClassWriter;
058    
059    /**
060     * <p>
061     * Add coverage instrumentation to existing classes.
062     * </p>
063     *
064     * <h3>What does that mean, exactly?</h3>
065     * <p>
066     * It means Cobertura will look at each class you give it.  It
067     * loads the bytecode into memory.  For each line of source,
068     * Cobertura adds a few extra instructions.  These instructions 
069     * do the following:
070     * </p>
071     * 
072     * <ol>
073     * <li>Get an instance of the ProjectData class.</li>
074     * <li>Call a method in this ProjectData class that increments
075     * a counter for this line of code.
076     * </ol>
077     *
078     * <p>
079     * After every line in a class has been "instrumented," Cobertura
080     * edits the bytecode for the class one more time and adds
081     * "implements net.sourceforge.cobertura.coveragedata.HasBeenInstrumented" 
082     * This is basically just a flag used internally by Cobertura to
083     * determine whether a class has been instrumented or not, so
084     * as not to instrument the same class twice.
085     * </p>
086     */
087    public class Main
088    {
089    
090            private static final Logger logger = Logger.getLogger(Main.class);
091    
092            private File destinationDirectory = null;
093    
094            private Collection ignoreRegexes = new Vector();
095    
096            private Collection ignoreBranchesRegexes = new Vector();
097    
098            private ClassPattern classPattern = new ClassPattern();
099    
100            private ProjectData projectData = null;
101    
102            /**
103             * @param entry A zip entry.
104             * @return True if the specified entry has "class" as its extension,
105             * false otherwise.
106             */
107            private static boolean isClass(ZipEntry entry)
108            {
109                    return entry.getName().endsWith(".class");
110            }
111    
112            private boolean addInstrumentationToArchive(CoberturaFile file, InputStream archive,
113                            OutputStream output) throws Exception
114            {
115                    ZipInputStream zis = null;
116                    ZipOutputStream zos = null;
117    
118                    try
119                    {
120                            zis = new ZipInputStream(archive);
121                            zos = new ZipOutputStream(output);
122                            return addInstrumentationToArchive(file, zis, zos);
123                    }
124                    finally
125                    {
126                            zis = (ZipInputStream)IOUtil.closeInputStream(zis);
127                            zos = (ZipOutputStream)IOUtil.closeOutputStream(zos);
128                    }
129            }
130    
131            private boolean addInstrumentationToArchive(CoberturaFile file, ZipInputStream archive,
132                            ZipOutputStream output) throws Exception
133            {
134                    /*
135                     * "modified" is returned and indicates that something was instrumented.
136                     * If nothing is instrumented, the original entry will be used by the
137                     * caller of this method.
138                     */
139                    boolean modified = false;
140                    ZipEntry entry;
141                    while ((entry = archive.getNextEntry()) != null)
142                    {
143                            try
144                            {
145                                    String entryName = entry.getName();
146    
147                                    /*
148                                     * If this is a signature file then don't copy it,
149                                     * but don't set modified to true.  If the only
150                                     * thing we do is strip the signature, just use
151                                     * the original entry.
152                                     */
153                                    if (ArchiveUtil.isSignatureFile(entry.getName()))
154                                    {
155                                            continue;
156                                    }
157                                    ZipEntry outputEntry = new ZipEntry(entry.getName());
158                                    outputEntry.setComment(entry.getComment());
159                                    outputEntry.setExtra(entry.getExtra());
160                                    outputEntry.setTime(entry.getTime());
161                                    output.putNextEntry(outputEntry);
162    
163                                    // Read current entry
164                                    byte[] entryBytes = IOUtil
165                                                    .createByteArrayFromInputStream(archive);
166    
167                                    // Instrument embedded archives if a classPattern has been specified
168                                    if ((classPattern.isSpecified()) && ArchiveUtil.isArchive(entryName))
169                                    {
170                                            Archive archiveObj = new Archive(file, entryBytes);
171                                            addInstrumentationToArchive(archiveObj);
172                                            if (archiveObj.isModified())
173                                            {
174                                                    modified = true;
175                                                    entryBytes = archiveObj.getBytes();
176                                                    outputEntry.setTime(System.currentTimeMillis());
177                                            }
178                                    }
179                                    else if (isClass(entry) && classPattern.matches(entryName))
180                                    {
181                                            try
182                                            {
183                                                    // Instrument class
184                                                    ClassReader cr = new ClassReader(entryBytes);
185                                                    ClassWriter cw = new ClassWriter(true);
186                                                    ClassInstrumenter cv = new ClassInstrumenter(projectData,
187                                                                    cw, ignoreRegexes, ignoreBranchesRegexes);
188                                                    cr.accept(cv, false);
189            
190                                                    // If class was instrumented, get bytes that define the
191                                                    // class
192                                                    if (cv.isInstrumented())
193                                                    {
194                                                            logger.debug("Putting instrumented entry: "
195                                                                            + entry.getName());
196                                                            entryBytes = cw.toByteArray();
197                                                            modified = true;
198                                                            outputEntry.setTime(System.currentTimeMillis());
199                                                    }
200                                            }
201                                            catch (Throwable t)
202                                            {
203                                                    if (entry.getName().endsWith("_Stub.class"))
204                                                    {
205                                                            //no big deal - it is probably an RMI stub, and they don't need to be instrumented
206                                                            logger.debug("Problems instrumenting archive entry: " + entry.getName(), t);
207                                                    }
208                                                    else
209                                                    {
210                                                            logger.warn("Problems instrumenting archive entry: " + entry.getName(), t);
211                                                    }
212                                            }
213                                    }
214    
215                                    // Add entry to the output
216                                    output.write(entryBytes);
217                                    output.closeEntry();
218                                    archive.closeEntry();
219                            }
220                            catch (Exception e)
221                            {
222                                    logger.warn("Problems with archive entry: " + entry.getName(), e);
223                            }
224                            catch (Throwable t)
225                            {
226                                    logger.warn("Problems with archive entry: " + entry.getName(), t);
227                            }
228                            output.flush();
229                    }
230                    return modified;
231            }
232    
233            private void addInstrumentationToArchive(Archive archive) throws Exception
234            {
235                    InputStream in = null;
236                    ByteArrayOutputStream out = null;
237                    try
238                    {
239                            in = archive.getInputStream();
240                            out = new ByteArrayOutputStream();
241                            boolean modified = addInstrumentationToArchive(archive.getCoberturaFile(), in, out);
242    
243                            if (modified)
244                            {
245                                    out.flush();
246                                    byte[] bytes = out.toByteArray();
247                                    archive.setModifiedBytes(bytes);
248                            }
249                    }
250                    finally
251                    {
252                            in = IOUtil.closeInputStream(in);
253                            out = (ByteArrayOutputStream)IOUtil.closeOutputStream(out);
254                    }
255            }
256    
257            private void addInstrumentationToArchive(CoberturaFile archive)
258            {
259                    logger.debug("Instrumenting archive " + archive.getAbsolutePath());
260    
261                    File outputFile = null;
262                    ZipInputStream input = null;
263                    ZipOutputStream output = null;
264                    boolean modified = false;
265                    try
266                    {
267                            // Open archive
268                            try
269                            {
270                                    input = new ZipInputStream(new FileInputStream(archive));
271                            }
272                            catch (FileNotFoundException e)
273                            {
274                                    logger.warn("Cannot open archive file: "
275                                                    + archive.getAbsolutePath(), e);
276                                    return;
277                            }
278    
279                            // Open output archive
280                            try
281                            {
282                                    // check if destination folder is set
283                                    if (destinationDirectory != null)
284                                    {
285                                            // if so, create output file in it
286                                            outputFile = new File(destinationDirectory, archive.getPathname());
287                                    }
288                                    else
289                                    {
290                                            // otherwise create output file in temporary location
291                                            outputFile = File.createTempFile(
292                                                            "CoberturaInstrumentedArchive", "jar");
293                                            outputFile.deleteOnExit();
294                                    }
295                                    output = new ZipOutputStream(new FileOutputStream(outputFile));
296                            }
297                            catch (IOException e)
298                            {
299                                    logger.warn("Cannot open file for instrumented archive: "
300                                                    + archive.getAbsolutePath(), e);
301                                    return;
302                            }
303    
304                            // Instrument classes in archive
305                            try
306                            {
307                                    modified = addInstrumentationToArchive(archive, input, output);
308                            }
309                            catch (Throwable e)
310                            {
311                                    logger.warn("Cannot instrument archive: "
312                                                    + archive.getAbsolutePath(), e);
313                                    return;
314                            }
315                    }
316                    finally
317                    {
318                            input = (ZipInputStream)IOUtil.closeInputStream(input);
319                            output = (ZipOutputStream)IOUtil.closeOutputStream(output);
320                    }
321    
322                    // If destination folder was not set, overwrite orginal archive with
323                    // instrumented one
324                    if (modified && (destinationDirectory == null))
325                    {
326                            try
327                            {
328                                    logger.debug("Moving " + outputFile.getAbsolutePath() + " to "
329                                                    + archive.getAbsolutePath());
330                                    IOUtil.moveFile(outputFile, archive);
331                            }
332                            catch (IOException e)
333                            {
334                                    logger.warn("Cannot instrument archive: "
335                                                    + archive.getAbsolutePath(), e);
336                                    return;
337                            }
338                    }
339                    if ((destinationDirectory != null) && (!modified))
340                    {
341                            outputFile.delete();
342                    }
343            }
344    
345            private void addInstrumentationToSingleClass(File file)
346            {
347                    logger.debug("Instrumenting class " + file.getAbsolutePath());
348    
349                    InputStream inputStream = null;
350                    ClassWriter cw;
351                    ClassInstrumenter cv;
352                    try
353                    {
354                            inputStream = new FileInputStream(file);
355                            ClassReader cr = new ClassReader(inputStream);
356                            cw = new ClassWriter(true);
357                            cv = new ClassInstrumenter(projectData, cw, ignoreRegexes, ignoreBranchesRegexes);
358                            cr.accept(cv, false);
359                    }
360                    catch (Throwable t)
361                    {
362                            logger.warn("Unable to instrument file " + file.getAbsolutePath(),
363                                            t);
364                            return;
365                    }
366                    finally
367                    {
368                            inputStream = IOUtil.closeInputStream(inputStream);
369                    }
370    
371                    OutputStream outputStream = null;
372                    try
373                    {
374                            if (cv.isInstrumented())
375                            {
376                                    // If destinationDirectory is null, then overwrite
377                                    // the original, uninstrumented file.
378                                    File outputFile;
379                                    if (destinationDirectory == null)
380                                            outputFile = file;
381                                    else
382                                            outputFile = new File(destinationDirectory, cv
383                                                            .getClassName().replace('.', File.separatorChar)
384                                                            + ".class");
385    
386                                    File parentFile = outputFile.getParentFile();
387                                    if (parentFile != null)
388                                    {
389                                            parentFile.mkdirs();
390                                    }
391    
392                                    byte[] instrumentedClass = cw.toByteArray();
393                                    outputStream = new FileOutputStream(outputFile);
394                                    outputStream.write(instrumentedClass);
395                            }
396                    }
397                    catch (Throwable t)
398                    {
399                            logger.warn("Unable to instrument file " + file.getAbsolutePath(),
400                                            t);
401                            return;
402                    }
403                    finally
404                    {
405                            outputStream = IOUtil.closeOutputStream(outputStream);
406                    }
407            }
408    
409            // TODO: Don't attempt to instrument a file if the outputFile already
410            //       exists and is newer than the input file, and the output and
411            //       input file are in different locations?
412            private void addInstrumentation(CoberturaFile coberturaFile)
413            {
414                    if (coberturaFile.isClass() && classPattern.matches(coberturaFile.getPathname()))
415                    {
416                            addInstrumentationToSingleClass(coberturaFile);
417                    }
418                    else if (coberturaFile.isDirectory())
419                    {
420                            String[] contents = coberturaFile.list();
421                            for (int i = 0; i < contents.length; i++)
422                            {
423                                    File relativeFile = new File(coberturaFile.getPathname(), contents[i]);
424                                    CoberturaFile relativeCoberturaFile = new CoberturaFile(coberturaFile.getBaseDir(),
425                                                    relativeFile.toString());
426                                    //recursion!
427                                    addInstrumentation(relativeCoberturaFile);
428                            }
429                    }
430            }
431    
432            private void parseArguments(String[] args)
433            {
434                    File dataFile = CoverageDataFileHandler.getDefaultDataFile();
435    
436                    // Parse our parameters
437                    List filePaths = new ArrayList();
438                    String baseDir = null;
439                    for (int i = 0; i < args.length; i++)
440                    {
441                            if (args[i].equals("--basedir"))
442                                    baseDir = args[++i];
443                            else if (args[i].equals("--datafile"))
444                                    dataFile = new File(args[++i]);
445                            else if (args[i].equals("--destination"))
446                                    destinationDirectory = new File(args[++i]);
447                            else if (args[i].equals("--ignore"))
448                            {
449                                    RegexUtil.addRegex(ignoreRegexes, args[++i]);
450                            }
451                            else if (args[i].equals("--ignoreBranches"))
452                            {
453                                    RegexUtil.addRegex(ignoreBranchesRegexes, args[++i]);
454                            }
455                            else if (args[i].equals("--includeClasses"))
456                            {
457                                    classPattern.addIncludeClassesRegex(args[++i]);
458                            }
459                            else if (args[i].equals("--excludeClasses"))
460                            {
461                                    classPattern.addExcludeClassesRegex(args[++i]);
462                            }
463                            else
464                            {
465                                    CoberturaFile coberturaFile = new CoberturaFile(baseDir, args[i]);
466                                    filePaths.add(coberturaFile);
467                            }
468                    }
469    
470                    // Load coverage data
471                    if (dataFile.isFile())
472                            projectData = CoverageDataFileHandler.loadCoverageData(dataFile);
473                    if (projectData == null)
474                            projectData = new ProjectData();
475                    
476                    // Instrument classes
477                    System.out.println("Instrumenting "     + filePaths.size() + " "
478                                    + (filePaths.size() == 1 ? "file" : "files")
479                                    + (destinationDirectory != null ? " to "
480                                                    + destinationDirectory.getAbsoluteFile() : ""));
481    
482                    Iterator iter = filePaths.iterator();
483                    while (iter.hasNext())
484                    {
485                            CoberturaFile coberturaFile = (CoberturaFile)iter.next();
486                            if (coberturaFile.isArchive())
487                            {
488                                    addInstrumentationToArchive(coberturaFile);
489                            }
490                            else
491                            {
492                                    addInstrumentation(coberturaFile);
493                            }
494                    }
495    
496                    // Save coverage data
497                    CoverageDataFileHandler.saveCoverageData(projectData, dataFile);
498            }
499    
500            public static void main(String[] args)
501            {
502                    Header.print(System.out);
503    
504                    long startTime = System.currentTimeMillis();
505    
506                    Main main = new Main();
507    
508                    try {
509                            args = CommandLineBuilder.preprocessCommandLineArguments( args);
510                    } catch( Exception ex) {
511                            System.err.println( "Error: Cannot process arguments: " + ex.getMessage());
512                            System.exit(1);
513                    }
514                    main.parseArguments(args);
515    
516                    long stopTime = System.currentTimeMillis();
517                    System.out.println("Instrument time: " + (stopTime - startTime) + "ms");
518            }
519    
520    }