Search code examples
javaannotationsbytecodejava-bytecode-asm

Adding RunTimeVisibleAnnotations to a Java class file using ASM


I am working on a project which requires me to add an annotation to a local variable directly in an existing class file so that it recreates the effect of the Java code below.

@Target({ElementType.LOCAL_VARIABLE, ElementType.TYPE, ElementType.TYPE_USE, ElementType.TYPE_PARAMETER})
@Retention(RetentionPolicy.RUNTIME)
public @interface MyAnnotation {}

public class MyClass {
    public static void main(String[] args) {
        int x = 512;
        @MyAnnotation int y = 5;
        ...
    }
}

The annotation part of the output of javap -p -v on the class file looks like this.

...
      RuntimeVisibleTypeAnnotations:
        0: #15(): LOCAL_VARIABLE, {start_pc=6, length=26, index=2}
          MyAnnotation

To do so, I have been looking into ASM. Now I have no experience with ASM whatsoever, but seeing some examples, I think I have got some idea about how to proceed. Here is my attempt.

public class ASMTransformer {
    public static void main(String[] args) throws Exception {
        FileInputStream fis = new FileInputStream(args[0]);
        ClassReader cr = new ClassReader(fis);
        ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
        cr.accept(new ASMClass(cw), ClassReader.EXPAND_FRAMES);
        FileOutputStream fos = new FileOutputStream(args[0]);
        fos.write(cw.toByteArray());
        fos.close();
    }

    public static class ASMClass extends ClassVisitor {
        public ASMClass(ClassVisitor cv) {
            super(Opcodes.ASM9, cv);
        }

        public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
            if (name == "main")
                return new ASMMethod(super.visitMethod(access, name, desc, signature, exceptions));
            else
                return super.visitMethod(access, name, desc, signature, exceptions);
        }
    }

    public static class ASMMethod extends MethodVisitor {
        public ASMMethod(MethodVisitor mv) {
            super(Opcodes.ASM9, mv);
        }

        public void visitEnd() {
            super.visitLocalVariableAnnotation(typeRef, typePath, start, end, index, descriptor, visible);
            super.visitEnd();
        }
    }
}

I figured out the values of some of the arguments in the visitLocalVariableAnnotation() by printing them out from an annotated class file. But the start Label[], the end Label[] and the index of the variable are the ones I am not able to figure out.

Can someone kindly verify if I am moving in the right direction and also help me with getting the values of these arguments?


Solution

  • These arrays allow you to to specify multiple variables and scopes to apply the annotation to. This allows shorter class files, especially for annotations with multiple (identical) values.

    The simplest approach is to pass arrays of length one, specifying the x variable and its scope. You can get the required information from the local variable table, assuming that the class has been compiled with debug information included.

    public static class ASMMethod extends MethodVisitor {
        public ASMMethod(MethodVisitor mv) {
            super(Opcodes.ASM9, mv);
        }
    
        @Override
        public void visitLocalVariable(String name, String descriptor, String signature,
                                       Label start, Label end, int index) {
    
            super.visitLocalVariable(name, descriptor, signature, start, end, index);
            if(name.equals("x")) {
                super.visitLocalVariableAnnotation(TypeReference.LOCAL_VARIABLE << 24, null,
                    new Label[] { start }, new Label[] { end }, new int[]{ index },
                    "LMyAnnotation;", true)
                    .visitEnd();
            }
        }
    }
    

    This produces this javap output, indicating that the annotation information has been stored once for x and once for y.

          RuntimeVisibleTypeAnnotations:
            0: #41(): LOCAL_VARIABLE, {start_pc=4, length=28, index=1}
              MyAnnotation
            1: #41(): LOCAL_VARIABLE, {start_pc=6, length=26, index=2}
    

    You can implement a “clone y’s annotation” logic instead, which also reduces the class file size:

    public static class ASMMethod extends MethodVisitor {
        private int xIndex, yIndex;
        private Label xStart, xEnd, yStart, yEnd;
    
        public ASMMethod(MethodVisitor mv) {
            super(Opcodes.ASM9, mv);
        }
    
        @Override
        public void visitLocalVariable(String name, String descriptor, String signature,
                                       Label start, Label end, int index) {
    
            super.visitLocalVariable(name, descriptor, signature, start, end, index);
            if(name.equals("x")) {
                xIndex = index;
                xStart = start;
                xEnd  = end;
            }
            else if(name.equals("y")) {
                yIndex = index;
                yStart = start;
                yEnd  = end;
            }
        }
    
        @Override
        public AnnotationVisitor visitLocalVariableAnnotation(int typeRef,
            TypePath typePath, Label[] start, Label[] end, int[] index,
            String descriptor, boolean visible) {
    
            if(Arrays.stream(index).anyMatch(ix -> ix == yIndex)) {
                int num = start.length, newSize = num + 1;
                index = Arrays.copyOf(index, newSize);
                index[num] = xIndex;
                start = Arrays.copyOf(start, newSize);
                start[num] = xStart;
                end = Arrays.copyOf(end, newSize);
                end[num] = xEnd;
            }
            return super.visitLocalVariableAnnotation(
                typeRef, typePath, start, end, index, descriptor, visible);
        }
    }
    

    This will copy all of y’s annotations to x, including their values. With your example class, javap now indicates that the annotation information is shared for x and y.

          RuntimeVisibleTypeAnnotations:
            0: #43(): LOCAL_VARIABLE, {start_pc=6, length=26, index=2; start_pc=4, length=28, index=1}
              MyAnnotation
    

    Some additional remarks on your code:

    You must not compare strings with ==. Replace the name == "main" with name.equals("main"), to make it work. But you can also reduce the code duplication of the identical super.visitMethod calls.

    Depending on the system you’re running on, not closing the FileInputStream can make your attempt of overwriting the same file fail. But it’s recommended to close resource as early as possible in general.

    The COMPUTE_FRAMES implies COMPUTE_MAXS, so you don’t need to combine the two. Further, COMPUTE_FRAMES means “compute all frames from scratch”, so it doesn’t use the original code’s frames. So, unless you’re performing some other processing like using Analyzer, you don’t need the original frames when using COMPUTE_FRAMES. So passing SKIP_FRAMES to the ClassReader is more appropriate then. In contrast, the option EXPAND_FRAMES will perform preparation work on the original frames for a processing that never happens here, it wastes CPU cycles.

    But when all you’re going to do, is to inject annotations, in other words, you’re not changing any executable code, there is no reason to meddle with stack map frames at all. The instrumented code can just keep the original frames, which is the simplest and most efficient processing.

    Incorporating all these points yields

    public class ASMTransformer {
        public static void main(String[] args) throws IOException {
            String clazz = args[0];
            byte[] code;
            try(InputStream fis = new FileInputStream(clazz)) {
                ClassReader cr = new ClassReader(fis);
                ClassWriter cw = new ClassWriter(cr, 0);
                cr.accept(new ASMClass(cw), 0);
                code = cw.toByteArray();
            }
            Files.write(Paths.get(clazz), code);
        }
    
        public static class ASMClass extends ClassVisitor {
            public ASMClass(ClassVisitor cv) {
                super(Opcodes.ASM9, cv);
            }
    
            @Override
            public MethodVisitor visitMethod(
                int access, String name, String desc, String sig, String[] exceptions) {
    
                MethodVisitor mv = super.visitMethod(access, name, desc, sig, exceptions);
                if(name.equals("main")) mv = new ASMMethod(mv);
                return mv;
            }
        }
    
        public static class ASMMethod extends MethodVisitor { // as above
        …
    }