com.ibm.wala.shrike.shrikeBT.tools.OfflineInstrumenterBase Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of com.ibm.wala.shrike Show documentation
T. J. Watson Libraries for Analysis
The newest version!
/*
 * Copyright (c) 2002,2006 IBM Corporation.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 *     IBM Corporation - initial API and implementation
 */
package com.ibm.wala.shrike.shrikeBT.tools;

import com.ibm.wala.shrike.shrikeBT.analysis.ClassHierarchyProvider;
import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.BitSet;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.jar.JarOutputStream;
import java.util.zip.ZipEntry;

/**
 * This class provides functionality for performing offline instrumentation. It is subclassed with
 * class-toolkit-specific functionality.
 */
public abstract class OfflineInstrumenterBase {
  private int inputIndex;

  private final HashSet entryNames = new HashSet<>();

  private final ArrayList inputs = new ArrayList<>();

  private final BitSet ignoringInputs = new BitSet();

  private File outputFile;

  private boolean passUnmodifiedClasses = false;

  private JarOutputStream outputJar;

  private JarFile cachedJar;

  private File cachedJarFile;

  private ManifestBuilder manifestBuilder;

  protected ClassHierarchyProvider cha;

  /**
   * This installs a ManifestBuilder callback that this class will notify whenever an entry has been
   * added to the output zip file.
   */
  public void setManifestBuilder(ManifestBuilder mb) {
    manifestBuilder = mb;
  }

  /** Thiscallback is notified whenever an entry has been added to the output zip file. */
  public interface ManifestBuilder {
    void addEntry(ZipEntry ze);
  }

  /**
   * This class represents a resource which can be opened and read; either a file or a JAR entry.
   */
  abstract static class Input {
    private String className;

    /** Tell us what the classname is supposed to be, if it's a class file. */
    public final void setClassName(String c) {
      className = c.intern();
    }

    /** Returns the classname if it has been set. */
    public final String getClassName() {
      return className;
    }

    /** get name of resource used for input */
    public abstract String getInputName();

    /** Open the resource for reading as a stream. */
    public abstract InputStream open() throws IOException;

    /**
     * @return true if this resource represents a class, false otherwise
     */
    public boolean isClass() {
      return true;
    }
  }

  /**
   * This class represents a JAR file entry. It might or might not be a class; we support non-class
   * JAR resources so that we can copy them to the output JAR if the client requests that.
   */
  final class JarInput extends Input {
    private final File file;

    private final String name;

    /** Select a particular entry from a JAR file on disk. */
    public JarInput(File f, String je) {
      file = f;
      name = je;
    }

    @Override
    @SuppressWarnings("resource")
    public InputStream open() throws IOException {
      JarFile cachedJar = openCachedJar(file);
      return cachedJar.getInputStream(cachedJar.getEntry(name));
    }

    @Override
    public String toString() {
      return file.getPath() + '#' + name;
    }

    @Override
    public boolean isClass() {
      return name.endsWith(".class");
    }

    @Override
    public String getInputName() {
      return name;
    }

    /** Get the underlying ZipEntry corresponding to this resource. */
    @SuppressWarnings("resource")
    public ZipEntry getEntry() throws IOException {
      JarFile cachedJar = openCachedJar(file);
      return cachedJar.getEntry(name);
    }
  }

  /**
   * Open a JAR/ZIP file. This routine caches the last JAR file opened to save effort when the same
   * file is accessed again and again. DO NOT close the file returned by this routine until you've
   * finished with this OfflineInstrumente completely. Also, this JarFile will be closed the next
   * time someone calls openCachedJar.
   */
  private JarFile openCachedJar(File file) throws IOException {
    if (cachedJarFile == null || !cachedJarFile.equals(file)) {
      if (cachedJar != null) {
        cachedJar.close();
      }
      cachedJarFile = file;
      cachedJar = new JarFile(file, false);
    }
    return cachedJar;
  }

  /**
   * This class represents a plain old class file in the filesystem. Non-class file resources are
   * not supported.
   */
  static final class ClassInput extends Input {
    private final File file;
    private final File baseDirectory;

    public ClassInput(File baseDirectory, File f) {
      file = f;
      this.baseDirectory = baseDirectory;
    }

    @Override
    public InputStream open() throws IOException {
      return new FileInputStream(file);
    }

    @Override
    public String toString() {
      return file.getPath();
    }

    @Override
    public String getInputName() {
      int base = baseDirectory.getPath().length() + 1;
      return file.getPath().substring(base);
    }
  }

  protected OfflineInstrumenterBase() {}

  public void setClassHierarchyProvider(ClassHierarchyProvider cha) {
    this.cha = cha;
  }

  /** Set the file in which instrumented classes will be deposited. */
  public final void setOutputJar(File f) {
    outputFile = f;
  }

  /** Indicate whether classes which are not modified will be put into the output jar anyway. */
  public final void setPassUnmodifiedClasses(boolean pass) {
    passUnmodifiedClasses = pass;
  }

  /** Add a JAR file containing source classes to instrument. */
  public final void addInputJar(File f) throws IOException {
    try (final JarFile jf = new JarFile(f, false)) {
      for (Enumeration e = jf.entries(); e.hasMoreElements(); ) {
        JarEntry entry = e.nextElement();
        String name = entry.getName();
        inputs.add(new JarInput(f, name));
      }
    }
  }

  /** Add a JAR entry containing a source class to instrument. */
  public final void addInputJarEntry(File f, String name) {
    inputs.add(new JarInput(f, name));
  }

  /** Add a class file containing a source class to instrument. */
  public final void addInputClass(File baseDirectory, File f) {
    inputs.add(new ClassInput(baseDirectory, f));
  }

  /**
   * Add a directory containing class files to instrument. All subdirectories are also scanned.
   *
   * @throws IllegalArgumentException if d is null
   */
  public final void addInputDirectory(File baseDirectory, File d)
      throws IOException, IllegalArgumentException {
    if (d == null) {
      throw new IllegalArgumentException("d is null");
    }
    File[] fs = d.listFiles(f -> f.isDirectory() || f.getName().endsWith(".class"));
    if (fs == null) {
      throw new IllegalArgumentException("bad directory " + d.getAbsolutePath());
    }
    for (File f : fs) {
      if (f.isDirectory()) {
        addInputDirectory(baseDirectory, f);
      } else {
        addInputClass(baseDirectory, f);
      }
    }
  }

  /**
   * Add something to instrument --- the name of a JAR file, a class file, a directory or an entry
   * within a jar file (as filename#entryname). If we can't identify it, nothing is added and we
   * return false.
   *
   * @throws IllegalArgumentException if a is null
   */
  public final boolean addInputElement(File baseDirectory, String a) throws IOException {
    if (a == null) {
      throw new IllegalArgumentException("a is null");
    }
    try {
      int poundIndex = a.indexOf('#');
      if (poundIndex > 0) {
        addInputJarEntry(new File(a.substring(0, poundIndex)), a.substring(poundIndex + 1));
        return true;
      }
      File f = new File(a);
      if (f.isDirectory()) {
        addInputDirectory(baseDirectory, f);
        return true;
      } else if (f.exists()) {
        if (a.endsWith(".class")) {
          addInputClass(baseDirectory, f);
          return true;
        } else if (a.endsWith(".jar") || a.endsWith(".zip")) {
          addInputJar(new File(a));
          return true;
        }
      }
    } catch (IOException ex) {
      throw new IOException("Error reading input element '" + a + "': " + ex.getMessage());
    }
    return false;
  }

  /**
   * Parse an argument list to find elements to instrument and the name of the output file. The "-o
   * filename" option selects the output JAR file name. Any other argument not starting with "-" is
   * added to the list of elements to instrument, if it appears to be the name of a class file, JAR
   * file, or directory. If any argument starting with "--" is encountered, the rest of the
   * command-line is considered leftover
   *
   * @return the arguments that were not understood
   * @throws IllegalArgumentException if args == null
   */
  public final String[] parseStandardArgs(String[] args)
      throws IllegalArgumentException, IOException {
    if (args == null) {
      throw new IllegalArgumentException("args == null");
    }
    ArrayList leftover = new ArrayList<>();

    for (int i = 0; i < args.length; i++) {
      String a = args[i];
      if (a == null) {
        throw new IllegalArgumentException("args[" + i + "] is null");
      }
      if (a.equals("-o") && i + 1 < args.length) {
        setOutputJar(new File(args[i + 1]));
        i++;
        continue;
      } else if (!a.startsWith("-")) {
        if (addInputElement(new File(a), a)) {
          continue;
        }
      } else if (a.startsWith("--")) {
        leftover.addAll(Arrays.asList(args).subList(i, args.length));
        break;
      }
      leftover.add(a);
    }

    String[] r = new String[leftover.size()];
    leftover.toArray(r);
    return r;
  }

  /**
   * @return the number of source classes to be instrumented
   */
  public final int getNumInputClasses() {
    return inputs.size();
  }

  /** Start traversing the source class list from the beginning. */
  public final void beginTraversal() {
    inputIndex = 0;
  }

  protected abstract Object makeClassFromStream(String inputName, BufferedInputStream s)
      throws IOException;

  protected abstract String getClassName(Object cl);

  protected abstract void writeClassTo(Object cl, Object mods, OutputStream s) throws IOException;

  protected final Object internalNextClass() throws IOException {
    while (true) {
      if (inputIndex >= inputs.size()) {
        return null;
      } else {
        Input in = inputs.get(inputIndex);
        inputIndex++;
        if (ignoringInputs.get(inputIndex - 1) || !in.isClass()) {
          continue;
        }
        try (final BufferedInputStream s = new BufferedInputStream(in.open())) {
          Object r = makeClassFromStream(in.getInputName(), s);
          String name = getClassName(r);
          in.setClassName(name);
          return r;
        }
      }
    }
  }

  private static String toEntryName(String className) {
    return className.replace('.', '/') + ".class";
  }

  /**
   * Get the name of the resource containing the last class returned. This is either a file name
   * (e.g., "com/ibm/Main.class"), or a JAR entry name (e.g., "apps/app.jar#com/ibm/Main.class").
   *
   * @return the resource name, or null if no class has been returned yet
   */
  public final String getLastClassResourceName() {
    if (inputIndex < 1) {
      return null;
    } else {
      Input in = inputs.get(inputIndex - 1);
      return in.toString();
    }
  }

  /** Returns the File we are storing classes into. */
  public final File getOutputFile() {
    return outputFile;
  }

  protected final boolean internalOutputModifiedClass(Object cf, String name, Object mods)
      throws IOException {
    makeOutputJar();
    if (entryNames.contains(name)) {
      return false;
    } else {
      putNextEntry(new ZipEntry(name));
      BufferedOutputStream s = new BufferedOutputStream(outputJar);
      writeClassTo(cf, mods, s);
      s.flush();
      outputJar.closeEntry();
      return true;
    }
  }

  /** Set the JAR Comment for the output JAR. */
  public final void setJARComment(String comment) throws IOException, IllegalStateException {
    makeOutputJar();
    outputJar.setComment(comment);
  }

  final void makeOutputJar() throws IOException, IllegalStateException {
    if (outputJar == null) {
      if (outputFile == null) {
        throw new IllegalStateException("Output file was not set");
      }

      final FileOutputStream out = new FileOutputStream(outputFile);
      outputJar = new JarOutputStream(out);
    }
  }

  /** Skip the last class returned in every future traversal of the class list. */
  public final void setIgnore() throws IllegalArgumentException {
    if (inputIndex == 0) {
      throw new IllegalArgumentException("Must get a class before ignoring it");
    }

    ignoringInputs.set(inputIndex - 1);
  }

  private static byte[] cachedBuf;

  private static synchronized byte[] makeBuf() {
    if (cachedBuf != null) {
      byte[] r = cachedBuf;
      cachedBuf = null;
      return r;
    } else {
      return new byte[60000];
    }
  }

  private static synchronized void releaseBuf(byte[] buf) {
    cachedBuf = buf;
  }

  public static void copyStream(InputStream in, OutputStream out)
      throws IllegalArgumentException, IOException {
    if (in == null) {
      throw new IllegalArgumentException("in == null");
    }
    byte[] buf = makeBuf();
    try {
      while (true) {
        int read = in.read(buf);
        if (read < 0) {
          return;
        }
        out.write(buf, 0, read);
      }
    } finally {
      releaseBuf(buf);
    }
  }

  /**
   * Add a raw ZipEntry to the output JAR. Call endOutputJarEntry() when you're done.
   *
   * @return the OutputStream to be used to write the entry contents
   */
  public final OutputStream addOutputJarEntry(ZipEntry ze)
      throws IOException, IllegalStateException {
    if (outputJar == null) {
      throw new IllegalStateException("output jar is null");
    }
    putNextEntry(ze);
    return outputJar;
  }

  /** Complete and flush the entry initiated by addOutputJarEntry. */
  public final void endOutputJarEntry() throws IOException, IllegalStateException {
    if (outputJar == null) {
      throw new IllegalStateException("output jar is null");
    }
    outputJar.closeEntry();
  }

  /**
   * Call this to copy any unmodified classes to the output. This is called automatically by
   * close(); you should only call this if you want to write an entry to the JAR file *after* the
   * unmodified classes. This will only ever be called once per output JAR.
   */
  public final void writeUnmodifiedClasses() throws IOException, IllegalStateException {
    passUnmodifiedClasses = false;
    makeOutputJar();
    for (Input in : inputs) {
      if (!in.isClass()) {
        if (in instanceof JarInput) {
          JarInput jin = (JarInput) in;
          ZipEntry entry = jin.getEntry();
          try (final InputStream s = jin.open()) {
            ZipEntry newEntry = new ZipEntry(entry.getName());
            newEntry.setComment(entry.getComment());
            newEntry.setExtra(entry.getExtra());
            newEntry.setTime(entry.getTime());
            putNextEntry(newEntry);
            copyStream(s, outputJar);
            outputJar.closeEntry();
          }
        } else {
          throw new Error("Unknown non-class input: " + in);
        }
      } else {
        String name = in.getClassName();
        if (name == null) {
          try (final BufferedInputStream s = new BufferedInputStream(in.open(), 65536)) {
            Object cl = makeClassFromStream(in.getInputName(), s);
            String entryName = toEntryName(getClassName(cl));
            if (!entryNames.contains(entryName)) {
              putNextEntry(new ZipEntry(entryName));
              BufferedOutputStream clOut = new BufferedOutputStream(outputJar);
              writeClassTo(cl, null, clOut);
              clOut.flush();
              outputJar.closeEntry();
            }
          }
        } else {
          String entryName = toEntryName(name);
          if (!entryNames.contains(entryName)) {
            try (final BufferedInputStream s = new BufferedInputStream(in.open())) {
              putNextEntry(new ZipEntry(entryName));
              BufferedOutputStream clOut = new BufferedOutputStream(outputJar);
              copyStream(s, clOut);
              clOut.flush();
              outputJar.closeEntry();
            }
          }
        }
      }
    }
  }

  /** Call this when you're done modifying classes. */
  public final void close() throws IOException, IllegalStateException {
    if (passUnmodifiedClasses) {
      writeUnmodifiedClasses();
    }

    if (outputJar != null) {
      outputJar.close();
    }

    if (cachedJar != null) {
      cachedJar.close();
    }
  }

  private void putNextEntry(ZipEntry newEntry) throws IOException, IllegalStateException {
    if (outputJar == null) {
      throw new IllegalStateException();
    }
    outputJar.putNextEntry(newEntry);
    entryNames.add(newEntry.getName());
    if (manifestBuilder != null) {
      manifestBuilder.addEntry(newEntry);
    }
  }
}