ClassEnhancer.java [plain text]
package com.sleepycat.persist.model;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.StringTokenizer;
import com.sleepycat.asm.ClassReader;
import com.sleepycat.asm.ClassVisitor;
import com.sleepycat.asm.ClassWriter;
public class ClassEnhancer implements ClassFileTransformer {
private static final String AGENT_PREFIX = "enhance:";
private Set<String> packagePrefixes;
private boolean verbose;
public static void main(String[] args) throws Exception {
try {
boolean verbose = false;
List<File> fileList = new ArrayList<File>();
for (int i = 0; i < args.length; i += 1) {
String arg = args[i];
if (arg.startsWith("-")) {
if ("-v".equals(args[i])) {
verbose = true;
} else {
throw new IllegalArgumentException
("Unknown arg: " + arg);
}
} else {
fileList.add(new File(arg));
}
}
ClassEnhancer enhancer = new ClassEnhancer();
enhancer.setVerbose(verbose);
int nFiles = 0;
for (File file : fileList) {
nFiles += enhancer.enhanceFile(file);
}
if (nFiles > 0) {
System.out.println("Enhanced: " + nFiles + " files");
}
} catch (Exception e) {
e.printStackTrace();
throw e;
}
}
public static void premain(String args, Instrumentation inst) {
if (!args.startsWith(AGENT_PREFIX)) {
throw new IllegalArgumentException
("Unknown javaagent args: " + args +
" Args must start with: \"" + AGENT_PREFIX + '"');
}
args = args.substring(AGENT_PREFIX.length());
Set<String> packageNames = null;
boolean verbose = false;
if (args.length() > 0) {
packageNames = new HashSet<String>();
StringTokenizer tokens = new StringTokenizer(args, ",");
while (tokens.hasMoreTokens()) {
String token = tokens.nextToken();
if (token.startsWith("-")) {
if (token.equals("-v")) {
verbose = true;
} else {
throw new IllegalArgumentException
("Unknown javaagent arg: " + token);
}
} else {
packageNames.add(token);
}
}
}
ClassEnhancer enhancer = new ClassEnhancer(packageNames);
enhancer.setVerbose(verbose);
inst.addTransformer(enhancer);
}
public ClassEnhancer() {
}
public void setVerbose(boolean verbose) {
this.verbose = verbose;
}
public boolean getVerbose() {
return verbose;
}
public ClassEnhancer(Set<String> packageNames) {
if (packageNames != null) {
packagePrefixes = new HashSet<String>();
for (String name : packageNames) {
packagePrefixes.add(name + '.');
}
}
}
public byte[] transform(ClassLoader loader,
String className,
Class<?> classBeingRedefined,
ProtectionDomain protectionDomain,
byte[] classfileBuffer) {
className = className.replace('/', '.');
byte[] bytes = enhance(className, classfileBuffer);
if (verbose && bytes != null) {
System.out.println("Enhanced: " + className);
}
return bytes;
}
public byte[] enhance(String className, byte[] classBytes) {
if (className != null && packagePrefixes != null) {
for (String prefix : packagePrefixes) {
if (className.startsWith(prefix)) {
return enhanceBytes(classBytes);
}
}
return null;
} else {
return enhanceBytes(classBytes);
}
}
int enhanceFile(File file)
throws IOException {
int nFiles = 0;
if (file.isDirectory()) {
String[] names = file.list();
if (names != null) {
for (int i = 0; i < names.length; i += 1) {
nFiles += enhanceFile(new File(file, names[i]));
}
}
} else if (file.getName().endsWith(".class")) {
byte[] newBytes = enhanceBytes(readFile(file));
if (newBytes != null) {
long modified = file.lastModified();
writeFile(file, newBytes);
file.setLastModified(modified);
nFiles += 1;
if (verbose) {
System.out.println("Enhanced: " + file);
}
}
}
return nFiles;
}
private byte[] readFile(File file)
throws IOException {
byte[] bytes = new byte[(int) file.length()];
FileInputStream in = new FileInputStream(file);
try {
in.read(bytes);
} finally {
in.close();
}
return bytes;
}
private void writeFile(File file, byte[] bytes)
throws IOException {
FileOutputStream out = new FileOutputStream(file);
try {
out.write(bytes);
} finally {
out.close();
}
}
private byte[] enhanceBytes(byte[] bytes) {
ClassWriter writer = new ClassWriter(true);
ClassVisitor visitor = writer;
visitor = new BytecodeEnhancer(visitor);
ClassReader reader = new ClassReader(bytes);
try {
reader.accept(visitor, false);
return writer.toByteArray();
} catch (BytecodeEnhancer.NotPersistentException e) {
return null;
}
}
}