support new instrument/define config format

hengyunabc · hengyunabc · commit f83a9e63bccb · 2025-10-23T11:47:40.000+08:00
diff --git a/bytekit-core/src/main/java/com/alibaba/bytekit/asm/instrument/InstrumentConfig.java b/bytekit-core/src/main/java/com/alibaba/bytekit/asm/instrument/InstrumentConfig.java
@@ -1,5 +1,8 @@
 package com.alibaba.bytekit.asm.instrument;
 
+import java.util.ArrayList;
+import java.util.List;
+
 import com.alibaba.bytekit.asm.matcher.ClassMatcher;
 import com.alibaba.deps.org.objectweb.asm.tree.ClassNode;
 
@@ -29,6 +32,11 @@ public class InstrumentConfig {
      */
     private boolean triggerRetransform = false;
 
+    /**
+     * 与这个 instrument 配置关联的 define 类列表，只有当这个 instrument 匹配时才会 define 这些类
+     */
+    private List<DefineConfig> defineConfigs = new ArrayList<DefineConfig>();
+
     public InstrumentConfig(ClassNode instrumentClassNode, ClassMatcher classMatcher) {
         this(instrumentClassNode, classMatcher, false, false);
     }
@@ -73,4 +81,16 @@ public void setTriggerRetransform(boolean triggerRetransform) {
         this.triggerRetransform = triggerRetransform;
     }
 
+    public List<DefineConfig> getDefineConfigs() {
+        return defineConfigs;
+    }
+
+    public void setDefineConfigs(List<DefineConfig> defineConfigs) {
+        this.defineConfigs = defineConfigs;
+    }
+
+    public void addDefineConfig(DefineConfig defineConfig) {
+        this.defineConfigs.add(defineConfig);
+    }
+
 }
diff --git a/bytekit-core/src/main/java/com/alibaba/bytekit/asm/instrument/InstrumentParseResult.java b/bytekit-core/src/main/java/com/alibaba/bytekit/asm/instrument/InstrumentParseResult.java
@@ -12,8 +12,16 @@ public class InstrumentParseResult {
 
     private List<InstrumentConfig> instrumentConfigs = new ArrayList<InstrumentConfig>();
 
+    /**
+     * @deprecated 已废弃，define 配置现在关联到具体的 InstrumentConfig 上。为了保持向后兼容暂时保留。
+     */
+    @Deprecated
     private List<DefineConfig> defineConfigs = new ArrayList<DefineConfig>();
 
+    /**
+     * @deprecated 已废弃，建议使用关联到 InstrumentConfig 的 define 配置
+     */
+    @Deprecated
     public void addDefineClass(String className, byte[] classBytes) {
         this.defineConfigs.add(new DefineConfig(classBytes, className));
     }
@@ -30,10 +38,18 @@ public void setInstrumentConfigs(List<InstrumentConfig> instrumentConfigs) {
         this.instrumentConfigs = instrumentConfigs;
     }
 
+    /**
+     * @deprecated 已废弃，define 配置现在关联到具体的 InstrumentConfig 上
+     */
+    @Deprecated
     public List<DefineConfig> getDefineConfigs() {
         return defineConfigs;
     }
 
+    /**
+     * @deprecated 已废弃，define 配置现在关联到具体的 InstrumentConfig 上
+     */
+    @Deprecated
     public void setDefineConfigs(List<DefineConfig> defineConfigs) {
         this.defineConfigs = defineConfigs;
     }
diff --git a/bytekit-core/src/main/java/com/alibaba/bytekit/asm/instrument/InstrumentTemplate.java b/bytekit-core/src/main/java/com/alibaba/bytekit/asm/instrument/InstrumentTemplate.java
@@ -6,7 +6,9 @@
 import java.lang.instrument.Instrumentation;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Properties;
 import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
@@ -85,14 +87,47 @@ public InstrumentParseResult build() throws IOException {
                     InputStream inputStream = jarFile.getInputStream(propertiesEntry);
                     Properties properties = PropertiesUtils.loadNotNull(inputStream);
 
-                    String triggerRetransformValue = properties.getProperty("triggerRetransform", "false"); // 使用默认值避免null值
+                    String triggerRetransformValue = properties.getProperty("triggerRetransform", "false");
                     boolean triggerRetransform = Boolean.parseBoolean(triggerRetransformValue);
-                    for (Pair<String, byte[]> pair : readClassBytes(properties, INSTRUMENT, jarFile)) {
-                        parse(result, pair.second, triggerRetransform);
+                    
+                    // 解析新格式：instrument.{key}=... 和 define.{key}=...
+                    Map<String, List<Pair<String, byte[]>>> instrumentMap = new HashMap<String, List<Pair<String, byte[]>>>();
+                    Map<String, List<Pair<String, byte[]>>> defineMap = new HashMap<String, List<Pair<String, byte[]>>>();
+                    
+                    for (String key : properties.stringPropertyNames()) {
+                        if (key.startsWith("instrument.")) {
+                            String configKey = key.substring("instrument.".length());
+                            List<Pair<String, byte[]>> classBytesList = readClassBytes(properties, key, jarFile);
+                            if (!classBytesList.isEmpty()) {
+                                instrumentMap.put(configKey, classBytesList);
+                            }
+                        } else if (key.startsWith("define.")) {
+                            String configKey = key.substring("define.".length());
+                            List<Pair<String, byte[]>> classBytesList = readClassBytes(properties, key, jarFile);
+                            if (!classBytesList.isEmpty()) {
+                                defineMap.put(configKey, classBytesList);
+                            }
+                        } else if (key.equals(INSTRUMENT)) {
+                            // 兼容旧格式：instrument=...
+                            for (Pair<String, byte[]> pair : readClassBytes(properties, INSTRUMENT, jarFile)) {
+                                parse(result, pair.second, triggerRetransform, null);
+                            }
+                        } else if (key.equals(DEFINE)) {
+                            // 兼容旧格式：define=...（全局 define，会被添加到所有 instrument）
+                            for (Pair<String, byte[]> pair : readClassBytes(properties, DEFINE, jarFile)) {
+                                result.addDefineClass(pair.first, pair.second);
+                            }
+                        }
                     }
-
-                    for (Pair<String, byte[]> pair : readClassBytes(properties, DEFINE, jarFile)) {
-                        result.addDefineClass(pair.first, pair.second);
+                    
+                    // 处理新格式的 instrument 和 define 关联
+                    for (Map.Entry<String, List<Pair<String, byte[]>>> entry : instrumentMap.entrySet()) {
+                        String configKey = entry.getKey();
+                        List<Pair<String, byte[]>> defineClassList = defineMap.get(configKey);
+                        
+                        for (Pair<String, byte[]> instrumentPair : entry.getValue()) {
+                            parse(result, instrumentPair.second, triggerRetransform, defineClassList);
+                        }
                     }
                 }
 
@@ -103,7 +138,7 @@ public InstrumentParseResult build() throws IOException {
 
         // 处理单独设置 byte[]
         for (byte[] classBytes : instrumentClassList) {
-            parse(result, classBytes, false);
+            parse(result, classBytes, false, null);
         }
 
         return result;
@@ -137,7 +172,7 @@ private List<Pair<String, byte[]>> readClassBytes(Properties properties, String
         return result;
     }
 
-    private void parse(InstrumentParseResult result, byte[] classBytes, boolean triggerRetransform) {
+    private void parse(InstrumentParseResult result, byte[] classBytes, boolean triggerRetransform, List<Pair<String, byte[]>> defineClassList) {
         ClassNode classNode = AsmUtils.toClassNode(classBytes);
 
         if (!AsmUtils.fitCurrentJvmMajorVersion(classNode)) {
@@ -158,28 +193,45 @@ private void parse(InstrumentParseResult result, byte[] classBytes, boolean trig
 
         if (matchClassList != null && !matchClassList.isEmpty()) {
             SimpleClassMatcher classMatcher = new SimpleClassMatcher(matchClassList);
-            result.addInstrumentConfig(new InstrumentConfig(classNode, classMatcher, updateMajorVersion, triggerRetransform));
+            InstrumentConfig config = new InstrumentConfig(classNode, classMatcher, updateMajorVersion, triggerRetransform);
+            addDefineConfigs(config, defineClassList);
+            result.addInstrumentConfig(config);
         }
 
         List<String> matchSuperclassList = AsmAnnotationUtils.queryAnnotationArrayValue(classNode.visibleAnnotations,
                 Type.getDescriptor(Instrument.class), "Superclass");
 
         if (!matchSuperclassList.isEmpty()) {
             SimpleSubclassMatcher matcher = new SimpleSubclassMatcher(matchSuperclassList);
-            result.addInstrumentConfig(new InstrumentConfig(classNode, matcher, updateMajorVersion, triggerRetransform));
+            InstrumentConfig config = new InstrumentConfig(classNode, matcher, updateMajorVersion, triggerRetransform);
+            addDefineConfigs(config, defineClassList);
+            result.addInstrumentConfig(config);
         }
 
         List<String> matchInterfaceList = AsmAnnotationUtils.queryAnnotationArrayValue(classNode.visibleAnnotations,
                 Type.getDescriptor(Instrument.class), "Interface");
 
         if (!matchInterfaceList.isEmpty()) {
             SimpleInterfaceMatcher matcher = new SimpleInterfaceMatcher(matchInterfaceList);
-            result.addInstrumentConfig(new InstrumentConfig(classNode, matcher, updateMajorVersion, triggerRetransform));
+            InstrumentConfig config = new InstrumentConfig(classNode, matcher, updateMajorVersion, triggerRetransform);
+            addDefineConfigs(config, defineClassList);
+            result.addInstrumentConfig(config);
         }
 
         // TODO 处理 @NewField
     }
 
+    /**
+     * 将 define 类列表添加到 InstrumentConfig
+     */
+    private void addDefineConfigs(InstrumentConfig config, List<Pair<String, byte[]>> defineClassList) {
+        if (defineClassList != null && !defineClassList.isEmpty()) {
+            for (Pair<String, byte[]> pair : defineClassList) {
+                config.addDefineConfig(new DefineConfig(pair.second, pair.first));
+            }
+        }
+    }
+
     public static List<Class<?>> matchedClass(Instrumentation instrumentation, InstrumentConfig instrumentConfig) {
         List<Class<?>> result = new ArrayList<Class<?>>();
 
diff --git a/bytekit-core/src/main/java/com/alibaba/bytekit/asm/instrument/InstrumentTransformer.java b/bytekit-core/src/main/java/com/alibaba/bytekit/asm/instrument/InstrumentTransformer.java
@@ -3,6 +3,7 @@
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.IllegalClassFormatException;
 import java.security.ProtectionDomain;
+import java.util.ArrayList;
 import java.util.List;
 
 import com.alibaba.bytekit.asm.inst.impl.InstrumentImpl;
@@ -97,8 +98,23 @@ public byte[] transform(ClassLoader loader, String className, Class<?> classBein
         if (targetClassNode != null) {
             // TODO 支持 bootstrap classloader?
             if (loader != null) {
-                List<DefineConfig> defineConfigs = instrumentParseResult.getDefineConfigs();
-                for (DefineConfig defineConfig : defineConfigs) {
+                // 收集所有匹配的 instrument config 关联的 define 类
+                List<DefineConfig> allDefineConfigs = new ArrayList<DefineConfig>();
+                for (InstrumentConfig config : instrumentConfigs) {
+                    if (config.getClassMatcher().match(loader, className, classBeingRedefined, protectionDomain,
+                            classfileBuffer)) {
+                        allDefineConfigs.addAll(config.getDefineConfigs());
+                    }
+                }
+                
+                // 兼容旧格式：添加全局的 define 配置（已废弃）
+                List<DefineConfig> globalDefineConfigs = instrumentParseResult.getDefineConfigs();
+                if (globalDefineConfigs != null && !globalDefineConfigs.isEmpty()) {
+                    allDefineConfigs.addAll(globalDefineConfigs);
+                }
+                
+                // 定义所有关联的类
+                for (DefineConfig defineConfig : allDefineConfigs) {
                     try {
                         ReflectUtils.defineClass(defineConfig.getClassName(), defineConfig.getClassBytes(), loader);
                     } catch (Throwable e) {
diff --git a/bytekit-core/src/test/java/com/alibaba/bytekit/asm/instrument/InstrumentDefineConfigTest.java b/bytekit-core/src/test/java/com/alibaba/bytekit/asm/instrument/InstrumentDefineConfigTest.java
diff --git a/bytekit-core/src/test/resources/instrument.properties b/bytekit-core/src/test/resources/instrument.properties
