Java: Find all direct & indirect references to annotated methods

Suppose I have defined the following API in api.jar:

public class API {
    // @RequiredActionsPermitted({"1","3"})
    public void higherApi1() { highApi1(); highApi3(); }
    
    public void highApi1() { lowApi1(); }
    public void highApi2() { lowApi1(); lowApi2(); }
    public void highApi3() { lowApi3(); }
    
    @RequiredActionsPermitted({"1"})    public void lowApi1() {}
    @RequiredActionsPermitted({"2"})    public void lowApi2() {}
    @RequiredActionsPermitted({"3"})    public void lowApi3() {}

    public void zMethod() { lowApi1(); lowApi2(); lowApi3(); }
}

And a main method like this in another client.jar (the above api.jar is included as a binary dependency):

public static void main(String[] args) {
    API api = new API();
    api.higherApi1();
}

Given that this main() method calls higherApi1(), which in turn indirectly calls lowApi1() and lowApi3(), I would like to generate some XML output like this:

<requiredActions>
   <action>1</action>
   <action>3</action>
</requiredActions>

Action 2 shouldn't be part of the output, since the main() method does not call lowApi2(), neither directly or indirectly.

Of course, I could easily annotate higherApi3() with the annotation that is commented out in the example, but this will be hard to maintain if you have a lot of indirect methods calls.

Edit: I now have a solution based on ASM that mostly works:

import java.io.BufferedInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

import org.objectweb.asm.AnnotationVisitor;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;

// TODO Add documentation
// TODO clean up this class
public class TokenSpecGenerator {
    public static void main(String[] jars) throws IOException {
        Map<String, Set<String>> methodsToAnnotationValuesMap = new HashMap<>();
        if (findAnnotatedMethods(jars, methodsToAnnotationValuesMap)) {
            while (findIndirectMethodInvocations(jars, methodsToAnnotationValuesMap)) {
            }
        }
        // TODO Should we also report any API method invocations outside of
        // methods?
        // (i.e. field initializers, static blocks, ...)
        for (Map.Entry<String, Set<String>> entry : methodsToAnnotationValuesMap.entrySet()) {
            if (!entry.getKey().startsWith("com/something/myapi")) {
                System.out.println(entry.getKey() + ": " + entry.getValue());
            }
        }
    }

    protected static boolean findAnnotatedMethods(String[] jars, Map<String, Set<String>> methodsToAnnotationValuesMap) {
        System.out.println("Scanning for methods annotated with @RequiredActionsPermitted");
        FindAnnotatedMethods classVisitor = new FindAnnotatedMethods(methodsToAnnotationValuesMap);
        visitClasses(jars, classVisitor);
        System.out.println("Methods and required actions permitted found until now:");
        System.out.println(methodsToAnnotationValuesMap);
        return classVisitor.hasFoundNew();
    }

    protected static boolean findIndirectMethodInvocations(String[] jars, Map<String, Set<String>> methodsToAnnotationValuesMap) {
        System.out.println("Next round of scanning for methods that indirectly call methods annotated with @RequiredActionsPermitted");
        FindMethodInvocations classVisitor = new FindMethodInvocations(methodsToAnnotationValuesMap);
        visitClasses(jars, classVisitor);
        System.out.println("Methods and required actions permitted found until now:");
        System.out.println(methodsToAnnotationValuesMap);
        return classVisitor.hasFoundNew();
    }

    protected static void visitClasses(String[] jars, ClassVisitor classVisitor) {
        for (String jar : jars) {
            JarFile jarFile = null;
            try {
                try {
                    jarFile = new JarFile(jar);
                    Enumeration<JarEntry> entries = jarFile.entries();
        
                    while (entries.hasMoreElements()) {
                        JarEntry entry = entries.nextElement();
        
                        if (entry.getName().endsWith(".class")) {
                            try {
                                InputStream stream = null;
                                try {
                                    stream = new BufferedInputStream(jarFile.getInputStream(entry), 1024);
                                    new ClassReader(stream).accept(classVisitor, 0);
                                } finally {
                                    stream.close();
                                }
                            } catch ( IOException e ) { e.printStackTrace(); }
                        }
                    }
                } finally {
                    jarFile.close();
                }
            } catch ( IOException e ) { e.printStackTrace(); }
        }
    }

    static abstract class AbstractClassVisitor extends ClassVisitor {
        private final Map<String, Set<String>> methodsToAnnotationValuesMap;
        private boolean foundNew = false;
        private String className;

        public AbstractClassVisitor(Map<String, Set<String>> methodsToAnnotationValuesMap) {
            super(Opcodes.ASM5);
            this.methodsToAnnotationValuesMap = methodsToAnnotationValuesMap;
        }
        
        @Override
        public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
            this.className = name;
            super.visit(version, access, name, signature, superName, interfaces);
        }

        protected void addValues(String method, Set<String> values) {
            Set<String> existingValues = methodsToAnnotationValuesMap.get(method);
            if (existingValues == null) {
                existingValues = new HashSet<String>();
                methodsToAnnotationValuesMap.put(method, existingValues);
            }
            foundNew |= existingValues.addAll(values);
        }

        public boolean hasFoundNew() {
            return foundNew;
        }
        
        public String getClassName() {
            return className;
        }
        
        public Map<String, Set<String>> getMethodsToAnnotationValuesMap() {
            return methodsToAnnotationValuesMap;
        }
    }

    static abstract class AbstractMethodVisitor extends MethodVisitor {
        private final String className;
        private final String methodName;
        
        public AbstractMethodVisitor(String className, String methodName) {
            super(Opcodes.ASM5);
            this.className = className;
            this.methodName = methodName;
        }

        // TODO Add method parameter types
        public String getMethodDescription() {
            return className+"."+methodName;
        }
    }

    static class FindAnnotatedMethods extends AbstractClassVisitor {
        public FindAnnotatedMethods(Map<String, Set<String>> methodsToAnnotationValuesMap) {
            super(methodsToAnnotationValuesMap);
        }

        @Override
        public MethodVisitor visitMethod(final int access, final String name, final String desc, final String signature, final String[] exceptions) {
            return new AbstractMethodVisitor(getClassName(), name) {
                private Set<String> values = new HashSet<>();
                @Override
                public AnnotationVisitor visitAnnotation(String desc, boolean visible) {
                    if ("Lcom/something/myapi/annotation/RequiredActionsPermitted;".equals(desc)) {
                        return new AnnotationVisitor(Opcodes.ASM5) {
                            @Override
                            public AnnotationVisitor visitArray(String name) {
                                if ( !"value".equals(name) ) {
                                    return null;
                                } else {
                                    return new AnnotationVisitor(Opcodes.ASM5) {
                                        @Override
                                        public void visit(String name, Object value) {
                                            values.add((String)value);
                                        }
                                    };
                                }
                            }
                        };
                    }
                    return null;
                }

                @Override
                public void visitEnd() {
                    if ( !values.isEmpty() ) {
                        addValues(getMethodDescription(), values);
                    }
                }
            };
        }
    }
    
    static class FindMethodInvocations extends AbstractClassVisitor {
        public FindMethodInvocations(Map<String, Set<String>> methodsToAnnotationValuesMap) {
            super(methodsToAnnotationValuesMap);
        }

        @Override
        public MethodVisitor visitMethod(final int access, final String name, final String desc, final String signature, final String[] exceptions) {
            return new AbstractMethodVisitor(getClassName(), name) {
                private Set<String> values = new HashSet<>();
                @Override
                public void visitMethodInsn(int opcode, String owner, String name, String desc, boolean itf) {
                    String method = owner+"."+name;
                    if ( getMethodsToAnnotationValuesMap().containsKey(method) ) {
                        values.addAll(getMethodsToAnnotationValuesMap().get(method));
                    }
                }

                @Override
                public void visitEnd() {
                    if ( !values.isEmpty() ) {
                        addValues(getMethodDescription(), values);
                    }
                }
            };
        }
    }

}

I still have the following questions/issues though:

• At the moment, methods are matched only on class and method name. To correctly handle overloaded methods, matching should also include method parameter types. Any idea how to get this information in the ASM visitors?

• Any better approaches for matching method invocations in visitMethodInsn() against method information previously stored by visitMethod() in the methodsToAnnotationValuesMap?

• Apart for clean-up and adding JavaDoc, any other suggestions for improving this implementation?

Thanks, Ruud

Previous approach

Just for reference, previously I attempted to implement this using Annotation Processors. Below is my annotation processor implementation that finds some indirect method calls but not all (depending on the order in which Javac compiles the methods and classes). I have now switched to the ASM-based approach listed above, but leaving this here as an example of annotation processing.

  @SupportedAnnotationTypes("com.something.annotation.RequiredActionsPermitted")
  @SupportedSourceVersion(SourceVersion.RELEASE_7)
  public class RequiredActionsPermittedProcessor extends AbstractProcessor implements TaskListener {
      private final Map<String, List<String>  >    callers = new HashMap<>    ();
      Trees trees;

      @Override
      public synchronized void init(ProcessingEnvironment processingEnv) {
          super.init(processingEnv);
          trees = Trees.instance(processingEnv);
          JavacTask.instance(processingEnv).setTaskListener(this);
      }

      @Override
      public boolean process(Set<? extends TypeElement>    annotations, RoundEnvironment roundEnv) {
          return true;
      }
      
      

      @Override
      public void finished(final TaskEvent taskEvt) {
          if (taskEvt.getKind() == TaskEvent.Kind.ANALYZE) {
              //System.out.println("!!!! TEST4 !!!!");
              taskEvt.getCompilationUnit().accept(new TreeScanner<Void, Void> () {
                  private MethodTree parentMethod = null;
                  
                  @Override
                  public Void visitMethod(MethodTree methodTree, Void arg1) {
                      this.parentMethod = methodTree;
                      
                      return super.visitMethod(methodTree, arg1);
                  }
                  
                  @Override
                  public Void visitMethodInvocation(MethodInvocationTree methodInv, Void v) {
                      //System.out.println("!!!! TEST5 !!!!: "+methodInv);
                      JCTree jcTree = (JCTree) methodInv.getMethodSelect();
                      Element method = TreeInfo.symbol(jcTree);
                      RequiredActionsPermitted RequiredActionsPermitted = method.getAnnotation(RequiredActionsPermitted.class);
                      //System.out.println("!!!! TEST0 !!!!: "+method.getSimpleName().toString());
                      String methodName = method.getSimpleName().toString();
                      if ( parentMethod != null ) {
                          String parentName = parentMethod.getName().toString();
                          List<String>     values = null;
                          if (RequiredActionsPermitted != null ) {
                              values = Arrays.asList(RequiredActionsPermitted.value());
                          } else if ( callers.containsKey(methodName) ) {
                              values = callers.get(methodName);
                          }
                          if ( values != null ) {
                              List<String>     currentValues = callers.get(parentName);
                              if ( currentValues == null ) {
                                  currentValues = new ArrayList<> ();
                                  callers.put(parentName, currentValues);
                              }
                              currentValues.addAll(values);
                              System.out.println("!!!! TEST4 !!!!: Parent "+parentName);
                              System.out.println("!!!! TEST5 !!!!: Calls "+jcTree+" - "+String.join(",", values));
                          }
                      }
                      return super.visitMethodInvocation(methodInv, v);
                  }
              }, null);
          }
      }

      @Override
      public void started(TaskEvent taskEvt) {}
  }