如何扫描程序中带有指定注解的类和方法

痛定思痛。 2022-04-04 11:15 253阅读 0赞

使用springboot的人基本都知道swagger，那么swagger是如何生成swagger-ui.html页面的呢？相信大家都能猜到，就是扫描程序中带有指定注解的类(带有@RestController和 @Controller)和方法(@RequestMapping等)，然后又根据方法上的@ApiOperation和@ApiImplicitParams去生成页面上要显示的一些元素。

实际项目中我们也会有类似需求，例如权限校验，我们需要先扫描程序中有哪些接口(就是找有哪些类有@RestController注解)，然后在根据自定义的一些注解，扫描是否这些controller的方法需要配置权限校验。那么如何来实现？

# 第一步， #

根据basePackage，开始遍历所有的类，然后测试该类是否添加了@RestController注解。  
需要注意的是，当我们的类是在jar文件中时，不需要递归，但是在普通的文件中(目录存放文件)需要递归循环。

private List<Class<?>> getClassesWithAnnotationFromPackage(String packageName, Class<? extends Annotation> annotation) { 
            List<Class<?>> classList = new ArrayList<Class<?>>();
            String packageDirName = packageName.replace('.', '/');
            Enumeration<URL> dirs = null;
    
            try { 
                dirs = Thread.currentThread().getContextClassLoader().getResources(packageDirName);
            }
            catch (IOException e) { 
                log.error("Failed to get resource", e);
                return null;
            }
    
            while (dirs.hasMoreElements()) { 
                URL url = dirs.nextElement();//file:/D:/E/workspaceGitub/springboot/JSONDemo/target/classes/com/yq/controller
                String protocol = url.getProtocol();//file
    
                //https://docs.oracle.com/javase/7/docs/api/java/net/URL.html
                //http, https, ftp, file, and jar
                //本文只需要处理file和jar
                if ("file".equals(protocol) ) { 
                    String filePath = null;
                    try { 
                        filePath = URLDecoder.decode(url.getFile(), "UTF-8");///D:/E/workspaceGitub/springboot/JSONDemo/target/classes/com/yq/controller
                    }
                    catch (UnsupportedEncodingException e) { 
                        log.error("Failed to decode class file", e);
                    }
    
                    filePath = filePath.substring(1);
                    getClassesWithAnnotationFromFilePath(packageName, filePath, classList, annotation);
                } else if ("jar".equals(protocol)) { 
                    JarFile jar = null;
                    try { 
                        jar = ((JarURLConnection) url.openConnection()).getJarFile();
                        //扫描jar包文件 并添加到集合中
                    }
                    catch (Exception e) { 
                        log.error("Failed to decode class jar", e);
                    }
    
                    List<Class<?>> alClassList = new ArrayList<Class<?>>();
                    findClassesByJar(packageName, jar, alClassList);
                    getClassesWithAnnotationFromAllClasses(alClassList, annotation, classList);
                }
                else { 
                    log.warn("can't process the protocol={}", protocol);
                }
            }
    
            return classList;
        }

private static void findClassesByJar(String pkgName, JarFile jar, List<Class<?>> classes) { 
            String pkgDir = pkgName.replace(".", "/");
            Enumeration<JarEntry> entry = jar.entries();
    
            while (entry.hasMoreElements()) { 
                // 获取jar里的一个实体 可以是目录 和一些jar包里的其他文件 如META-INF等文
                JarEntry jarEntry = entry.nextElement();
                String name = jarEntry.getName();
                // 如果是以/开头的
                if (name.charAt(0) == '/') { 
                    // 获取后面的字符串
                    name = name.substring(1);
                }
    
                if (jarEntry.isDirectory() || !name.startsWith(pkgDir) || !name.endsWith(".class")) { 
                    continue;
                }
                //如果是一个.class文件 而且不是目录
                // 去掉后面的".class" 获取真正的类名
                String className = name.substring(0, name.length() - 6);
                Class<?> tempClass = loadClass(className.replace("/", "."));
                // 添加到集合中去
                if (tempClass != null) { 
                    classes.add(tempClass);
                }
            }
        }
    
        /** * 加载类 * @param fullClsName 类全名 * @return */
        private static Class<?> loadClass(String fullClsName ) { 
            try { 
                return Thread.currentThread().getContextClassLoader().loadClass(fullClsName );
            } catch (ClassNotFoundException e) { 
                log.error("PkgClsPath loadClass", e);
            }
            return null;
        }
    
    
        //filePath is like this 'D:/E/workspaceGitub/springboot/JSONDemo/target/classes/com/yq/controller'
        private void getClassesWithAnnotationFromFilePath(String packageName, String filePath, List<Class<?>> classList,
                                               Class<? extends Annotation> annotation) { 
            Path dir = Paths.get(filePath);//D:\E\workspaceGitub\springboot\JSONDemo\target\classes\com\yq\controller
    
            try (DirectoryStream<Path> stream = Files.newDirectoryStream(dir)) { 
                for(Path path : stream) { 
                    String fileName = String.valueOf(path.getFileName()); // for current dir , it is 'helloworld'
                    //如果path是目录的话， 此处需要递归，
                    boolean isDir = Files.isDirectory(path);
                    if(isDir) { 
                        getClassesWithAnnotationFromFilePath(packageName + "." + fileName , path.toString(), classList, annotation);
                    }
                    else  { 
                        String className = fileName.substring(0, fileName.length() - 6);
    
                        Class<?> classes = null;
                        String fullClassPath = packageName + "." + className;
                        try { 
                            log.info("fullClassPath={}", fullClassPath);
                            classes = Thread.currentThread().getContextClassLoader().loadClass(fullClassPath);
                        }
                        catch (ClassNotFoundException e) { 
                            log.error("Failed to find class={}", fullClassPath, e);
                        }
    
                        if (null != classes && null != classes.getAnnotation(annotation)) { 
                            classList.add(classes);
                        }
                    }
                }
            }
            catch (IOException e) { 
                log.error("Failed to read class file", e);
            }
        }
    
        private void getClassesWithAnnotationFromAllClasses(List<Class<?>> inClassList,
                                                          Class<? extends Annotation> annotation, List<Class<?>> outClassList) { 
            for(Class<?> myClasss : inClassList) { 
                if (null != myClasss && null != myClasss.getAnnotation(annotation)) { 
                    outClassList.add(myClasss);
                }
            }
        }

# 第二步 #

根据第一步提供的controller，遍历该类的所有方法，检查该方法是否添加MyChecker注解(MyChecker为自定义的注解)

private void geMethodWithAnnotationFromFilePath(String fullClassPath, Map<String, String> checkIdMethodMap,
                                                                 Class<? extends Annotation> methodAnnotation) { 
            Class<?> classes = null;
            try { 
                log.info("fullClassPath={}", fullClassPath);
                classes = Thread.currentThread().getContextClassLoader().loadClass(fullClassPath);
    
                Method[] methods = classes.getDeclaredMethods();
    
                for (Method method : methods) { 
                    MyChecker myAnnotation = method.getAnnotation(MyChecker.class);
                    if (null != myAnnotation) { 
                        checkIdMethodMap.put (myAnnotation.id(), method.getName() );
                    }
                }
            }
            catch (ClassNotFoundException e) { 
                log.error("Failed to find class={}", fullClassPath, e);
            }
        }

# 效果截图 #

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3J1c3NsZQ_size_16_color_FFFFFF_t_70]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3J1c3NsZQ_size_16_color_FFFFFF_t_70]: /images/20220404/aa1dbea53a4240ef93ecd627ac383403.png