From 82a3d69cc3d86bd21343bfb176d61cef52e80321 Mon Sep 17 00:00:00 2001
From: hengyunabc <hengyunabc@gmail.com>
Date: Mon, 4 May 2020 14:00:38 +0800
Subject: [PATCH] add AsmAnnotationUtils

---
 .../bytekit/utils/AsmAnnotationUtils.java     | 77 ++++++++++++++++++
 .../bytekit/utils/AsmAnnotationUtilsTest.java | 78 +++++++++++++++++++
 2 files changed, 155 insertions(+)
 create mode 100644 bytekit/src/main/java/com/taobao/arthas/bytekit/utils/AsmAnnotationUtils.java
 create mode 100644 bytekit/src/test/java/com/taobao/arthas/bytekit/utils/AsmAnnotationUtilsTest.java

diff --git a/bytekit/src/main/java/com/taobao/arthas/bytekit/utils/AsmAnnotationUtils.java b/bytekit/src/main/java/com/taobao/arthas/bytekit/utils/AsmAnnotationUtils.java
new file mode 100644
index 000000000..70d381d39
--- /dev/null
+++ b/bytekit/src/main/java/com/taobao/arthas/bytekit/utils/AsmAnnotationUtils.java
@@ -0,0 +1,77 @@
+package com.taobao.arthas.bytekit.utils;
+
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+
+import com.alibaba.arthas.deps.org.objectweb.asm.tree.AnnotationNode;
+
+/**
+ * 
+ * @author hengyunabc 2020-05-04
+ *
+ */
+public class AsmAnnotationUtils {
+
+    public static List<String> queryAnnotationInfo(List<AnnotationNode> annotations, String annotationType,
+            String key) {
+        List<String> result = new ArrayList<String>();
+        if (annotations != null) {
+            for (AnnotationNode annotationNode : annotations) {
+                if (annotationNode.desc.equals(annotationType)) {
+                    if (annotationNode.values != null) {
+                        Iterator<Object> iterator = annotationNode.values.iterator();
+                        while (iterator.hasNext()) {
+                            String name = (String) iterator.next();
+                            Object values = iterator.next();
+                            if (key.equals(name)) {
+                                result.addAll((List<String>) values);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        return result;
+    }
+
+    public static void addAnnotationInfo(List<AnnotationNode> annotations, String annotationType, String key,
+            String value) {
+
+        AnnotationNode annotationNode = null;
+        for (AnnotationNode tmp : annotations) {
+            if (tmp.desc.equals(annotationType)) {
+                annotationNode = tmp;
+            }
+        }
+
+        if (annotationNode == null) {
+            annotationNode = new AnnotationNode(annotationType);
+            annotations.add(annotationNode);
+        }
+
+        if (annotationNode.values == null) {
+            annotationNode.values = new ArrayList<Object>();
+        }
+
+        // 查找有没有对应的key
+        String name = null;
+        List<String> values = null;
+        Iterator<Object> iterator = annotationNode.values.iterator();
+        while (iterator.hasNext()) {
+            if (key.equals(iterator.next())) {
+                values = (List<String>) iterator.next();
+            } else {
+                iterator.next();
+            }
+        }
+        if (values == null) {
+            values = new ArrayList<String>();
+            annotationNode.values.add(key);
+            annotationNode.values.add(values);
+        }
+        if (!values.contains(values)) {
+            values.add(value);
+        }
+    }
+}
diff --git a/bytekit/src/test/java/com/taobao/arthas/bytekit/utils/AsmAnnotationUtilsTest.java b/bytekit/src/test/java/com/taobao/arthas/bytekit/utils/AsmAnnotationUtilsTest.java
new file mode 100644
index 000000000..366fac9a7
--- /dev/null
+++ b/bytekit/src/test/java/com/taobao/arthas/bytekit/utils/AsmAnnotationUtilsTest.java
@@ -0,0 +1,78 @@
+package com.taobao.arthas.bytekit.utils;
+
+import java.io.IOException;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import java.util.Arrays;
+
+import org.assertj.core.api.Assertions;
+import org.junit.Test;
+import org.springframework.stereotype.Service;
+
+import com.alibaba.arthas.deps.org.objectweb.asm.Type;
+import com.alibaba.arthas.deps.org.objectweb.asm.tree.ClassNode;
+import com.taobao.arthas.bytekit.utils.AsmAnnotationUtils;
+import com.taobao.arthas.bytekit.utils.AsmUtils;
+
+/**
+ * 
+ * @author hengyunabc 2020-05-04
+ *
+ */
+public class AsmAnnotationUtilsTest {
+
+    @Target(value = { ElementType.TYPE, ElementType.METHOD })
+    @Retention(value = RetentionPolicy.RUNTIME)
+    public @interface AdviceInfo {
+
+        public String[] adviceInfos();
+    }
+
+    @Service
+    @AdviceInfo(adviceInfos = { "xxxx", "yyy" })
+    static class AAA {
+
+        @AdviceInfo(adviceInfos = { "mmm", "yyy" })
+        public void test() {
+
+        }
+
+    }
+
+    @Service
+    static class BBB {
+        public void test() {
+        }
+    }
+
+    @Test
+    public void test() throws IOException {
+        ClassNode classNodeA = AsmUtils.loadClass(AAA.class);
+
+        ClassNode classNodeB = AsmUtils.loadClass(BBB.class);
+
+        Assertions.assertThat(AsmAnnotationUtils.queryAnnotationInfo(classNodeA.visibleAnnotations,
+                Type.getDescriptor(AdviceInfo.class), "adviceInfos")).isEqualTo(Arrays.asList("xxxx", "yyy"));
+
+        AsmAnnotationUtils.addAnnotationInfo(classNodeA.visibleAnnotations, Type.getDescriptor(AdviceInfo.class),
+                "adviceInfos", "fff");
+
+        Assertions
+                .assertThat(AsmAnnotationUtils.queryAnnotationInfo(classNodeA.visibleAnnotations,
+                        Type.getDescriptor(AdviceInfo.class), "adviceInfos"))
+                .isEqualTo(Arrays.asList("xxxx", "yyy", "fff"));
+
+        Assertions.assertThat(AsmAnnotationUtils.queryAnnotationInfo(classNodeB.visibleAnnotations,
+                Type.getDescriptor(AdviceInfo.class), "adviceInfos")).isEmpty();
+
+        AsmAnnotationUtils.addAnnotationInfo(classNodeB.visibleAnnotations, Type.getDescriptor(AdviceInfo.class),
+                "adviceInfos", "fff");
+
+        Assertions.assertThat(AsmAnnotationUtils.queryAnnotationInfo(classNodeB.visibleAnnotations,
+                Type.getDescriptor(AdviceInfo.class), "adviceInfos")).isEqualTo(Arrays.asList("fff"));
+
+    }
+
+}