/*
* Copyright (c) 2014, Oracle and/or its affiliates. All rights reserved.
* DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
*
* This code is free software; you can redistribute it and/or modify it
* under the terms of the GNU General Public License version 2 only, as
* published by the Free Software Foundation.
*
* This code is distributed in the hope that it will be useful, but WITHOUT
* ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
* FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
* version 2 for more details (a copy is included in the LICENSE file that
* accompanied this code).
*
* You should have received a copy of the GNU General Public License version
* 2 along with this work; if not, write to the Free Software Foundation,
* Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
*
* Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
* or visit www.oracle.com if you need additional information or have any
* questions.
*/
package asmlib;
import java.io.PrintStream;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import jdk.internal.org.objectweb.asm.ClassReader;
import jdk.internal.org.objectweb.asm.ClassVisitor;
import jdk.internal.org.objectweb.asm.ClassWriter;
import jdk.internal.org.objectweb.asm.MethodVisitor;
import jdk.internal.org.objectweb.asm.Opcodes;
import java.util.function.Consumer;
import jdk.internal.org.objectweb.asm.Type;
public class Instrumentor {
public static class InstrHelper {
private final MethodVisitor mv;
private final String name;
InstrHelper(MethodVisitor mv, String name) {
this.mv = mv;
this.name = name;
}
public String getName() {
return this.name;
}
public void invokeStatic(String owner, String name, String desc, boolean itf) {
mv.visitMethodInsn(Opcodes.INVOKESTATIC, owner, name, desc, itf);
}
public void invokeSpecial(String owner, String name, String desc) {
mv.visitMethodInsn(Opcodes.INVOKESPECIAL, owner, name, desc, false);
}
public void invokeVirtual(String owner, String name, String desc) {
mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, owner, name, desc, false);
}
public void push(int val) {
if (val >= -1 && val <= 5) {
mv.visitInsn(Opcodes.ICONST_0 + val);
} else if (val >= Byte.MIN_VALUE && val <= Byte.MAX_VALUE) {
mv.visitIntInsn(Opcodes.BIPUSH, val);
} else if (val >= Short.MIN_VALUE && val <= Short.MAX_VALUE) {
mv.visitIntInsn(Opcodes.SIPUSH, val);
} else {
mv.visitLdcInsn(val);
}
}
public void push(Object val) {
mv.visitLdcInsn(val);
}
public void println(String s) {
mv.visitFieldInsn(Opcodes.GETSTATIC, Type.getInternalName(System.class), "out", Type.getDescriptor(PrintStream.class));
mv.visitLdcInsn(s);
mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, Type.getInternalName(PrintStream.class), "println", Type.getMethodDescriptor(Type.VOID_TYPE, Type.getType(String.class)), false);
}
}
public static Instrumentor instrFor(byte[] classData) {
return new Instrumentor(classData);
}
private final ClassReader cr;
private final ClassWriter output;
private ClassVisitor instrumentingVisitor = null;
private final AtomicInteger matches = new AtomicInteger(0);
private Instrumentor(byte[] classData) {
cr = new ClassReader(classData);
output = new ClassWriter(ClassWriter.COMPUTE_MAXS);
instrumentingVisitor = output;
}
public synchronized Instrumentor addMethodEntryInjection(String methodName, Consumer<InstrHelper> injector) {
instrumentingVisitor = new ClassVisitor(Opcodes.ASM5, instrumentingVisitor) {
@Override
public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
if (name.equals(methodName)) {
matches.getAndIncrement();
mv = new MethodVisitor(Opcodes.ASM5, mv) {
@Override
public void visitCode() {
injector.accept(new InstrHelper(mv, name));
}
};
}
return mv;
}
};
return this;
}
public synchronized Instrumentor addNativeMethodTrackingInjection(String prefix, Consumer<InstrHelper> injector) {
instrumentingVisitor = new ClassVisitor(Opcodes.ASM5, instrumentingVisitor) {
private final Set<Consumer<ClassVisitor>> wmGenerators = new HashSet<>();
private String className;
@Override
public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
this.className = name;
super.visit(version, access, name, signature, superName, interfaces);
}
@Override
public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
if ((access & Opcodes.ACC_NATIVE) != 0) {
matches.getAndIncrement();
String newName = prefix + name;
wmGenerators.add((v)->{
MethodVisitor mv = v.visitMethod(access & ~Opcodes.ACC_NATIVE, name, desc, signature, exceptions);
mv.visitCode();
injector.accept(new InstrHelper(mv, name));
Type[] argTypes = Type.getArgumentTypes(desc);
Type retType = Type.getReturnType(desc);
boolean isStatic = (access & Opcodes.ACC_STATIC) != 0;
if (!isStatic) {
mv.visitIntInsn(Opcodes.ALOAD, 0); // load "this"
}
// load the method parameters
if (argTypes.length > 0) {
int ptr = isStatic ? 0 : 1;
for(Type argType : argTypes) {
mv.visitIntInsn(argType.getOpcode(Opcodes.ILOAD), ptr);
ptr += argType.getSize();
}
}
mv.visitMethodInsn(isStatic ? Opcodes.INVOKESTATIC : Opcodes.INVOKESPECIAL, className, newName, desc, false);
mv.visitInsn(retType.getOpcode(Opcodes.IRETURN));
mv.visitMaxs(1, 1); // dummy call; let ClassWriter to deal with this
mv.visitEnd();
});
return super.visitMethod(access, newName, desc, signature, exceptions);
}
return super.visitMethod(access, name, desc, signature, exceptions);
}
@Override
public void visitEnd() {
wmGenerators.stream().forEach((e) -> {
e.accept(cv);
});
super.visitEnd();
}
};
return this;
}
public synchronized byte[] apply() {
cr.accept(instrumentingVisitor, ClassReader.SKIP_DEBUG + ClassReader.EXPAND_FRAMES);
return matches.get() == 0 ? null : output.toByteArray();
}
}