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
028package net.sourceforge.cobertura.instrument;
029
030import java.io.ByteArrayOutputStream;
031import java.io.File;
032import java.io.FileInputStream;
033import java.io.FileNotFoundException;
034import java.io.FileOutputStream;
035import java.io.IOException;
036import java.io.InputStream;
037import java.io.OutputStream;
038import java.util.ArrayList;
039import java.util.Collection;
040import java.util.Iterator;
041import java.util.List;
042import java.util.Vector;
043import java.util.zip.ZipEntry;
044import java.util.zip.ZipInputStream;
045import java.util.zip.ZipOutputStream;
046
047import net.sourceforge.cobertura.coveragedata.CoverageDataFileHandler;
048import net.sourceforge.cobertura.coveragedata.ProjectData;
049import net.sourceforge.cobertura.util.ArchiveUtil;
050import net.sourceforge.cobertura.util.CommandLineBuilder;
051import net.sourceforge.cobertura.util.Header;
052import net.sourceforge.cobertura.util.IOUtil;
053import net.sourceforge.cobertura.util.RegexUtil;
054
055import org.apache.log4j.Logger;
056import org.objectweb.asm.ClassReader;
057import 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 */
087public 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(ClassWriter.COMPUTE_MAXS);
186                                                ClassInstrumenter cv = new ClassInstrumenter(projectData,
187                                                                cw, ignoreRegexes, ignoreBranchesRegexes);
188                                                cr.accept(cv, 0);
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(ClassWriter.COMPUTE_MAXS);
357                        cv = new ClassInstrumenter(projectData, cw, ignoreRegexes, ignoreBranchesRegexes);
358                        cr.accept(cv, 0);
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}