Search code examples
javabytecodejava-bytecode-asmbytecode-manipulationjvm-bytecode

Injecting try/catch block for serializable check in bytecode through ASM


I am new to ASM and I want some help related to bytecode transformation.

I would like to add print function with try/catch block for every local variable in bytecode through ASM. I found that previous questions about adding try/catch block were about the entire method. I know little about the stack map frame, so any pointers would be highly appreciated. Thanks in advance.

What I expected for every object, e.g. someObject: print the serialized representation for it if this object is serializable, if not, use toString() to print:

try {
  ByteArrayOutputStream bos = new ByteArrayOutputStream();
  ObjectOutputStream oos = new ObjectOutputStream(bos);
  oos.writeObject(someObject);
  String serializedObject = Base64.getEncoder().encodeToString(bos.toByteArray());
  oos.close();
  System.out.println(serializedObject);
} catch (IOException ex) {
  System.out.println(someObject.toString());
}

Since I'm trying to do this for every object, so I override visitVarInsn() in MethodVisitor, as following:

@Override
public void visitVarInsn(int opcode, int var) {
  super.visitVarInsn(opcode, var);
  switch (opcode) {
    case Opcodes.ASTORE:
      Label tryStart = new Label ();
      Label tryEnd = new Label ();
      Label catchStart = new Label ();
      Label catchEnd = new Label ();
      mv.visitTryCatchBlock(tryStart, tryEnd, catchStart, "java/io/IOException");

      mv.visitLabel(tryStart);
      // ==> ByteArrayOutputStream bos = new ByteArrayOutputStream();
      mv.visitTypeInsn(Opcodes.NEW, "java/io/ByteArrayOutputStream");
      mv.visitInsn(Opcodes.DUP);
      mv.visitMethodInsn(INVOKESPECIAL, "java/io/ByteArrayOutputStream", "<init>", "()V", false);
      mv.visitVarInsn(Opcodes.ASTORE, var + 1);
      // ==> ObjectOutputStream oos = new ObjectOutputStream(bos);
      mv.visitTypeInsn(Opcodes.NEW, "java/io/ObjectOutputStream");
      mv.visitInsn(Opcodes.DUP);
      mv.visitVarInsn(Opcodes.ALOAD, var + 1);
      mv.visitMethodInsn(INVOKESPECIAL, "java/io/ObjectOutputStream", "<init>", "(Ljava/io/OutputStream;)V", false);
      mv.visitVarInsn(Opcodes.ASTORE, var + 2);
      // ==> oos.writeObject(someObject);
      mv.visitVarInsn(Opcodes.ALOAD, var + 2);
      mv.visitVarInsn(Opcodes.ALOAD, var);
      mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/ObjectOutputStream", "writeObject", "(Ljava/lang/Object;)V", false);
      // ==> String serializedObject = Base64.getEncoder().encodeToString(bos.toByteArray());
      mv.visitMethodInsn(INVOKESTATIC, "java/util/Base64", "getEncoder", "()Ljava/util/Base64$Encoder;", false);
      mv.visitVarInsn(Opcodes.ALOAD, var + 1);
      mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/ByteArrayOutputStream", "toByteArray", "()[B", false);
      mv.visitMethodInsn(INVOKEVIRTUAL, "java/util/Base64$Encoder", "encodeToString", "([B)Ljava/lang/String;", false);
      mv.visitVarInsn(Opcodes.ASTORE, var + 3);
      // ==> oos.close();
      mv.visitVarInsn(Opcodes.ALOAD, var + 2);
      mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/ObjectOutputStream", "close", "()V", false);
      // ==> System.out.println
      mv.visitFieldInsn(Opcodes.GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
      mv.visitVarInsn(Opcodes.ALOAD, var + 3);
      mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V", false);

      mv.visitLabel(tryEnd);
      mv.visitJumpInsn(Opcodes.GOTO, catchEnd);

      mv.visitLabel(catchStart);
      mv.visitVarInsn(ASTORE, var + 1); // store exception
      // ==> System.out.println
      mv.visitFieldInsn(Opcodes.GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
      mv.visitVarInsn(Opcodes.ALOAD, var);
      mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/Object", "toString", "()Ljava/lang/String;", false);
      mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V", false);

      mv.visitLabel(catchEnd);

      // not sure whether I should add this.
      mv.visitLocalVariable("e", "Ljava/io/IOException;", null, catchStart, catchEnd, var + 1);
      break;
    default: // do nothing
  }
}

But when I tested, I keep getting NotSerializableException -- I thought I used try-catch to catch this exception.

I'm not sure whether I should add visitFrame for try-catch block (and I also don't know how to do it).

PS -- Any pointers about other better methods to do logging for every local variable would also be highly appreciated!


Solution

  • Your logic of building a try-catch block is correct, except that you are using variables var + 1 to var + 3 which may clash with uses by the original code. When I try your code to instrument an example specifically chosen such that it has no such variable clashes, it works.

    You could work-around such issues using a LocalVariablesSorter but it requires calls to newLocal to declare a variable for your injected code and since there’s no such call in your code, I assume, you’re not using LocalVariablesSorter.

    Generally, injecting code of such complexity, even potentially multiple times, not only is error prone, it might raise the code size significantly, up to the point that it exceeds the maximum code size of a method.

    The preferable approach is to move the complex code in a method on its own, which might even be delivered in precompiled form, i.e. created using ordinary Java source code, and only inject an invocation of that method.

    So, assuming a helper class like

    package mypackage;
    
    import java.io.*;
    import java.util.Base64;
    
    public class MyUtil {
        public static void printSerializedWithToStringFallback(Object someObject) {
            try {
                ByteArrayOutputStream bos = new ByteArrayOutputStream();
                ObjectOutputStream oos = new ObjectOutputStream(bos);
                oos.writeObject(someObject);
                oos.close();
                System.out.println(Base64.getEncoder().encodeToString(bos.toByteArray()));
              } catch(IOException ex) {
                System.out.println(someObject.toString());
              }
        }
    }
    

    You can inject the invocation like

    @Override
    public void visitVarInsn(int opcode, int var) {
        super.visitVarInsn(opcode, var);
        if(opcode == Opcodes.ASTORE) {
            super.visitVarInsn(Opcodes.ALOAD, var);
            super.visitMethodInsn(Opcodes.INVOKESTATIC, "mypackage/MyUtil",
                "printSerializedWithToStringFallback", "(Ljava/lang/Object;)V", false);
        }
    }
    

    Injecting this invocation does not introduce any branches, so the stack map table does not need to be recalculated. Even the requirements for the stack do not change. The injected code does not introduce new local variables and its highest stack size, after the aload, is identical to the stack size before the astore. So this simple instrumentation does not need the COMPUTE_FRAMES option and not even COMPUTE_MAXS.