test: support advisor and asm agent unit test

2026-09-21 22:50:42 +08:00 · 2025-03-30 15:33:31 +08:00
parent 3d2f65a808
commit e3a59e9f21
8 changed files with 285 additions and 6 deletions
@@ -13,6 +13,20 @@ java {
 group = 'com.reajason.javaweb.memsell'
 version = rootProject.version

+tasks.withType(Test).configureEach {
+    javaLauncher = javaToolchains.launcherFor {
+        languageVersion = JavaLanguageVersion.of(17)
+    }
+}
+
+tasks.named('compileTestJava') {
+    javaCompiler = javaToolchains.compilerFor {
+        languageVersion = JavaLanguageVersion.of(17)
+    }
+    sourceCompatibility = JavaVersion.VERSION_17
+    targetCompatibility = JavaVersion.VERSION_17
+}
+
 test {
    useJUnitPlatform()
    finalizedBy jacocoTestReport
@@ -30,7 +44,6 @@ dependencies {
    implementation 'javax.websocket:javax.websocket-api'
    implementation 'jakarta.servlet:jakarta.servlet-api'

-//    implementation 'xalan:xalan'
    implementation 'org.apache.bcel:bcel'

    implementation 'commons-io:commons-io'
@@ -39,7 +52,6 @@ dependencies {
    implementation 'com.squareup.okhttp3:okhttp'
    implementation 'ch.qos.logback:logback-classic'
    implementation 'com.alibaba.fastjson2:fastjson2'
-//    implementation 'org.java-websocket:Java-WebSocket'

    implementation 'org.springframework:spring-webmvc'
    implementation 'org.springframework:spring-webflux'
@@ -49,4 +61,5 @@ dependencies {
    testImplementation 'org.junit.jupiter:junit-jupiter'
    testRuntimeOnly 'org.junit.platform:junit-platform-launcher'
    testImplementation "org.mockito:mockito-core"
+    testImplementation 'org.mockito:mockito-junit-jupiter'
 }
@@ -0,0 +1,28 @@
+package com.reajason.javaweb.memshell.shelltool;
+
+import javax.servlet.ServletOutputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+
+public class DelegatingServletOutputStream extends ServletOutputStream {
+    private final ByteArrayOutputStream delegate;
+
+    public DelegatingServletOutputStream(ByteArrayOutputStream delegate) {
+        this.delegate = delegate;
+    }
+
+    @Override
+    public void write(int b) throws IOException {
+        delegate.write(b);
+    }
+
+    @Override
+    public void write(byte[] b) throws IOException {
+        delegate.write(b);
+    }
+
+    @Override
+    public void write(byte[] b, int off, int len) throws IOException {
+        delegate.write(b, off, len);
+    }
+}
@@ -0,0 +1,17 @@
+package com.reajason.javaweb.memshell.shelltool;
+
+import javax.servlet.FilterChain;
+import javax.servlet.ServletException;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import java.io.IOException;
+
+/**
+ * @author ReaJason
+ * @since 2025/3/30
+ */
+public interface FilterChainInterface {
+    void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException;
+
+    void doFilterInternal();
+}
@@ -0,0 +1,22 @@
+package com.reajason.javaweb.memshell.shelltool;
+
+import javax.servlet.FilterChain;
+import javax.servlet.ServletException;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import java.io.IOException;
+
+/**
+ * @author ReaJason
+ * @since 2025/3/30
+ */
+public class TestFilterChain implements FilterChainInterface {
+
+    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
+        doFilterInternal();
+    }
+
+    public void doFilterInternal() {
+        System.out.println("doFilterInternal");
+    }
+}
@@ -0,0 +1,106 @@
+package com.reajason.javaweb.memshell.shelltool.command;
+
+import com.reajason.javaweb.asm.ClassRenameUtils;
+import com.reajason.javaweb.memshell.shelltool.DelegatingServletOutputStream;
+import com.reajason.javaweb.memshell.shelltool.FilterChainInterface;
+import com.reajason.javaweb.memshell.shelltool.TestFilterChain;
+import com.reajason.javaweb.util.ClassUtils;
+import lombok.SneakyThrows;
+import me.n1ar4.clazz.obfuscator.api.ClassObf;
+import me.n1ar4.clazz.obfuscator.api.Result;
+import me.n1ar4.clazz.obfuscator.config.BaseConfig;
+import org.apache.commons.io.IOUtils;
+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 org.objectweb.asm.*;
+
+import javax.servlet.FilterChain;
+import javax.servlet.ServletOutputStream;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import java.io.ByteArrayOutputStream;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.nio.charset.StandardCharsets;
+import java.util.Objects;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.*;
+
+/**
+ * @author ReaJason
+ * @since 2025/3/30
+ */
+@ExtendWith(MockitoExtension.class)
+public class CommandFilterChainASMTest {
+
+    @Mock
+    ServletRequest mockRequest;
+
+    @Mock
+    ServletResponse mockResponse;
+
+    Object instance;
+
+    @BeforeEach
+    @SneakyThrows
+    void setUp() {
+        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)) {
+                    Type[] argumentTypes = Type.getArgumentTypes(descriptor);
+                    return new CommandFilterChainAsmMethodVisitor(mv, argumentTypes);
+                }
+                return mv;
+            }
+        };
+        cr.accept(cv, ClassReader.EXPAND_FRAMES);
+        byte[] bytes2 = ClassRenameUtils.renameClass(cw.toByteArray(), TestFilterChain.class.getName() + "Asm");
+        BaseConfig config = BaseConfig.Default();
+        config.setIgnorePublic(true);
+        config.setEnableMethodName(false);
+        config.setEnableParamName(false);
+        config.setEnableAES(false);
+        config.setEnableAdvanceString(false);
+        ClassObf classObf = new ClassObf(config);
+        Result run = classObf.run(bytes2);
+        bytes2 = run.getData();
+        IOUtils.write(bytes2, new FileOutputStream(new File("godzilla2.class")));
+        Class<?> clazz = ClassUtils.defineClass(bytes2);
+        instance = spy(clazz.newInstance());
+    }
+
+    @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,93 @@
+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();
+    }
+}