From d7eaf5c06880a1ef68261399480b70d256211841 Mon Sep 17 00:00:00 2001 From: ReaJason Date: Sun, 24 Aug 2025 14:55:19 +0800 Subject: [PATCH] fix: custom shell not supported for listener and valve vistor --- .../memshell/generator/AntSwordGenerator.java | 7 +- .../memshell/generator/BehinderGenerator.java | 8 +- .../generator/ByteBuddyShellGenerator.java | 12 +-- .../generator/CustomShellGenerator.java | 29 ++++--- .../memshell/generator/GodzillaGenerator.java | 7 +- .../generator/NeoreGeorgGenerator.java | 7 +- .../memshell/generator/Suo5Generator.java | 6 +- .../generator/command/CommandGenerator.java | 9 +- .../generator/CustomShellGeneratorTest.java | 45 +++++++++- .../javaweb/asm/ClassReferenceVisitor.java | 83 +++++++++++++++++++ 10 files changed, 178 insertions(+), 35 deletions(-) create mode 100644 memshell-party-common/src/main/java/com/reajason/javaweb/asm/ClassReferenceVisitor.java diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/generator/AntSwordGenerator.java b/generator/src/main/java/com/reajason/javaweb/memshell/generator/AntSwordGenerator.java index e5b80c77..039477e1 100644 --- a/generator/src/main/java/com/reajason/javaweb/memshell/generator/AntSwordGenerator.java +++ b/generator/src/main/java/com/reajason/javaweb/memshell/generator/AntSwordGenerator.java @@ -2,6 +2,7 @@ package com.reajason.javaweb.memshell.generator; import com.reajason.javaweb.memshell.config.AntSwordConfig; import com.reajason.javaweb.memshell.config.ShellConfig; +import net.bytebuddy.ByteBuddy; import net.bytebuddy.dynamic.DynamicType; import static net.bytebuddy.matcher.ElementMatchers.named; @@ -16,8 +17,10 @@ public class AntSwordGenerator extends ByteBuddyShellGenerator { } @Override - protected DynamicType.Builder build(DynamicType.Builder builder) { - return builder.field(named("pass")).value(shellToolConfig.getPass()) + protected DynamicType.Builder getBuilder() { + return new ByteBuddy() + .redefine(shellToolConfig.getShellClass()) + .field(named("pass")).value(shellToolConfig.getPass()) .field(named("headerName")).value(shellToolConfig.getHeaderName()) .field(named("headerValue")).value(shellToolConfig.getHeaderValue()); } diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/generator/BehinderGenerator.java b/generator/src/main/java/com/reajason/javaweb/memshell/generator/BehinderGenerator.java index bc966d5c..a3537ba0 100644 --- a/generator/src/main/java/com/reajason/javaweb/memshell/generator/BehinderGenerator.java +++ b/generator/src/main/java/com/reajason/javaweb/memshell/generator/BehinderGenerator.java @@ -2,6 +2,7 @@ package com.reajason.javaweb.memshell.generator; import com.reajason.javaweb.memshell.config.BehinderConfig; import com.reajason.javaweb.memshell.config.ShellConfig; +import net.bytebuddy.ByteBuddy; import net.bytebuddy.dynamic.DynamicType; import org.apache.commons.codec.digest.DigestUtils; @@ -16,9 +17,12 @@ public class BehinderGenerator extends ByteBuddyShellGenerator { super(shellConfig, shellToolConfig); } - public DynamicType.Builder build(DynamicType.Builder builder) { + @Override + public DynamicType.Builder getBuilder() { String md5Key = DigestUtils.md5Hex(shellToolConfig.getPass()).substring(0, 16); - return builder.field(named("pass")).value(md5Key) + return new ByteBuddy() + .redefine(shellToolConfig.getShellClass()) + .field(named("pass")).value(md5Key) .field(named("headerName")).value(shellToolConfig.getHeaderName()) .field(named("headerValue")).value(shellToolConfig.getHeaderValue()); } diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/generator/ByteBuddyShellGenerator.java b/generator/src/main/java/com/reajason/javaweb/memshell/generator/ByteBuddyShellGenerator.java index 8bedfe77..a6393ff3 100644 --- a/generator/src/main/java/com/reajason/javaweb/memshell/generator/ByteBuddyShellGenerator.java +++ b/generator/src/main/java/com/reajason/javaweb/memshell/generator/ByteBuddyShellGenerator.java @@ -10,7 +10,6 @@ import com.reajason.javaweb.memshell.ShellType; import com.reajason.javaweb.memshell.config.ShellConfig; import com.reajason.javaweb.memshell.config.ShellToolConfig; import com.reajason.javaweb.memshell.server.AbstractServer; -import net.bytebuddy.ByteBuddy; import net.bytebuddy.dynamic.DynamicType; /** @@ -26,16 +25,13 @@ public abstract class ByteBuddyShellGenerator impleme this.shellToolConfig = shellToolConfig; } - protected abstract DynamicType.Builder build(DynamicType.Builder builder); + protected abstract DynamicType.Builder getBuilder(); @Override public byte[] getBytes() { Class shellClass = shellToolConfig.getShellClass(); String shellClassName = shellToolConfig.getShellClassName(); - DynamicType.Builder builder = build(new ByteBuddy() - .redefine(shellClass) - .name(shellClassName) - .visit(new TargetJreVersionVisitorWrapper(shellConfig.getTargetJreVersion()))); + DynamicType.Builder builder = getBuilder(); String shellType = shellConfig.getShellType(); AbstractServer server = ServerFactory.getServer(shellConfig.getServer()); @@ -56,6 +52,10 @@ public abstract class ByteBuddyShellGenerator impleme builder = LogRemoveMethodVisitor.extend(builder); } + builder = builder + .name(shellClassName) + .visit(new TargetJreVersionVisitorWrapper(shellConfig.getTargetJreVersion())); + try (DynamicType.Unloaded unloaded = builder.make()) { return ClassBytesShrink.shrink(unloaded.getBytes(), shellConfig.isShrink()); } diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java b/generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java index 8266fa53..a8ebfe01 100644 --- a/generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java +++ b/generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java @@ -1,11 +1,13 @@ package com.reajason.javaweb.memshell.generator; -import com.reajason.javaweb.ClassBytesShrink; -import com.reajason.javaweb.GenerationException; -import com.reajason.javaweb.asm.ClassRenameUtils; import com.reajason.javaweb.memshell.config.CustomConfig; import com.reajason.javaweb.memshell.config.ShellConfig; -import org.apache.commons.lang3.StringUtils; +import net.bytebuddy.ByteBuddy; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.dynamic.ClassFileLocator; +import net.bytebuddy.dynamic.DynamicType; +import net.bytebuddy.jar.asm.ClassReader; +import net.bytebuddy.pool.TypePool; import java.util.Base64; @@ -13,21 +15,24 @@ import java.util.Base64; * @author ReaJason * @since 2025/3/18 */ -public class CustomShellGenerator extends ASMShellGenerator { +public class CustomShellGenerator extends ByteBuddyShellGenerator { public CustomShellGenerator(ShellConfig shellConfig, CustomConfig customConfig) { super(shellConfig, customConfig); } @Override - public byte[] getBytes() { + protected DynamicType.Builder getBuilder() { String shellClassBase64 = shellToolConfig.getShellClassBase64(); - - if (StringUtils.isBlank(shellClassBase64)) { - throw new GenerationException("Custom shell class is empty"); - } byte[] classBytes = Base64.getDecoder().decode(shellClassBase64); - byte[] bytes = ClassRenameUtils.renameClass(classBytes, shellToolConfig.getShellClassName()); - return ClassBytesShrink.shrink(bytes, shellConfig.isShrink()); + ClassReader classReader = new ClassReader(classBytes); + String className = classReader.getClassName().replace('/', '.'); + ClassFileLocator classFileLocator = ClassFileLocator.Simple.of(className, classBytes); + TypeDescription typeDescription = new TypePool.Default( + new TypePool.CacheProvider.Simple(), classFileLocator, + TypePool.Default.ReaderMode.FAST, TypePool.Default.ofSystemLoader() + ).describe(className).resolve(); + return new ByteBuddy() + .redefine(typeDescription, classFileLocator); } } diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/generator/GodzillaGenerator.java b/generator/src/main/java/com/reajason/javaweb/memshell/generator/GodzillaGenerator.java index ae61e400..8d1a722d 100644 --- a/generator/src/main/java/com/reajason/javaweb/memshell/generator/GodzillaGenerator.java +++ b/generator/src/main/java/com/reajason/javaweb/memshell/generator/GodzillaGenerator.java @@ -2,6 +2,7 @@ package com.reajason.javaweb.memshell.generator; import com.reajason.javaweb.memshell.config.GodzillaConfig; import com.reajason.javaweb.memshell.config.ShellConfig; +import net.bytebuddy.ByteBuddy; import net.bytebuddy.dynamic.DynamicType; import org.apache.commons.codec.digest.DigestUtils; @@ -18,10 +19,12 @@ public class GodzillaGenerator extends ByteBuddyShellGenerator { } @Override - public DynamicType.Builder build(DynamicType.Builder builder) { + public DynamicType.Builder getBuilder() { String md5Key = DigestUtils.md5Hex(shellToolConfig.getKey()).substring(0, 16); String md5 = DigestUtils.md5Hex(shellToolConfig.getPass() + md5Key).toUpperCase(); - return builder.field(named("pass")).value(shellToolConfig.getPass()) + return new ByteBuddy() + .redefine(shellToolConfig.getShellClass()) + .field(named("pass")).value(shellToolConfig.getPass()) .field(named("key")).value(md5Key) .field(named("md5")).value(md5) .field(named("headerName")).value(shellToolConfig.getHeaderName()) diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/generator/NeoreGeorgGenerator.java b/generator/src/main/java/com/reajason/javaweb/memshell/generator/NeoreGeorgGenerator.java index 62974ec5..2b72a602 100644 --- a/generator/src/main/java/com/reajason/javaweb/memshell/generator/NeoreGeorgGenerator.java +++ b/generator/src/main/java/com/reajason/javaweb/memshell/generator/NeoreGeorgGenerator.java @@ -2,6 +2,7 @@ package com.reajason.javaweb.memshell.generator; import com.reajason.javaweb.memshell.config.NeoreGeorgConfig; import com.reajason.javaweb.memshell.config.ShellConfig; +import net.bytebuddy.ByteBuddy; import net.bytebuddy.dynamic.DynamicType; import static net.bytebuddy.matcher.ElementMatchers.named; @@ -16,8 +17,10 @@ public class NeoreGeorgGenerator extends ByteBuddyShellGenerator build(DynamicType.Builder builder) { - return builder.field(named("headerName")).value(shellToolConfig.getHeaderName()) + protected DynamicType.Builder getBuilder() { + return new ByteBuddy() + .redefine(shellToolConfig.getShellClass()) + .field(named("headerName")).value(shellToolConfig.getHeaderName()) .field(named("headerValue")).value(shellToolConfig.getHeaderValue()); } } diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/generator/Suo5Generator.java b/generator/src/main/java/com/reajason/javaweb/memshell/generator/Suo5Generator.java index e133acec..258849b2 100644 --- a/generator/src/main/java/com/reajason/javaweb/memshell/generator/Suo5Generator.java +++ b/generator/src/main/java/com/reajason/javaweb/memshell/generator/Suo5Generator.java @@ -2,6 +2,7 @@ package com.reajason.javaweb.memshell.generator; import com.reajason.javaweb.memshell.config.ShellConfig; import com.reajason.javaweb.memshell.config.Suo5Config; +import net.bytebuddy.ByteBuddy; import net.bytebuddy.dynamic.DynamicType; import static net.bytebuddy.matcher.ElementMatchers.named; @@ -17,8 +18,9 @@ public class Suo5Generator extends ByteBuddyShellGenerator { } @Override - protected DynamicType.Builder build(DynamicType.Builder builder) { - return builder + protected DynamicType.Builder getBuilder() { + return new ByteBuddy() + .redefine(shellToolConfig.getShellClass()) .field(named("headerName")).value(shellToolConfig.getHeaderName()) .field(named("headerValue")).value(shellToolConfig.getHeaderValue()); } diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/generator/command/CommandGenerator.java b/generator/src/main/java/com/reajason/javaweb/memshell/generator/command/CommandGenerator.java index 2d0908af..b275b077 100644 --- a/generator/src/main/java/com/reajason/javaweb/memshell/generator/command/CommandGenerator.java +++ b/generator/src/main/java/com/reajason/javaweb/memshell/generator/command/CommandGenerator.java @@ -7,6 +7,7 @@ import com.reajason.javaweb.memshell.config.CommandConfig; import com.reajason.javaweb.memshell.config.ShellConfig; import com.reajason.javaweb.memshell.generator.ByteBuddyShellGenerator; import com.reajason.javaweb.utils.ShellCommonUtil; +import net.bytebuddy.ByteBuddy; import net.bytebuddy.asm.Advice; import net.bytebuddy.description.modifier.Ownership; import net.bytebuddy.description.modifier.Visibility; @@ -26,9 +27,11 @@ public class CommandGenerator extends ByteBuddyShellGenerator { } @Override - public DynamicType.Builder build(DynamicType.Builder builder) { - - builder = builder.field(named("paramName")).value(shellToolConfig.getParamName()); + public DynamicType.Builder getBuilder() { + DynamicType.Builder builder = new ByteBuddy() + .redefine(shellToolConfig.getShellClass()) + .field(named("paramName")) + .value(shellToolConfig.getParamName()); if (shellConfig.isJakarta()) { builder = builder.visit(ServletRenameVisitorWrapper.INSTANCE); diff --git a/generator/src/test/java/com/reajason/javaweb/memshell/generator/CustomShellGeneratorTest.java b/generator/src/test/java/com/reajason/javaweb/memshell/generator/CustomShellGeneratorTest.java index 890af712..de9a4910 100644 --- a/generator/src/test/java/com/reajason/javaweb/memshell/generator/CustomShellGeneratorTest.java +++ b/generator/src/test/java/com/reajason/javaweb/memshell/generator/CustomShellGeneratorTest.java @@ -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 referencedClasses = classVisitor.getReferencedClasses(); + assertEquals(className, classReader.getClassName().replace("/", ".")); + assertTrue(referencedClasses.contains("com/bes/enterprise/webtier/Valve")); + assertFalse(referencedClasses.contains("org/apache/catalina/Valve")); + } } \ No newline at end of file diff --git a/memshell-party-common/src/main/java/com/reajason/javaweb/asm/ClassReferenceVisitor.java b/memshell-party-common/src/main/java/com/reajason/javaweb/asm/ClassReferenceVisitor.java new file mode 100644 index 00000000..70fc8f40 --- /dev/null +++ b/memshell-party-common/src/main/java/com/reajason/javaweb/asm/ClassReferenceVisitor.java @@ -0,0 +1,83 @@ +package com.reajason.javaweb.asm; + +import lombok.Getter; +import org.objectweb.asm.*; + +import java.util.Collections; +import java.util.HashSet; +import java.util.Set; + +@Getter +public class ClassReferenceVisitor extends ClassVisitor { + private final Set referencedClasses = new HashSet<>(); + + public ClassReferenceVisitor() { + super(Opcodes.ASM9); + } + + @Override + public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) { + if (superName != null) { + referencedClasses.add(superName); + } + if (interfaces != null) { + Collections.addAll(referencedClasses, interfaces); + } + super.visit(version, access, name, signature, superName, interfaces); + } + + @Override + public FieldVisitor visitField(int access, String name, String descriptor, String signature, Object value) { + addType(Type.getType(descriptor)); + return super.visitField(access, name, descriptor, signature, value); + } + + @Override + public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) { + addType(Type.getMethodType(descriptor)); + if (exceptions != null) { + Collections.addAll(referencedClasses, exceptions); + } + return new MethodReferenceVisitor(super.visitMethod(access, name, descriptor, signature, exceptions)); + } + + private void addType(Type type) { + if (type.getSort() == Type.OBJECT) { + referencedClasses.add(type.getInternalName()); + } else if (type.getSort() == Type.ARRAY) { + addType(type.getElementType()); + } else if (type.getSort() == Type.METHOD) { + addType(type.getReturnType()); + for (Type argType : type.getArgumentTypes()) { + addType(argType); + } + } + } + + class MethodReferenceVisitor extends MethodVisitor { + public MethodReferenceVisitor(MethodVisitor methodVisitor) { + super(Opcodes.ASM9, methodVisitor); + } + + @Override + public void visitMethodInsn(int opcode, String owner, String name, String descriptor, boolean isInterface) { + referencedClasses.add(owner); + addType(Type.getMethodType(descriptor)); + super.visitMethodInsn(opcode, owner, name, descriptor, isInterface); + } + + @Override + public void visitLdcInsn(Object value) { + if (value instanceof Type) { + addType((Type) value); + } + super.visitLdcInsn(value); + } + + @Override + public void visitTypeInsn(int opcode, String type) { + referencedClasses.add(type); + super.visitTypeInsn(opcode, type); + } + } +} \ No newline at end of file