From bc70853b0c9991c1ba9e1b8ddea1f5376a0ee9e4 Mon Sep 17 00:00:00 2001
From: hengyunabc <hengyunabc@gmail.com>
Date: Wed, 2 Dec 2020 00:55:00 +0800
Subject: [PATCH] support arthas.enhanceLoaders config/add jboss
 ModuleClassLoader testcase. #1596

---
 core/pom.xml                                  |  8 +++
 core/src/main/java/arthas.properties          |  2 +
 .../taobao/arthas/core/config/Configure.java  | 13 +++++
 .../arthas/core/server/ArthasBootstrap.java   | 42 +++++++++++---
 .../core/util/InstrumentationUtils.java       | 14 +++++
 .../core/server/ArthasBootstrapTest.java      | 55 +++++++++++++++++++
 6 files changed, 127 insertions(+), 7 deletions(-)
 create mode 100644 core/src/test/java/com/taobao/arthas/core/server/ArthasBootstrapTest.java

diff --git a/core/pom.xml b/core/pom.xml
index 684e3d822..494a423bc 100644
--- a/core/pom.xml
+++ b/core/pom.xml
@@ -235,6 +235,14 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>org.jboss.modules</groupId>
+            <artifactId>jboss-modules</artifactId>
+            <version>1.11.0.Final</version>
+            <scope>test</scope>
+            <optional>true</optional>
+        </dependency>
+
         <dependency>
             <groupId>org.benf</groupId>
             <artifactId>cfr</artifactId>
diff --git a/core/src/main/java/arthas.properties b/core/src/main/java/arthas.properties
index aab9953fe..99d9f256a 100644
--- a/core/src/main/java/arthas.properties
+++ b/core/src/main/java/arthas.properties
@@ -6,6 +6,8 @@ arthas.ip=127.0.0.1
 # seconds
 arthas.sessionTimeout=1800
 
+arthas.enhanceLoaders=java.lang.ClassLoader
+
 #arthas.appName=demoapp
 #arthas.tunnelServer=ws://127.0.0.1:7777/ws
 #arthas.agentId=mmmmmmyiddddd
diff --git a/core/src/main/java/com/taobao/arthas/core/config/Configure.java b/core/src/main/java/com/taobao/arthas/core/config/Configure.java
index c0982b01b..1ebe0db8b 100644
--- a/core/src/main/java/com/taobao/arthas/core/config/Configure.java
+++ b/core/src/main/java/com/taobao/arthas/core/config/Configure.java
@@ -28,6 +28,11 @@ public class Configure {
     private String tunnelServer;
     private String agentId;
 
+    /**
+     * 需要被增强的ClassLoader的全类名，多个用英文 , 分隔
+     */
+    private String enhanceLoaders;
+
     /**
      * <pre>
      * 1. 如果显式传入 arthas.agentId ，则直接使用
@@ -135,6 +140,14 @@ public class Configure {
         this.appName = appName;
     }
 
+    public String getEnhanceLoaders() {
+        return enhanceLoaders;
+    }
+
+    public void setEnhanceLoaders(String enhanceLoaders) {
+        this.enhanceLoaders = enhanceLoaders;
+    }
+
     /**
      * 序列化成字符串
      *
diff --git a/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java b/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java
index ef6e77de1..ab758a121 100644
--- a/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java
+++ b/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java
@@ -9,10 +9,12 @@ import java.lang.reflect.Method;
 import java.security.CodeSource;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Properties;
+import java.util.Set;
 import java.util.Timer;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
@@ -25,9 +27,12 @@ import com.alibaba.arthas.deps.ch.qos.logback.classic.LoggerContext;
 import com.alibaba.arthas.deps.org.slf4j.Logger;
 import com.alibaba.arthas.deps.org.slf4j.LoggerFactory;
 import com.alibaba.arthas.tunnel.client.TunnelClient;
+import com.alibaba.bytekit.asm.instrument.InstrumentConfig;
 import com.alibaba.bytekit.asm.instrument.InstrumentParseResult;
 import com.alibaba.bytekit.asm.instrument.InstrumentTemplate;
 import com.alibaba.bytekit.asm.instrument.InstrumentTransformer;
+import com.alibaba.bytekit.asm.matcher.SimpleClassMatcher;
+import com.alibaba.bytekit.utils.AsmUtils;
 import com.alibaba.bytekit.utils.IOUtils;
 import com.taobao.arthas.common.AnsiLog;
 import com.taobao.arthas.common.ArthasConstants;
@@ -59,7 +64,9 @@ import com.taobao.arthas.core.shell.term.impl.http.api.HttpApiHandler;
 import com.taobao.arthas.core.shell.term.impl.httptelnet.HttpTelnetTermServer;
 import com.taobao.arthas.core.util.ArthasBanner;
 import com.taobao.arthas.core.util.FileUtils;
+import com.taobao.arthas.core.util.InstrumentationUtils;
 import com.taobao.arthas.core.util.LogUtil;
+import com.taobao.arthas.core.util.StringUtils;
 import com.taobao.arthas.core.util.UserStatUtil;
 import com.taobao.arthas.core.util.affect.EnhancerAffect;
 import com.taobao.arthas.core.util.matcher.WildcardMatcher;
@@ -126,10 +133,12 @@ public class ArthasBootstrap {
         // 3. init logger
         loggerContext = LogUtil.initLooger(arthasEnvironment);
 
-        // 4. init beans
+        // 4. 增强ClassLoader
+        enhanceClassLoader();
+        // 5. init beans
         initBeans();
 
-        // 5. start agent server
+        // 6. start agent server
         bind(configure);
 
         executorService = Executors.newScheduledThreadPool(1, new ThreadFactory() {
@@ -182,19 +191,38 @@ public class ArthasBootstrap {
                 throw new IllegalStateException("can not find " + ARTHAS_SPY_JAR);
             }
         }
+    }
+
+    void enhanceClassLoader() throws IOException, UnmodifiableClassException {
+        if (configure.getEnhanceLoaders() == null) {
+            return;
+        }
+        Set<String> loaders = new HashSet<String>();
+        for (String s : configure.getEnhanceLoaders().split(",")) {
+            loaders.add(s.trim());
+        }
 
         // 增强 ClassLoader#loadClsss ，解决一些ClassLoader加载不到 SpyAPI的问题
         // https://github.com/alibaba/arthas/issues/1596
-        InstrumentTemplate template = new InstrumentTemplate();
         byte[] classBytes = IOUtils.getBytes(ArthasBootstrap.class.getClassLoader()
                 .getResourceAsStream(ClassLoader_Instrument.class.getName().replace('.', '/') + ".class"));
-        template.addInstrumentClass(classBytes);
 
-        InstrumentParseResult instrumentParseResult = template.build();
+        SimpleClassMatcher matcher = new SimpleClassMatcher(loaders);
+        InstrumentConfig instrumentConfig = new InstrumentConfig(AsmUtils.toClassNode(classBytes), matcher);
+
+        InstrumentParseResult instrumentParseResult = new InstrumentParseResult();
+        instrumentParseResult.addInstrumentConfig(instrumentConfig);
         classLoaderInstrumentTransformer = new InstrumentTransformer(instrumentParseResult);
-        instrumentation.addTransformer(classLoaderInstrumentTransformer);
-    }
+        instrumentation.addTransformer(classLoaderInstrumentTransformer, true);
 
+        if (loaders.size() == 1 && loaders.contains(ClassLoader.class.getName())) {
+            // 如果只增强 java.lang.ClassLoader，可以减少查找过程
+            instrumentation.retransformClasses(ClassLoader.class);
+        } else {
+            InstrumentationUtils.trigerRetransformClasses(instrumentation, loaders);
+        }
+    }
+    
     private void initArthasEnvironment(Map<String, String> argsMap) throws IOException {
         if (arthasEnvironment == null) {
             arthasEnvironment = new ArthasEnvironment();
diff --git a/core/src/main/java/com/taobao/arthas/core/util/InstrumentationUtils.java b/core/src/main/java/com/taobao/arthas/core/util/InstrumentationUtils.java
index 51b781a1f..3bcaa1c09 100644
--- a/core/src/main/java/com/taobao/arthas/core/util/InstrumentationUtils.java
+++ b/core/src/main/java/com/taobao/arthas/core/util/InstrumentationUtils.java
@@ -2,6 +2,7 @@ package com.taobao.arthas.core.util;
 
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.Instrumentation;
+import java.util.Collection;
 import java.util.Set;
 
 import com.alibaba.arthas.deps.org.slf4j.Logger;
@@ -38,4 +39,17 @@ public class InstrumentationUtils {
             inst.removeTransformer(transformer);
         }
     }
+
+    public static void trigerRetransformClasses(Instrumentation inst, Collection<String> classes) {
+        for (Class<?> clazz : inst.getAllLoadedClasses()) {
+            if (classes.contains(clazz.getName())) {
+                try {
+                    inst.retransformClasses(clazz);
+                } catch (Throwable e) {
+                    String errorMsg = "retransformClasses class error, name: " + clazz.getName();
+                    logger.error(errorMsg, e);
+                }
+            }
+        }
+    }
 }
diff --git a/core/src/test/java/com/taobao/arthas/core/server/ArthasBootstrapTest.java b/core/src/test/java/com/taobao/arthas/core/server/ArthasBootstrapTest.java
new file mode 100644
index 000000000..00a1177e5
--- /dev/null
+++ b/core/src/test/java/com/taobao/arthas/core/server/ArthasBootstrapTest.java
@@ -0,0 +1,55 @@
+package com.taobao.arthas.core.server;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.lang.instrument.Instrumentation;
+
+import org.jboss.modules.ModuleClassLoader;
+import org.junit.Test;
+import org.mockito.Mockito;
+import org.mockito.internal.util.reflection.FieldSetter;
+
+import com.taobao.arthas.core.bytecode.TestHelper;
+import com.taobao.arthas.core.config.Configure;
+
+import net.bytebuddy.agent.ByteBuddyAgent;
+
+/**
+ * 
+ * @author hengyunabc 2020-12-02
+ *
+ */
+public class ArthasBootstrapTest {
+    @Test
+    public void test() throws Exception {
+        Instrumentation instrumentation = ByteBuddyAgent.install();
+        TestHelper.appendSpyJar(instrumentation);
+
+        ArthasBootstrap arthasBootstrap = Mockito.mock(ArthasBootstrap.class);
+        Mockito.doCallRealMethod().when(arthasBootstrap).enhanceClassLoader();
+
+        Configure configure = Mockito.mock(Configure.class);
+        Mockito.when(configure.getEnhanceLoaders())
+                .thenReturn("java.lang.ClassLoader,org.jboss.modules.ConcurrentClassLoader");
+        FieldSetter.setField(arthasBootstrap, ArthasBootstrap.class.getDeclaredField("configure"), configure);
+        FieldSetter.setField(arthasBootstrap, ArthasBootstrap.class.getDeclaredField("instrumentation"),
+                instrumentation);
+
+        org.jboss.modules.ModuleClassLoader moduleClassLoader = Mockito.mock(ModuleClassLoader.class);
+
+        boolean flag = false;
+        try {
+            moduleClassLoader.loadClass("java.arthas.SpyAPI");
+        } catch (Exception e) {
+            flag = true;
+        }
+        assertThat(flag).isTrue();
+
+        arthasBootstrap.enhanceClassLoader();
+
+        Class<?> loadClass = moduleClassLoader.loadClass("java.arthas.SpyAPI");
+
+        System.err.println(loadClass);
+
+    }
+}