Write ReturnInsnProcessor tests

newrelic · Jan 28, 2025 · f2e4e52 · f2e4e52
1 parent f0ef5b8
commit f2e4e52
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 46 deletions.
diff --git a/newrelic-weaver/src/main/java/com/newrelic/weave/utils/ReturnInsnProcessor.java b/newrelic-weaver/src/main/java/com/newrelic/weave/utils/ReturnInsnProcessor.java
@@ -74,7 +74,7 @@ private static InsnList insnsBeforeReturn(int opcode, int stackSize, int varInde
     }
 
     /*
-     * Compute the size of the operand stack at each return instruction.
+     * Compute the size of the operand stack at each return instruction exceeding the EXPECTED_RETURN_STACK_SIZE of 1.
      * Only applies to methods returning I (int type) or A (reference type).
      *
      * @param owner The owning class
@@ -88,7 +88,7 @@ private static Map<AbstractInsnNode, Integer> getReturnStacks(String owner, Meth
         Map<AbstractInsnNode, Integer> rtStacks = new HashMap<>();
         for (int j = 0; j < method.instructions.size(); ++j) {
             AbstractInsnNode insn = method.instructions.get(j);
-            if (insn.getOpcode() == Opcodes.IRETURN && insn.getOpcode() == Opcodes.ARETURN) {
+            if (insn.getOpcode() == Opcodes.IRETURN || insn.getOpcode() == Opcodes.ARETURN) {
                 Frame<BasicValue> f = frames[j];
                 if (f != null && f.getStackSize() > EXPECTED_RETURN_STACK_SIZE) {
                     rtStacks.put(insn, f.getStackSize());

diff --git a/newrelic-weaver/src/test/java/com/newrelic/weave/utils/ReturnInsnProcessorTest.java b/newrelic-weaver/src/test/java/com/newrelic/weave/utils/ReturnInsnProcessorTest.java
@@ -1,74 +1,117 @@
 package com.newrelic.weave.utils;
 
-import com.sun.org.apache.bcel.internal.generic.ALOAD;
+import net.bytebuddy.asm.Advice;
 import org.junit.Test;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
 import org.objectweb.asm.tree.*;
 
-import static org.junit.Assert.*;
-import static org.objectweb.asm.Opcodes.ACC_PUBLIC;
-import static org.objectweb.asm.Type.DOUBLE_TYPE;
-import static org.objectweb.asm.Type.INT_TYPE;
-
-public class ReturnInsnProcessorTest {
+import java.util.Iterator;
 
+import static org.junit.Assert.*;
+import static org.objectweb.asm.Opcodes.*;
 
 
+public class ReturnInsnProcessorTest {
     //the nodes in here need to be constructed by hand since Java compilers don't have this issue
-
     @Test
-    public void singleReturnStackClears() {
-       MethodNode mn = buildMethodNode(INT_TYPE);
-       buildOneUntidyReturn(mn);
-       WeaveUtils.printAllInstructions(mn);
-
-       ReturnInsnProcessor.clearReturnStacks("MyClazz", mn);
-       WeaveUtils.printAllInstructions(mn);
+    public void doesNotAlterAlreadyClearStack(){
+        int[] originalSeq = {ICONST_1, DUP, POP, IRETURN};
+        MethodNode mn = basicMethodFromSequence(Type.INT, originalSeq, 2);
+        ReturnInsnProcessor.clearReturnStacks("MyClazz", mn);
+        assertInsnSequence(originalSeq, mn.instructions);
+
+        int[] originalSeqRefType = {ACONST_NULL, ARETURN};
+        MethodNode mn2 = basicMethodFromSequence(Type.OBJECT, originalSeqRefType, 1);
+        ReturnInsnProcessor.clearReturnStacks("MyClazz", mn2);
+        assertInsnSequence(originalSeqRefType, mn2.instructions);
     }
 
     @Test
-    public void stackAlreadyClearDoesNothing() {
+    public void addsPopsForIntReturnType(){
+        int[] originalInsnSequence = {ICONST_2, ICONST_1, DUP, DUP, ICONST_3, IRETURN};
+        MethodNode mn = basicMethodFromSequence(Type.INT, originalInsnSequence, 5);
 
+        ReturnInsnProcessor.clearReturnStacks("MyClazz", mn);
+        int[] expectedInsnSequence = {ICONST_2, ICONST_1, DUP, DUP, ICONST_3, ISTORE, POP, POP, POP, POP, ILOAD, IRETURN};
+        assertInsnSequence(expectedInsnSequence, mn.instructions);
     }
 
+    @Test
+    public void addsPopsForReferenceReturnType(){
+        int[] originalInsnSequence = {ACONST_NULL, DUP, DUP, DUP, ARETURN};
+        MethodNode mn = basicMethodFromSequence(Type.OBJECT, originalInsnSequence, 4);
 
-    public MethodNode buildMethodNode(Type type) {
-        String desc = "";
-        desc = "()" + type.getDescriptor();
-        MethodNode mn = new MethodNode(Opcodes.ASM9, ACC_PUBLIC, "myMethod", desc, null, null);
-        return mn;
+        ReturnInsnProcessor.clearReturnStacks("MyClazz", mn);
+        int[] expectedInsnSequence = {ACONST_NULL, DUP, DUP, DUP, ASTORE, POP, POP, POP, ALOAD, ARETURN};
+        assertInsnSequence(expectedInsnSequence, mn.instructions);
     }
 
-    private void buildOneUntidyReturn(MethodNode mn) {
-        //Java compilers will resist untidy return stacks.
-        //But we can generate unnecessary bytecode using ASM if we want to!
+    @Test
+    public void handlesMultipleReturns(){
+        //This Insn List is more complicated, so we have to build it directly here
         InsnList insns = new InsnList();
-        //setup insns
-        insns.add(new LabelNode());
-        insns.add(new InsnNode(Opcodes.DCONST_1));
-        insns.add(new VarInsnNode(Opcodes.DSTORE, 1));
-        insns.add(new InsnNode(Opcodes.DCONST_0));
-        insns.add(new InsnNode(Opcodes.DUP2));
-        insns.add(new InsnNode(Opcodes.DRETURN));
-
+        insns.add(new InsnNode(ICONST_1));
+        insns.add(new InsnNode(DUP));
+        LabelNode label = new LabelNode();
+        insns.add(new JumpInsnNode(IFEQ, label));
+        insns.add(new InsnNode(IRETURN));
+        insns.add(label);
+        insns.add(new InsnNode(DUP));
+        insns.add(new InsnNode(IRETURN));
+        MethodNode mn = buildMethodNode(Type.INT);
         mn.instructions.add(insns);
-        mn.maxLocals += 4;
-        mn.maxStack += 4;
-    }
-
-    private void buildOneTidyReturn(MethodNode mn) {
-        InsnList insns = new InsnList();
-        //setup insns
-        insns.add(new InsnNode(Opcodes.ICONST_3));
-        insns.add(new VarInsnNode(Opcodes.ISTORE, 1));
-        insns.add(new InsnNode(Opcodes.ICONST_4));
-        insns.add(new InsnNode(Opcodes.IRETURN));
+        mn.maxLocals = 1;
+        mn.maxStack = 2;
 
-        mn.instructions.add(insns);
+        ReturnInsnProcessor.clearReturnStacks("MyClazz", mn);
+        int[] expectedInsns = {ICONST_1, DUP, IFEQ, IRETURN, -1, DUP, ISTORE, POP, ILOAD, IRETURN};
+        assertInsnSequence(expectedInsns, mn.instructions);
     }
 
+    @Test
+    public void failedAnalyzerDoesNothing() {
+        //these instructions are bad and will fail the analyzer - we shouldn't touch them.
+        int[] originalInsns = {ICONST_1, POP, POP, IRETURN};
+        MethodNode mn = basicMethodFromSequence(Type.INT, originalInsns, 1);
+        ReturnInsnProcessor.clearReturnStacks("MyClazz", mn);
+        assertInsnSequence(originalInsns, mn.instructions);
+    }
 
+    private void assertInsnSequence( int[] expectedSequence, InsnList insns) {
+        Iterator<AbstractInsnNode> it = insns.iterator();
+        int i = 0;
+        while (it.hasNext()) {
+            AbstractInsnNode insn = it.next();
+            int expectedOpcode = expectedSequence[i];
+            assertEquals(expectedOpcode, insn.getOpcode());
+            i++;
+        }
+    }
 
+    private MethodNode basicMethodFromSequence(int type, int[] insnSequence, int maxStack) {
+        MethodNode mn = buildMethodNode(type);
+        for(int opcode: insnSequence){
+            mn.instructions.add(new InsnNode(opcode));
+        }
+        mn.maxStack = maxStack;
+        mn.maxLocals = 1;
+        return mn;
+    }
 
+    public MethodNode buildMethodNode(int type) {
+        String desc = "";
+        switch (type) {
+            case Type.INT:
+                desc = "()I";
+                break;
+            case Type.OBJECT:
+                desc = "()L";
+                break;
+            default:
+                fail();
+        }
+        MethodNode mn = new MethodNode(Opcodes.ASM9, ACC_PUBLIC, "myMethod", desc, null, null);
+        return mn;
+    }
 }