diff --git a/agent/src/main/java/com/alibaba/testable/agent/handler/SourceClassHandler.java b/agent/src/main/java/com/alibaba/testable/agent/handler/SourceClassHandler.java index 0514a6e..6ea2f43 100644 --- a/agent/src/main/java/com/alibaba/testable/agent/handler/SourceClassHandler.java +++ b/agent/src/main/java/com/alibaba/testable/agent/handler/SourceClassHandler.java @@ -2,6 +2,7 @@ package com.alibaba.testable.agent.handler; import com.alibaba.testable.agent.constant.ConstPool; import com.alibaba.testable.agent.model.MethodInfo; +import com.alibaba.testable.agent.util.BytecodeUtil; import com.alibaba.testable.agent.util.ClassUtil; import com.alibaba.testable.agent.util.CollectionUtil; import com.alibaba.testable.agent.util.StringUtil; @@ -9,6 +10,7 @@ import org.objectweb.asm.Opcodes; import org.objectweb.asm.tree.*; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; import java.util.Set; @@ -18,6 +20,10 @@ import java.util.Set; public class SourceClassHandler extends BaseClassHandler { private final List injectMethods; + private final Set invokeOps = new HashSet() {{ + add(Opcodes.INVOKESPECIAL); + add(Opcodes.INVOKEVIRTUAL); + }}; public SourceClassHandler(List injectMethods) { this.injectMethods = injectMethods; @@ -31,11 +37,19 @@ public class SourceClassHandler extends BaseClassHandler { protected void transform(ClassNode cn) { List methods = new ArrayList(); for (MethodNode m : cn.methods) { + // record all member method names if (!ConstPool.CONSTRUCTOR.equals(m.name)) { methods.add(new MethodInfo(cn.name, m.name, m.desc)); } } + // member methods which has injection stub Set memberInjectMethods = CollectionUtil.getCrossSet(methods, injectMethods); + for (MethodInfo mi : injectMethods) { + if (!mi.getClazz().equals(cn.name)) { + memberInjectMethods.add(mi); + } + } + // new operations which has injection stub Set newOperatorInjectMethods = CollectionUtil.getMinusSet(injectMethods, memberInjectMethods); for (MethodNode m : cn.methods) { transformMethod(cn, m, memberInjectMethods, newOperatorInjectMethods); @@ -47,10 +61,10 @@ public class SourceClassHandler extends BaseClassHandler { AbstractInsnNode[] instructions = mn.instructions.toArray(); int i = 0; do { - if (instructions[i].getOpcode() == Opcodes.INVOKESPECIAL) { + if (invokeOps.contains(instructions[i].getOpcode())) { MethodInsnNode node = (MethodInsnNode)instructions[i]; - if (cn.name.equals(node.owner) && memberInjectMethods.contains(new MethodInfo(cn.name, node.name, node.desc))) { - // it's a member method and an inject method for it exist + if (memberInjectMethods.contains(new MethodInfo(node.owner, node.name, node.desc))) { + // it's a member method of current class and an inject method for it exist int rangeStart = getMemberMethodStart(instructions, i); if (rangeStart >= 0) { instructions = replaceMemberCallOps(cn, mn, instructions, rangeStart, i); @@ -96,8 +110,20 @@ public class SourceClassHandler extends BaseClassHandler { } private int getMemberMethodStart(AbstractInsnNode[] instructions, int rangeEnd) { + int stackLevel = ClassUtil.getParameterTypes(((MethodInsnNode)instructions[rangeEnd]).desc).size(); for (int i = rangeEnd - 1; i >= 0; i--) { - if (instructions[i].getOpcode() == Opcodes.ALOAD && ((VarInsnNode)instructions[i]).var == 0) { + switch (instructions[i].getOpcode()) { + case Opcodes.INVOKEINTERFACE: + case Opcodes.INVOKEVIRTUAL: + case Opcodes.INVOKESPECIAL: + case Opcodes.INVOKEDYNAMIC: + case Opcodes.INVOKESTATIC: + stackLevel += ClassUtil.getParameterTypes(((MethodInsnNode)instructions[i]).desc).size(); + break; + default: + stackLevel -= BytecodeUtil.stackEffect(instructions[i].getOpcode()); + } + if (stackLevel < 0) { return i; } } diff --git a/agent/src/main/java/com/alibaba/testable/agent/util/BytecodeUtil.java b/agent/src/main/java/com/alibaba/testable/agent/util/BytecodeUtil.java new file mode 100644 index 0000000..550e810 --- /dev/null +++ b/agent/src/main/java/com/alibaba/testable/agent/util/BytecodeUtil.java @@ -0,0 +1,180 @@ +package com.alibaba.testable.agent.util; + +import java.util.HashMap; +import java.util.Map; + +import static org.objectweb.asm.Opcodes.*; + +/** + * @author flin + */ +public class BytecodeUtil { + + /** + * refer to https://en.wikipedia.org/wiki/Java_bytecode_instruction_listings + */ + private static Map bytecodeStackEffect = new HashMap() {{ + put(NOP, 0); + put(ACONST_NULL, 1); + put(ICONST_M1, 1); + put(ICONST_0, 1); + put(ICONST_1, 1); + put(ICONST_2, 1); + put(ICONST_3, 1); + put(ICONST_4, 1); + put(ICONST_5, 1); + put(LCONST_0, 1); + put(LCONST_1, 1); + put(FCONST_0, 1); + put(FCONST_1, 1); + put(FCONST_2, 1); + put(DCONST_0, 1); + put(DCONST_1, 1); + put(BIPUSH, 1); + put(SIPUSH, 1); + put(LDC, 1); + put(ILOAD, 1); + put(LLOAD, 1); + put(FLOAD, 1); + put(DLOAD, 1); + put(ALOAD, 1); + put(IALOAD, 1); + put(LALOAD, 1); + put(FALOAD, 1); + put(DALOAD, 1); + put(AALOAD, 1); + put(BALOAD, 1); + put(CALOAD, 1); + put(SALOAD, 1); + put(ISTORE, -1); + put(LSTORE, -1); + put(FSTORE, -1); + put(DSTORE, -1); + put(ASTORE, -1); + put(IASTORE, -1); + put(LASTORE, -1); + put(FASTORE, -1); + put(DASTORE, -1); + put(AASTORE, -1); + put(BASTORE, -1); + put(CASTORE, -1); + put(SASTORE, -1); + put(POP, -1); + put(POP2, -2); + put(DUP, 1); + put(DUP_X1, 1); + put(DUP_X2, 1); + put(DUP2, 2); + put(DUP2_X1, 2); + put(DUP2_X2, 2); + put(SWAP, 0); + put(IADD, -1); + put(LADD, -1); + put(FADD, -1); + put(DADD, -1); + put(ISUB, -1); + put(LSUB, -1); + put(FSUB, -1); + put(DSUB, -1); + put(IMUL, -1); + put(LMUL, -1); + put(FMUL, -1); + put(DMUL, -1); + put(IDIV, -1); + put(LDIV, -1); + put(FDIV, -1); + put(DDIV, -1); + put(IREM, -1); + put(LREM, -1); + put(FREM, -1); + put(DREM, -1); + put(INEG, 0); + put(LNEG, 0); + put(FNEG, 0); + put(DNEG, 0); + put(ISHL, -1); + put(LSHL, -1); + put(ISHR, -1); + put(LSHR, -1); + put(IUSHR, -1); + put(LUSHR, -1); + put(IAND, -1); + put(LAND, -1); + put(IOR, -1); + put(LOR, -1); + put(IXOR, -1); + put(LXOR, -1); + put(IINC, 0); + put(I2L, 0); + put(I2F, 0); + put(I2D, 0); + put(L2I, 0); + put(L2F, 0); + put(L2D, 0); + put(F2I, 0); + put(F2L, 0); + put(F2D, 0); + put(D2I, 0); + put(D2L, 0); + put(D2F, 0); + put(I2B, 0); + put(I2C, 0); + put(I2S, 0); + put(LCMP, -1); + put(FCMPL, -1); + put(FCMPG, -1); + put(DCMPL, -1); + put(DCMPG, -1); + put(IFEQ, -1); + put(IFNE, -1); + put(IFLT, -1); + put(IFGE, -1); + put(IFGT, -1); + put(IFLE, -1); + put(IF_ICMPEQ, -2); + put(IF_ICMPNE, -2); + put(IF_ICMPLT, -2); + put(IF_ICMPGE, -2); + put(IF_ICMPGT, -2); + put(IF_ICMPLE, -2); + put(IF_ACMPEQ, -2); + put(IF_ACMPNE, -2); + put(GOTO, 0); + put(JSR, 1); + put(RET, 0); + put(TABLESWITCH, -1); + put(LOOKUPSWITCH, -1); + put(IRETURN, 0); + put(LRETURN, 0); + put(FRETURN, 0); + put(DRETURN, 0); + put(ARETURN, 0); + put(RETURN, 1); + put(GETSTATIC, 1); + put(PUTSTATIC, -1); + put(GETFIELD, 0); + put(PUTFIELD, -2); + put(INVOKEVIRTUAL, 0); // variable + put(INVOKESPECIAL, 0); // variable + put(INVOKESTATIC, 0); // variable + put(INVOKEINTERFACE, 0); // variable + put(INVOKEDYNAMIC, 0); // variable + put(NEW, 1); + put(NEWARRAY, 0); + put(ANEWARRAY, 0); + put(ARRAYLENGTH, 0); + put(ATHROW, 1); + put(CHECKCAST, 0); + put(INSTANCEOF, 0); + put(MONITORENTER, -1); + put(MONITOREXIT, -1); + put(MULTIANEWARRAY, 0); // variable + put(IFNULL, -1); + put(IFNONNULL, -1); + }}; + + public static int stackEffect(int bytecode) { + return bytecodeStackEffect.get(bytecode); + } + +} diff --git a/agent/src/test/java/com/alibaba/testable/agent/util/ClassUtilTest.java b/agent/src/test/java/com/alibaba/testable/agent/util/ClassUtilTest.java index 0ba877e..5949119 100644 --- a/agent/src/test/java/com/alibaba/testable/agent/util/ClassUtilTest.java +++ b/agent/src/test/java/com/alibaba/testable/agent/util/ClassUtilTest.java @@ -15,6 +15,7 @@ class ClassUtilTest { @Test void should_able_to_get_parameter_count() { + assertEquals(0, ClassUtil.getParameterTypes("()V").size()); assertEquals(1, ClassUtil.getParameterTypes("(Ljava/lang/String;)V").size()); assertEquals(6, ClassUtil.getParameterTypes("(Ljava/lang/String;IDLjava/lang/String;ZLjava/net/URL;)V").size()); assertEquals(10, ClassUtil.getParameterTypes("(ZLjava/lang/String;IJFDCSBZ)V").size());