EvictingClassLoader.java

/*
 * The MIT License
 * Copyright © 2016 AdvisedTesting
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
package com.github.advisedtesting.classloader;

import java.io.DataInputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class EvictingClassLoader extends ClassLoader {

  //Spring's shadowing classloader had this info...
  public static final String[] DEFAULT_EXCLUDED_PACKAGES =
          new String[] {"java.", "javax.", "sun.", "oracle.", "com.sun.", "com.ibm.", "COM.ibm.",
              "org.w3c.", "org.xml.", "org.dom4j.", "org.eclipse", "org.aspectj.", "net.sf.cglib",
              "org.springframework.cglib", "org.apache.xerces.", "org.apache.commons.logging."};
  
  private final List<String> whiteList;

  private final ClassFileTransformer transformer;

  private final Map<String, String> classNameToError = new HashMap<>();
  
  public EvictingClassLoader(List<String> whiteList, ClassFileTransformer transformer, ClassLoader parent) {
    super(parent);
    this.whiteList = whiteList;
    whiteList.addAll(Arrays.asList(DEFAULT_EXCLUDED_PACKAGES));
    this.transformer = transformer;
  }

  

  
  private Class<?> getClass(String name) throws ClassNotFoundException {
    String file = name.replace('.', File.separatorChar) + ".class";
    byte[] bytes = null;
    try {
      bytes = loadClassData(file);
      try {
        transformer.transform(null, name, null, null, bytes);
      } catch (ClassFormatError error) {
        classNameToError.put(name, error.getMessage());
        throw error;
      } catch (IllegalClassFormatException icfe) {
        throw new ClassNotFoundException(name, icfe);
      }
      Class<?> loaded = super.findLoadedClass(name);
      if (loaded != null) {
        return loaded;
      } else {   
        Class<?> cl = defineClass(name, bytes, 0, bytes.length);
        resolveClass(cl);
        return cl;
      }
    } catch (IOException ioe) {
      ioe.printStackTrace();
      return null;
    }
  }

  @Override
  public Class<?> loadClass(String name) throws ClassNotFoundException {
    boolean shouldLoad = true;
    for (String prefix : whiteList) {
      shouldLoad = shouldLoad && !name.startsWith(prefix);
    }
    if (shouldLoad) {
      return getClass(name);
    }
    return super.loadClass(name);
  }

  @Override
  public Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
    boolean shouldLoad = true;
    for (String prefix : whiteList) {
      shouldLoad = shouldLoad && !name.startsWith(prefix);
    }
    if (shouldLoad) {
      return getClass(name);
    }
    return super.loadClass(name, resolve);
  }
  
  /**
   * Loads a given file (presumably .class) into a byte array. The file should be accessible as a resource, for example it could be
   * located on the classpath.
   *
   * @param name
   *          File name to load
   * @return Byte array read from the file
   * @throws IOException
   *           Is thrown when there was some problem reading the file
   */
  private byte[] loadClassData(String name) throws IOException {
    InputStream stream = getClass().getClassLoader().getResourceAsStream(name);
    int size = stream.available();
    byte[] buff = new byte[size];
    DataInputStream in = new DataInputStream(stream);
    in.readFully(buff);
    in.close();
    return buff;
  }
  
  /**
   * <p>
   * If this restrictive class loader didn't allow the class this will return the reason the class was evicted. 
   * </p>
   * <p> 
   * Useful for transforming NoClassDefFoundErrors in subsequent calls to a class that attempted to
   * link to a class that triggered a ClassFormatError in an imported class.
   * </p>
   * @param className the class we suspect was evicted.
   * @return errorMessage of eviction, or null.
   */
  public String getError(String className) {
    return classNameToError.get(className);
  }
}