How to record all parameters of any invoked java method dynamically using java agent and ASM?

Question

What I want to do is to record those parameters and do some comparison with previously input parameters. I need to record parameters for each method invoked, so the length of the parameter list is uncertain.

I think I can parse the method descriptor to know how many parameters this method has. But the problem is, How to record arbitrary number of values from the operand stack?

Now I can only write some sample code. In my MethodVisitorAdapter Class:

public void visitMethodInsn(int opc, String owner, String name, String desc, boolean isInterface){
    int n = getParameterCount(desc);
    // How to duplicate arbitrary number of values at operand stack?
    // ...
    mv.visitMethodInsn(INVOKESTATIC, MyClass, "recordParameters",
                    /* should be what kind of descriptor (may have arbitrary number of parameters)? */, 
                    false);
    mv.visitMethodInsn(opc, owner, name, desc, isInterface);
}

Answer 1

You can make recordParameters a variable arity method (varargs). It will accept all parameters as a single Object[] argument.

So, you'll need to create an array and fill it with (possibly boxed) arguments. The following builder-style class will help to collect the arguments from the stack. It contains push methods for all possible types, and also handles autoboxing.

After recording the parameters, the helper can put arguments back onto the stack by calling the corresponding pop methods in the reverse order.

public class ArgCollector {
    private final Object[] args;
    private int index;

    public ArgCollector(int length) {
        this.args = new Object[length];
        this.index = length;
    }

    public ArgCollector push(Object o) {
        args[--index] = o;
        return this;
    }

    public Object pop() {
        return args[index++];
    }

    public static ArgCollector push(boolean a, ArgCollector c) { return c.push(a); }
    public static ArgCollector push(byte    a, ArgCollector c) { return c.push(a); }
    public static ArgCollector push(char    a, ArgCollector c) { return c.push(a); }
    public static ArgCollector push(short   a, ArgCollector c) { return c.push(a); }
    public static ArgCollector push(int     a, ArgCollector c) { return c.push(a); }
    public static ArgCollector push(long    a, ArgCollector c) { return c.push(a); }
    public static ArgCollector push(float   a, ArgCollector c) { return c.push(a); }
    public static ArgCollector push(double  a, ArgCollector c) { return c.push(a); }
    public static ArgCollector push(Object  a, ArgCollector c) { return c.push(a); }

    public boolean popZ() { return (boolean) pop(); }
    public byte    popB() { return (byte)    pop(); }
    public char    popC() { return (char)    pop(); }
    public short   popS() { return (short)   pop(); }
    public int     popI() { return (int)     pop(); }
    public long    popJ() { return (long)    pop(); }
    public float   popF() { return (float)   pop(); }
    public double  popD() { return (double)  pop(); }

    public Object[] toArray() {
        return args;
    }
}

Now, the task is to generate the bytecode equivalent for the following Java code:

    ArgCollector collector = new ArgCollector(N);
    recordParameters(
            ArgCollector.push(arg1,
                ArgCollector.push(arg2,
                    ArgCollector.push(argN, collector)))
            .toArray()
    );

    originalMethod(
            collector.popI(),
            collector.popJ(),
            (String) collector.pop()
    );

Here is how to do this with ASM:

    Type[] args = Type.getArgumentTypes(desc);
    String collector = Type.getInternalName(ArgCollector.class);

    // new ArgCollector(argCount)
    mv.visitTypeInsn(NEW, collector);
    mv.visitInsn(DUP);
    mv.visitIntInsn(SIPUSH, args.length);
    mv.visitMethodInsn(INVOKESPECIAL, collector, "<init>", "(I)V", false);

    // For each argument call ArgCollector.push(arg, collector)
    for (int i = args.length; --i >= 0; ) {
        Type arg = args[i];
        String argDesc = arg.getDescriptor().length() == 1 ? arg.getDescriptor() : "Ljava/lang/Object;";
        mv.visitMethodInsn(INVOKESTATIC, collector, "push",
                "(" + argDesc + "L" + collector + ";)L" + collector + ";", false);
    }

    // Call recordParameters(collector.toArray())
    mv.visitInsn(DUP);
    mv.visitMethodInsn(INVOKEVIRTUAL, collector, "toArray", "()[Ljava/lang/Object;", false);
    mv.visitMethodInsn(INVOKESTATIC, MyClass, "recordParameters", "([Ljava/lang/Object;)V", false);

    // Push original arguments back on stack
    for (Type arg : args) {
        String argDesc = arg.getDescriptor().length() == 1 ? arg.getDescriptor() : "Ljava/lang/Object;";

        mv.visitInsn(DUP);
        if (argDesc.length() == 1) {
            mv.visitMethodInsn(INVOKEVIRTUAL, collector, "pop" + argDesc, "()" + argDesc, false);
        } else {
            mv.visitMethodInsn(INVOKEVIRTUAL, collector, "pop", "()Ljava/lang/Object;", false);
            if (!arg.getDescriptor().equals("Ljava/lang/Object;")) {
                // Need to cast object arguments to the original type
                mv.visitTypeInsn(CHECKCAST, arg.getDescriptor());
            }
        }

        // Swap the last argument with ArgCollector, so that ArgCollector is on top again
        if (arg.getSize() == 1) {
            mv.visitInsn(SWAP);
        } else {
            mv.visitInsn(DUP2_X1);
            mv.visitInsn(POP2);
        }
    }

    // Pop off the remaining ArgCollector, and call the original method
    mv.visitInsn(POP);
    mv.visitMethodInsn(opc, owner, name, desc, isInterface);

Answer 2

To duplicate an arbitrary sequence of values on the operand stack, there is no way around storing them into new local variables temporarily. Then, push all these values to the operand stack, invoke your reporting method, push them again, and invoke the original method.

A simple approach, avoiding the expenses of LocalVariablesSorter, is to check each instruction that uses a local variable and remember the first free index. This requires that the source code has already valid stack map frames, to handle backward branches correctly within a single visiting pass. For code targeting Java 7 or newer, this is mandatory anyway.

Since the code to report the invocation is free of branches and doesn’t need the temporary variables afterwards, we don’t even need an expensive recalculation of stack map frames, as at branch merge points, only the original variables are needed. Only the max stack and locals need recalculation.

To invoke a report method, we can use a holder object, like in apangin’s answer, but alternatively, we can use the java.lang.invoke package, which allows us to generate varargs collectors with precise method descriptors on-the-fly.

The following code will acquire a MethodHandle to PrintStream.printf(String,Object...) via a single ldc instruction, followed by binding System.out as first argument, followed by binding a constant String suitable for the current number of arguments (and the receiver for non-static methods), then adapt the handle calling .asVarargsCollector(Object[].class)​.asType(targetType). The targetType is the method type descriptor, with an additional first parameter type for non-static method invocations. This MethodType is also loaded with a single ldc instruction. Then, the handle can be used by calling invokeExact with the same arguments as the actual invocation on the stack.

Only for constructor invocations, the receiver object is omitted, as we are not allowed to use an object before its initialization.

{𝓍store           n + i } for each argumentᵢ
 ldc              MethodHandle invokeVirtual java/io/PrintStream.printf(Ljava/lang/String;[Ljava/lang/Object;)Ljava/io/PrintStream;
 getstatic        java/lang/System.out Ljava/io/PrintStream;
 invokevirtual    java/lang/invoke/MethodHandle.bindTo(Ljava/lang/Object;)Ljava/lang/invoke/MethodHandle;
 ldc              String containing <method name> and as many %s place holders as needed
 invokevirtual    java/lang/invoke/MethodHandle.bindTo(Ljava/lang/Object;)Ljava/lang/invoke/MethodHandle;
 ldc              class [Ljava/lang/Object;
 invokevirtual    java/lang/invoke/MethodHandle.asVarargsCollector(Ljava/lang/Class;)Ljava/lang/invoke/MethodHandle;
 ldc              MethodType («actual argument types»)Ljava/io/PrintStream;
 invokevirtual    java/lang/invoke/MethodHandle.asType(Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/MethodHandle;
{𝓍load            n + i } for each argumentᵢ
 invokevirtual    java/lang/invoke/MethodHandle.invokeExact(«actual argument types»)Ljava/io/PrintStream;
 pop              // remove the PrintStream return by printf
{𝓍load            n + i } for each argumentᵢ
 invoke...        original method

The transforming code and example

public class LogMethodCalls {
    public static void main(String[] args) throws IOException, IllegalAccessException {
        MethodHandles.lookup().defineClass(instrument(LogMethodCalls.class
            .getResourceAsStream("LogMethodCalls$ToInstrument.class")));
        runInstrumented();
    }
    private static void runInstrumented() {
        new ToInstrument().run();
    }
    static class ToInstrument implements Runnable {
        @Override
        public void run() {
            double min = Integer.MAX_VALUE, max = Integer.MIN_VALUE;
            for(double i: List.of(4, 2, 9, 6)) {
                min = Math.min(min, i);
                max = Math.max(max, i);
            }
            System.out.printf("min %.0f, max %.0f%n", min, max);
        }
    }
    static byte[] instrument(InputStream is) throws IOException {
        ClassReader cr = new ClassReader(is);
        ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS);

        cr.accept(new ClassVisitor(Opcodes.ASM7, cw) {
            @Override
            public MethodVisitor visitMethod(int access, String name,
                    String descriptor, String signature, String[] exceptions) {
                return new LogInjector(
                    super.visitMethod(access, name, descriptor, signature, exceptions),
                    access, descriptor);
            }
        }, ClassReader.EXPAND_FRAMES);

        return cw.toByteArray();
    }

    static class LogInjector extends MethodVisitor {
        static final String PS_T = "java/io/PrintStream", PS_S = "L" + PS_T + ";";
        static final String PRINTF_DESC="(Ljava/lang/String;[Ljava/lang/Object;)"+PS_S;
        static final String MH_T="java/lang/invoke/MethodHandle", MH_S="L" + MH_T + ";";

        private int firstUnusedVar;
        public LogInjector(MethodVisitor mv, int acc, String desc) {
            super(Opcodes.ASM7, mv);
            int vars = Type.getArgumentsAndReturnSizes(desc) >> 2;
            if((acc & Opcodes.ACC_STATIC) != 0) vars--;
            firstUnusedVar = vars;
        }
        @Override
        public void visitFrame(int type,
            int numLocal, Object[] local, int numStack, Object[] stack) {
            super.visitFrame(type, numLocal, local, numStack, stack);
            firstUnusedVar = Math.max(firstUnusedVar, numLocal);
        }
        @Override
        public void visitVarInsn(int opcode, int var) {
            super.visitVarInsn(opcode, var);
            if(opcode == Opcodes.LSTORE || opcode == Opcodes.DSTORE) var++;
            if(var >= firstUnusedVar) firstUnusedVar = var + 1;
        }
        @Override
        public void visitMethodInsn(int opcode,
            String owner, String name, String descriptor, boolean isInterface) {
            Type[] arg = Type.getArgumentTypes(descriptor);

            int[] vars = storeArguments(arg, opcode, name, owner);

            String reportDesc = getReportDescriptor(owner, descriptor, arg, vars);

            mv.visitLdcInsn(new Handle(Opcodes.H_INVOKEVIRTUAL,
                PS_T, "printf", PRINTF_DESC, false));
            mv.visitFieldInsn(Opcodes.GETSTATIC, "java/lang/System", "out", PS_S);
            bindTo();
            mv.visitLdcInsn(messageFormat(opcode, owner, name, arg));
            bindTo();
            mv.visitLdcInsn(Type.getObjectType("[Ljava/lang/Object;"));
            mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, MH_T,
                "asVarargsCollector", "(Ljava/lang/Class;)"+MH_S, false);
            mv.visitLdcInsn(Type.getMethodType(reportDesc));
            mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, MH_T,
                "asType", "(Ljava/lang/invoke/MethodType;)"+MH_S, false);
            pushArguments(arg, vars);
            mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL,
                MH_T, "invokeExact", reportDesc, false);
            mv.visitInsn(Opcodes.POP);

            pushArguments(arg, vars);
            super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
        }

        String getReportDescriptor(
            String owner, String descriptor, Type[] arg, int[] vars) {
            StringBuilder sb = new StringBuilder(owner.length()+descriptor.length()+2);
            sb.append('(');
            if(arg.length != vars.length) {
                if(owner.charAt(0) == '[') sb.append(owner);
                else sb.append('L').append(owner).append(';');
            }
            sb.append(descriptor, 1, descriptor.lastIndexOf(')')+1);
            return sb.append(PS_S).toString();
        }

        int[] storeArguments(Type[] arg, int opcode, String name, String owner) {
            int nArg = arg.length;
            boolean withThis = opcode != Opcodes.INVOKESTATIC && !name.equals("<init>");
            if(withThis) nArg++;
            int[] vars = new int[nArg];
            int slot = firstUnusedVar;
            for(int varIx = nArg-1, argIx = arg.length-1; argIx >= 0; varIx--,argIx--) {
                Type t = arg[argIx];
                mv.visitVarInsn(t.getOpcode(Opcodes.ISTORE), vars[varIx] = slot);
                slot += t.getSize();
            }
            if(withThis)
                mv.visitVarInsn(Opcodes.ASTORE, vars[0] = slot);
            return vars;
        }
        private void bindTo() {
            mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, MH_T,
                "bindTo", "(Ljava/lang/Object;)"+MH_S, false);
        }
        private void pushArguments(Type[] arg, int[] vars) {
            int vIx = 0;
            if(arg.length != vars.length)
                mv.visitVarInsn(Opcodes.ALOAD, vars[vIx++]);
            for(Type t: arg)
                mv.visitVarInsn(t.getOpcode(Opcodes.ILOAD), vars[vIx++]);
        }
        private String messageFormat(int opcode, String owner, String name, Type[] arg){
            StringBuilder sb = new StringBuilder();
            switch(opcode) {
                case Opcodes.INVOKESPECIAL:
                    if(name.equals("<init>")) {
                        name = Type.getObjectType(owner).getClassName();
                        break;
                    }
                    // else no break
                case Opcodes.INVOKEINTERFACE: // no break
                case Opcodes.INVOKEVIRTUAL:
                    sb.append("[%s].");
                    break;
                case Opcodes.INVOKESTATIC:
                    sb.append(Type.getObjectType(owner).getClassName()).append('.');
                    break;
            }
            sb.append(name);
            if(arg.length == 0) sb.append("()%n");
            else {
                sb.append('(');
                for(int i = arg.length; i > 1; i--) sb.append("%s, ");
                sb.append("%s)%n");
            }
            return sb.toString();
        }
    }
}

The example invocation uses Java 9 and relies on a JVM with lazy loading, so it can (re)define the class before actually using it. It might get replaced with an actual Instrumentation scenario, as it is not relevant for the actual transformation logic. In my setup, the example prints

java.lang.Object()
java.lang.Integer.valueOf(4)
java.lang.Integer.valueOf(2)
java.lang.Integer.valueOf(9)
java.lang.Integer.valueOf(6)
java.util.List.of(4, 2, 9, 6)
[[4, 2, 9, 6]].iterator()
[java.util.ImmutableCollections$ListItr@5ce65a89].hasNext()
[java.util.ImmutableCollections$ListItr@5ce65a89].next()
[4].intValue()
java.lang.Math.min(2.147483647E9, 4.0)
java.lang.Math.max(-2.147483648E9, 4.0)
[java.util.ImmutableCollections$ListItr@5ce65a89].hasNext()
[java.util.ImmutableCollections$ListItr@5ce65a89].next()
[2].intValue()
java.lang.Math.min(4.0, 2.0)
java.lang.Math.max(4.0, 2.0)
[java.util.ImmutableCollections$ListItr@5ce65a89].hasNext()
[java.util.ImmutableCollections$ListItr@5ce65a89].next()
[9].intValue()
java.lang.Math.min(2.0, 9.0)
java.lang.Math.max(4.0, 9.0)
[java.util.ImmutableCollections$ListItr@5ce65a89].hasNext()
[java.util.ImmutableCollections$ListItr@5ce65a89].next()
[6].intValue()
java.lang.Math.min(2.0, 6.0)
java.lang.Math.max(9.0, 6.0)
[java.util.ImmutableCollections$ListItr@5ce65a89].hasNext()
java.lang.Double.valueOf(2.0)
java.lang.Double.valueOf(9.0)
[java.io.PrintStream@6e5e91e4].printf(min %.0f, max %.0f%n, [Ljava.lang.Object;@2cdf8d8a)
min 2, max 9

---

Note that your use case might be simpler. If your logging method is a static method, not needing a PrintStream, you don’t need binding it. When it doesn’t return a value, you also don’t need to pop it. It would be even simpler when it accepts variable arguments, including the format string or method name. Then, we can pass the string like an ordinary argument, instead of bind it, and since the method handle now is unmodified, it will already be a varargs collector when the target method is a varargs method:

static void yourLog(Object... arg) {
    String name = (String) arg[0]; // or format string
    arg = Arrays.copyOfRange(arg, 1, arg.length);
    …
}

@Override
public void visitMethodInsn(int opcode,
    String owner, String name, String descriptor, boolean isInterface) {

    Type[] arg = Type.getArgumentTypes(descriptor);
    int[] vars = storeArguments(arg, opcode, name, owner);
    String reportDesc = getReportDescriptor(owner, descriptor, arg, vars);
    mv.visitLdcInsn(new Handle(Opcodes.H_INVOKESTATIC,
        YOUR_TARGET_TYPE, "yourLog", "([Ljava/lang/Object;)V", false));
    mv.visitLdcInsn(Type.getMethodType(reportDesc));
    mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, MH_T,
        "asType", "(Ljava/lang/invoke/MethodType;)"+MH_S, false);
    mv.visitLdcInsn(messageFormat(opcode, owner, name, arg));
    pushArguments(arg, vars);
    mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, MH_T, "invokeExact", reportDesc, false);

    pushArguments(arg, vars);
    super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
}

String getReportDescriptor(String owner, String descriptor, Type[] arg, int[] vars) {
    StringBuilder sb = new StringBuilder(owner.length()+descriptor.length()+2);
    sb.append("(Ljava/lang/String;");
    if(arg.length != vars.length) {
        if(owner.charAt(0) == '[') sb.append(owner);
        else sb.append('L').append(owner).append(';');
    }
    sb.append(descriptor, 1, descriptor.lastIndexOf(')')+1);
    return sb.append('V').toString();
}