fix: custom shell not supported for listener and valve vistor

This commit is contained in:
ReaJason
2025-08-31 17:39:57 +08:00
parent 81e9068a89
commit cc5f3a18c8
10 changed files with 178 additions and 35 deletions
@@ -1,17 +1,21 @@
package com.reajason.javaweb.memshell.generator;
import com.reajason.javaweb.Server;
import com.reajason.javaweb.asm.ClassReferenceVisitor;
import com.reajason.javaweb.memshell.ShellType;
import com.reajason.javaweb.memshell.config.CustomConfig;
import com.reajason.javaweb.memshell.config.ShellConfig;
import com.reajason.javaweb.memshell.shelltool.godzilla.GodzillaValve;
import com.reajason.javaweb.utils.CommonUtil;
import lombok.SneakyThrows;
import net.bytebuddy.ByteBuddy;
import net.bytebuddy.jar.asm.ClassReader;
import org.junit.jupiter.api.Test;
import org.objectweb.asm.ClassReader;
import java.util.Base64;
import java.util.Set;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.*;
/**
@@ -27,9 +31,42 @@ class CustomShellGeneratorTest {
.subclass(Object.class)
.name(CommonUtil.generateShellClassName()).make().getBytes();
String className = CommonUtil.generateShellClassName();
byte[] bytes1 = new CustomShellGenerator(ShellConfig.builder().build(), CustomConfig.builder().shellClassName(className).shellClassBase64(Base64.getEncoder().encodeToString(bytes)).build()).getBytes();
ShellConfig shellConfig = ShellConfig.builder()
.shellType(ShellType.FILTER)
.build();
CustomConfig customConfig = CustomConfig.builder()
.shellClassName(className)
.shellClassBase64(Base64.getEncoder().encodeToString(bytes))
.build();
byte[] bytes1 = new CustomShellGenerator(shellConfig, customConfig).getBytes();
ClassReader classReader = new ClassReader(bytes1);
assertEquals(className, classReader.getClassName().replace("/", "."));
}
@Test
@SneakyThrows
void testValue() {
byte[] bytes = new ByteBuddy()
.redefine(GodzillaValve.class)
.name(CommonUtil.generateShellClassName()).make().getBytes();
String className = CommonUtil.generateShellClassName();
ShellConfig shellConfig = ShellConfig.builder()
.server(Server.BES)
.shellType(ShellType.VALVE)
.build();
CustomConfig customConfig = CustomConfig.builder()
.shellClassName(className)
.shellClassBase64(Base64.getEncoder().encodeToString(bytes))
.build();
byte[] bytes1 = new CustomShellGenerator(shellConfig, customConfig).getBytes();
ClassReader classReader = new ClassReader(bytes1);
ClassReferenceVisitor classVisitor = new ClassReferenceVisitor();
classReader.accept(classVisitor, 0);
Set<String> referencedClasses = classVisitor.getReferencedClasses();
assertEquals(className, classReader.getClassName().replace("/", "."));
assertTrue(referencedClasses.contains("com/bes/enterprise/webtier/Valve"));
assertFalse(referencedClasses.contains("org/apache/catalina/Valve"));
}
}