Skip to content

Commit f49c4e2

Browse files
ReaJasonReaJason
committed
test: support advisor and asm agent unit test
1 parent cc6d79a commit f49c4e2

File tree

8 files changed

+285
-6
lines changed

8 files changed

+285
-6
lines changed

‎bom/build.gradle‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ dependencies {
2020
api 'commons-codec:commons-codec:1.+'
2121
api 'ch.qos.logback:logback-classic:1.+'
2222

23-
api 'xalan:xalan:2.7.0'
2423
api 'org.apache.bcel:bcel:5.2'
2524

2625
api 'org.java-websocket:Java-WebSocket:1.5.7'
@@ -29,7 +28,8 @@ dependencies {
2928

3029
api 'org.jetbrains:annotations:26.0.1'
3130

32-
api "org.mockito:mockito-core:5.15.2"
31+
api 'org.mockito:mockito-core:5.15.2'
32+
api 'org.mockito:mockito-junit-jupiter:5.15.2'
3333
api 'org.hamcrest:hamcrest:3.0'
3434
api 'org.junit:junit-bom:5.11.4'
3535
api 'org.testcontainers:testcontainers:1.20.5'

‎generator/build.gradle‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,20 @@ java {
1313
group = 'com.reajason.javaweb.memsell'
1414
version = rootProject.version
1515

16+
tasks.withType(Test).configureEach {
17+
javaLauncher = javaToolchains.launcherFor {
18+
languageVersion = JavaLanguageVersion.of(17)
19+
}
20+
}
21+
22+
tasks.named('compileTestJava') {
23+
javaCompiler = javaToolchains.compilerFor {
24+
languageVersion = JavaLanguageVersion.of(17)
25+
}
26+
sourceCompatibility = JavaVersion.VERSION_17
27+
targetCompatibility = JavaVersion.VERSION_17
28+
}
29+
1630
test {
1731
useJUnitPlatform()
1832
finalizedBy jacocoTestReport
@@ -30,7 +44,6 @@ dependencies {
3044
implementation 'javax.websocket:javax.websocket-api'
3145
implementation 'jakarta.servlet:jakarta.servlet-api'
3246

33-
// implementation 'xalan:xalan'
3447
implementation 'org.apache.bcel:bcel'
3548

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

4456
implementation 'org.springframework:spring-webmvc'
4557
implementation 'org.springframework:spring-webflux'
@@ -49,4 +61,5 @@ dependencies {
4961
testImplementation 'org.junit.jupiter:junit-jupiter'
5062
testRuntimeOnly 'org.junit.platform:junit-platform-launcher'
5163
testImplementation "org.mockito:mockito-core"
64+
testImplementation 'org.mockito:mockito-junit-jupiter'
5265
}

‎generator/src/test/java/com/reajason/javaweb/memshell/shelltool/DelegatingServletOutputStream.java‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
package com.reajason.javaweb.memshell.shelltool;
2+
3+
import javax.servlet.ServletOutputStream;
4+
import java.io.ByteArrayOutputStream;
5+
import java.io.IOException;
6+
7+
public class DelegatingServletOutputStream extends ServletOutputStream {
8+
private final ByteArrayOutputStream delegate;
9+
10+
public DelegatingServletOutputStream(ByteArrayOutputStream delegate) {
11+
this.delegate = delegate;
12+
}
13+
14+
@Override
15+
public void write(int b) throws IOException {
16+
delegate.write(b);
17+
}
18+
19+
@Override
20+
public void write(byte[] b) throws IOException {
21+
delegate.write(b);
22+
}
23+
24+
@Override
25+
public void write(byte[] b, int off, int len) throws IOException {
26+
delegate.write(b, off, len);
27+
}
28+
}

‎generator/src/test/java/com/reajason/javaweb/memshell/shelltool/FilterChainInterface.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package com.reajason.javaweb.memshell.shelltool;
2+
3+
import javax.servlet.FilterChain;
4+
import javax.servlet.ServletException;
5+
import javax.servlet.ServletRequest;
6+
import javax.servlet.ServletResponse;
7+
import java.io.IOException;
8+
9+
/**
10+
* @author ReaJason
11+
* @since 2025/3/30
12+
*/
13+
public interface FilterChainInterface {
14+
void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException;
15+
16+
void doFilterInternal();
17+
}

‎generator/src/test/java/com/reajason/javaweb/memshell/shelltool/TestFilterChain.java‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
package com.reajason.javaweb.memshell.shelltool;
2+
3+
import javax.servlet.FilterChain;
4+
import javax.servlet.ServletException;
5+
import javax.servlet.ServletRequest;
6+
import javax.servlet.ServletResponse;
7+
import java.io.IOException;
8+
9+
/**
10+
* @author ReaJason
11+
* @since 2025/3/30
12+
*/
13+
public class TestFilterChain implements FilterChainInterface {
14+
15+
public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
16+
doFilterInternal();
17+
}
18+
19+
public void doFilterInternal() {
20+
System.out.println("doFilterInternal");
21+
}
22+
}

‎generator/src/test/java/com/reajason/javaweb/memshell/shelltool/command/CommandFilterChainASMTest.java‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
package com.reajason.javaweb.memshell.shelltool.command;
2+
3+
import com.reajason.javaweb.asm.ClassRenameUtils;
4+
import com.reajason.javaweb.memshell.shelltool.DelegatingServletOutputStream;
5+
import com.reajason.javaweb.memshell.shelltool.FilterChainInterface;
6+
import com.reajason.javaweb.memshell.shelltool.TestFilterChain;
7+
import com.reajason.javaweb.util.ClassUtils;
8+
import lombok.SneakyThrows;
9+
import me.n1ar4.clazz.obfuscator.api.ClassObf;
10+
import me.n1ar4.clazz.obfuscator.api.Result;
11+
import me.n1ar4.clazz.obfuscator.config.BaseConfig;
12+
import org.apache.commons.io.IOUtils;
13+
import org.junit.jupiter.api.BeforeEach;
14+
import org.junit.jupiter.api.Test;
15+
import org.junit.jupiter.api.extension.ExtendWith;
16+
import org.mockito.Mock;
17+
import org.mockito.junit.jupiter.MockitoExtension;
18+
import org.objectweb.asm.*;
19+
20+
import javax.servlet.FilterChain;
21+
import javax.servlet.ServletOutputStream;
22+
import javax.servlet.ServletRequest;
23+
import javax.servlet.ServletResponse;
24+
import java.io.ByteArrayOutputStream;
25+
import java.io.File;
26+
import java.io.FileOutputStream;
27+
import java.nio.charset.StandardCharsets;
28+
import java.util.Objects;
29+
30+
import static org.junit.jupiter.api.Assertions.assertTrue;
31+
import static org.mockito.Mockito.*;
32+
33+
/**
34+
* @author ReaJason
35+
* @since 2025/3/30
36+
*/
37+
@ExtendWith(MockitoExtension.class)
38+
public class CommandFilterChainASMTest {
39+
40+
@Mock
41+
ServletRequest mockRequest;
42+
43+
@Mock
44+
ServletResponse mockResponse;
45+
46+
Object instance;
47+
48+
@BeforeEach
49+
@SneakyThrows
50+
void setUp() {
51+
byte[] bytes = IOUtils.toByteArray(Objects.requireNonNull(TestFilterChain.class.getClassLoader().getResource(TestFilterChain.class.getName().replace('.', '/') + ".class")));
52+
ClassReader cr = new ClassReader(bytes);
53+
ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES);
54+
ClassVisitor cv = new ClassVisitor(Opcodes.ASM9, cw) {
55+
@Override
56+
public MethodVisitor visitMethod(int access, String name, String descriptor,
57+
String signature, String[] exceptions) {
58+
MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
59+
if ("doFilter".equals(name)) {
60+
Type[] argumentTypes = Type.getArgumentTypes(descriptor);
61+
return new CommandFilterChainAsmMethodVisitor(mv, argumentTypes);
62+
}
63+
return mv;
64+
}
65+
};
66+
cr.accept(cv, ClassReader.EXPAND_FRAMES);
67+
byte[] bytes2 = ClassRenameUtils.renameClass(cw.toByteArray(), TestFilterChain.class.getName() + "Asm");
68+
BaseConfig config = BaseConfig.Default();
69+
config.setIgnorePublic(true);
70+
config.setEnableMethodName(false);
71+
config.setEnableParamName(false);
72+
config.setEnableAES(false);
73+
config.setEnableAdvanceString(false);
74+
ClassObf classObf = new ClassObf(config);
75+
Result run = classObf.run(bytes2);
76+
bytes2 = run.getData();
77+
IOUtils.write(bytes2, new FileOutputStream(new File("godzilla2.class")));
78+
Class<?> clazz = ClassUtils.defineClass(bytes2);
79+
instance = spy(clazz.newInstance());
80+
}
81+
82+
@Test
83+
@SneakyThrows
84+
void testInvokeParam() {
85+
when(mockRequest.getParameter("paramName")).thenReturn("id");
86+
ByteArrayOutputStream capturedOutput = new ByteArrayOutputStream();
87+
ServletOutputStream servletOutputStream = new DelegatingServletOutputStream(capturedOutput);
88+
when(mockResponse.getOutputStream()).thenReturn(servletOutputStream);
89+
90+
instance.getClass().getMethod("doFilter", ServletRequest.class, ServletResponse.class, FilterChain.class).invoke(instance, mockRequest, mockResponse, null);
91+
String output = capturedOutput.toString(StandardCharsets.UTF_8);
92+
assertTrue(output.contains("uid="));
93+
94+
verify(((FilterChainInterface) instance), never()).doFilterInternal();
95+
}
96+
97+
@Test
98+
@SneakyThrows
99+
void testNotParameter() {
100+
when(mockRequest.getParameter("paramName")).thenReturn(null);
101+
102+
instance.getClass().getMethod("doFilter", ServletRequest.class, ServletResponse.class, FilterChain.class).invoke(instance, mockRequest, mockResponse, null);
103+
104+
verify(((FilterChainInterface) instance), atLeastOnce()).doFilterInternal();
105+
}
106+
}

‎generator/src/test/java/com/reajason/javaweb/memshell/shelltool/command/CommandFilterChainAdvisorTest.java‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
package com.reajason.javaweb.memshell.shelltool.command;
2+
3+
import com.reajason.javaweb.memshell.shelltool.DelegatingServletOutputStream;
4+
import com.reajason.javaweb.memshell.shelltool.FilterChainInterface;
5+
import com.reajason.javaweb.memshell.shelltool.TestFilterChain;
6+
import lombok.SneakyThrows;
7+
import net.bytebuddy.agent.ByteBuddyAgent;
8+
import net.bytebuddy.agent.builder.AgentBuilder;
9+
import net.bytebuddy.asm.Advice;
10+
import net.bytebuddy.dynamic.ClassFileLocator;
11+
import net.bytebuddy.dynamic.loading.ByteArrayClassLoader;
12+
import net.bytebuddy.matcher.ElementMatchers;
13+
import org.junit.jupiter.api.AfterEach;
14+
import org.junit.jupiter.api.BeforeEach;
15+
import org.junit.jupiter.api.Test;
16+
import org.junit.jupiter.api.extension.ExtendWith;
17+
import org.mockito.Mock;
18+
import org.mockito.junit.jupiter.MockitoExtension;
19+
20+
import javax.servlet.FilterChain;
21+
import javax.servlet.ServletOutputStream;
22+
import javax.servlet.ServletRequest;
23+
import javax.servlet.ServletResponse;
24+
import java.io.ByteArrayOutputStream;
25+
import java.lang.instrument.ClassFileTransformer;
26+
import java.nio.charset.StandardCharsets;
27+
28+
import static net.bytebuddy.matcher.ElementMatchers.none;
29+
import static org.junit.jupiter.api.Assertions.assertTrue;
30+
import static org.mockito.Mockito.*;
31+
32+
/**
33+
* @author ReaJason
34+
* @since 2025/3/30
35+
*/
36+
@ExtendWith(MockitoExtension.class)
37+
public class CommandFilterChainAdvisorTest {
38+
39+
@Mock
40+
ServletRequest mockRequest;
41+
42+
@Mock
43+
ServletResponse mockResponse;
44+
45+
Object instance;
46+
47+
static ClassFileTransformer classFileTransformer;
48+
49+
@BeforeEach
50+
@SneakyThrows
51+
void setUp() {
52+
ByteBuddyAgent.install();
53+
ClassLoader classLoader = new ByteArrayClassLoader.ChildFirst(CommandFilterChainAdvisorTest.class.getClassLoader(),
54+
ClassFileLocator.ForClassLoader.readToNames(TestFilterChain.class),
55+
ByteArrayClassLoader.PersistenceHandler.MANIFEST);
56+
classFileTransformer = new AgentBuilder.Default()
57+
.ignore(none())
58+
.type(ElementMatchers.is(TestFilterChain.class), ElementMatchers.is(classLoader)).transform((
59+
(builder, typeDescription, c, module, protectionDomain) ->
60+
builder.visit(Advice.to(CommandFilterChainAdvisor.class).on(ElementMatchers.named("doFilter")))))
61+
.installOnByteBuddyAgent();
62+
Class<?> clazz = classLoader.loadClass(TestFilterChain.class.getName());
63+
instance = spy(clazz.newInstance());
64+
}
65+
66+
@AfterEach
67+
void tearDown() {
68+
ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
69+
}
70+
71+
@Test
72+
@SneakyThrows
73+
void testInvokeParam() {
74+
when(mockRequest.getParameter("paramName")).thenReturn("id");
75+
ByteArrayOutputStream capturedOutput = new ByteArrayOutputStream();
76+
ServletOutputStream servletOutputStream = new DelegatingServletOutputStream(capturedOutput);
77+
when(mockResponse.getOutputStream()).thenReturn(servletOutputStream);
78+
79+
instance.getClass().getMethod("doFilter", ServletRequest.class, ServletResponse.class, FilterChain.class).invoke(instance, mockRequest, mockResponse, null);
80+
String output = capturedOutput.toString(StandardCharsets.UTF_8);
81+
assertTrue(output.contains("uid="));
82+
83+
verify(((FilterChainInterface) instance), never()).doFilterInternal();
84+
}
85+
86+
@Test
87+
@SneakyThrows
88+
void testNotParameter() {
89+
when(mockRequest.getParameter("paramName")).thenReturn(null);
90+
instance.getClass().getMethod("doFilter", ServletRequest.class, ServletResponse.class, FilterChain.class).invoke(instance, mockRequest, mockResponse, null);
91+
verify(((FilterChainInterface) instance), atLeastOnce()).doFilterInternal();
92+
}
93+
}

‎integration-test/src/test/java/com/reajason/javaweb/integration/jetty/Jetty92ContainerTest.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,8 @@ static Stream<Arguments> casesProvider() {
5454
Server server = Server.Jetty;
5555
List<String> supportedShellTypes = List.of(
5656
ShellType.SERVLET, ShellType.FILTER, ShellType.LISTENER,
57-
ShellType.JETTY_AGENT_HANDLER
58-
// ShellType.JETTY_AGENT_HANDLER_ASM // 内置 ASM，但是版本较低 5.0.1, API 不兼容
57+
ShellType.JETTY_AGENT_HANDLER,
58+
ShellType.JETTY_AGENT_HANDLER_ASM
5959
);
6060
List<Packers> testPackers = List.of(Packers.JSP, Packers.JSPX, Packers.JavaDeserialize);
6161
return TestCasesProvider.getTestCases(imageName, server, supportedShellTypes, testPackers);

0 commit comments

Comments
 (0)