fix: custom shell not supported for listener and valve vistor

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<AntSwordConfig> {
     }
 
     @Override
-    protected DynamicType.Builder<?> build(DynamicType.Builder<?> builder) {
+    protected DynamicType.Builder<?> getBuilder() {
-        return builder.field(named("pass")).value(shellToolConfig.getPass())
+        return new ByteBuddy()
+                .redefine(shellToolConfig.getShellClass())
+                .field(named("pass")).value(shellToolConfig.getPass())
                 .field(named("headerName")).value(shellToolConfig.getHeaderName())
                 .field(named("headerValue")).value(shellToolConfig.getHeaderValue());
     }

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<BehinderConfig> {
         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());
     }

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<T extends ShellToolConfig> 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()
+        DynamicType.Builder<?> builder = getBuilder();
-                .redefine(shellClass)
-                .name(shellClassName)
-                .visit(new TargetJreVersionVisitorWrapper(shellConfig.getTargetJreVersion())));
 
         String shellType = shellConfig.getShellType();
         AbstractServer server = ServerFactory.getServer(shellConfig.getServer());
@@ -56,6 +52,10 @@ public abstract class ByteBuddyShellGenerator<T extends ShellToolConfig> 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());
         }

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<CustomConfig> {
+public class CustomShellGenerator extends ByteBuddyShellGenerator<CustomConfig> {
 
     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());
+        ClassReader classReader = new ClassReader(classBytes);
-        return ClassBytesShrink.shrink(bytes, shellConfig.isShrink());
+        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);
     }
 }

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<GodzillaConfig> {
     }
 
     @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())

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<NeoreGeorgConfi
     }
 
     @Override
-    protected DynamicType.Builder<?> build(DynamicType.Builder<?> builder) {
+    protected DynamicType.Builder<?> getBuilder() {
-        return builder.field(named("headerName")).value(shellToolConfig.getHeaderName())
+        return new ByteBuddy()
+                .redefine(shellToolConfig.getShellClass())
+                .field(named("headerName")).value(shellToolConfig.getHeaderName())
                 .field(named("headerValue")).value(shellToolConfig.getHeaderValue());
     }
 }

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<Suo5Config> {
     }
 
     @Override
-    protected DynamicType.Builder<?> build(DynamicType.Builder<?> builder) {
+    protected DynamicType.Builder<?> getBuilder() {
-        return builder
+        return new ByteBuddy()
+                .redefine(shellToolConfig.getShellClass())
                 .field(named("headerName")).value(shellToolConfig.getHeaderName())
                 .field(named("headerValue")).value(shellToolConfig.getHeaderValue());
     }

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<CommandConfig> {
     }
 
     @Override
-    public DynamicType.Builder<?> build(DynamicType.Builder<?> builder) {
+    public DynamicType.Builder<?> getBuilder() {
-
+        DynamicType.Builder<?> builder = new ByteBuddy()
-        builder = builder.field(named("paramName")).value(shellToolConfig.getParamName());
+                .redefine(shellToolConfig.getShellClass())
+                .field(named("paramName"))
+                .value(shellToolConfig.getParamName());
 
         if (shellConfig.isJakarta()) {
             builder = builder.visit(ServletRenameVisitorWrapper.INSTANCE);

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<String> referencedClasses = classVisitor.getReferencedClasses();
+        assertEquals(className, classReader.getClassName().replace("/", "."));
+        assertTrue(referencedClasses.contains("com/bes/enterprise/webtier/Valve"));
+        assertFalse(referencedClasses.contains("org/apache/catalina/Valve"));
+    }
 }

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<String> 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);
+        }
+    }
+}