refactor: use processors

2026-09-22 07:00:43 +08:00 · 2025-12-08 01:43:41 +08:00
parent b29d83d3b4
commit 5d28df0e39
13 changed files with 235 additions and 134 deletions
@@ -1,17 +1,10 @@
 package com.reajason.javaweb.memshell.generator;

-import com.reajason.javaweb.ClassBytesShrink;
 import com.reajason.javaweb.GenerationException;
 import com.reajason.javaweb.ShellGenerator;
-import com.reajason.javaweb.buddy.LogRemoveMethodVisitor;
-import com.reajason.javaweb.buddy.ServletRenameVisitorWrapper;
 import com.reajason.javaweb.buddy.TargetJreVersionVisitorWrapper;
-import com.reajason.javaweb.memshell.ServerFactory;
-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 com.reajason.javaweb.memshell.server.Jetty;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.DynamicType;

@@ -35,43 +28,21 @@ public abstract class ByteBuddyShellGenerator<T extends ShellToolConfig> impleme
        DynamicType.Builder<?> builder = getBuilder();
        String shellClassName = shellToolConfig.getShellClassName();
        Class<?> shellClass = shellToolConfig.getShellClass();
+
        if (shellClass != null) {
            shellToolConfig.setShellTypeDescription(TypeDescription.ForLoadedType.of(shellClass));
        }
+
        if (shellToolConfig.getShellTypeDescription() == null) {
            throw new GenerationException("shellClass or shellTypeDescription could not be null.");
        }

-        String shellType = shellConfig.getShellType();
-        AbstractServer server = ServerFactory.getServer(shellConfig.getServer());
-
-        if (ShellType.LISTENER.equals(shellType) || ShellType.JAKARTA_LISTENER.equals(shellType)) {
-            builder = ListenerGenerator.build(builder, server.getListenerInterceptor(), shellToolConfig.getShellTypeDescription(), shellClassName);
-        }
-
-        if (ShellType.VALVE.equals(shellType) || ShellType.JAKARTA_VALVE.equals(shellType)) {
-            builder = ValveGenerator.build(builder, server, shellConfig.getServerVersion());
-        }
-
-        if (server instanceof Jetty
-                && (ShellType.HANDLER.equals(shellType) || ShellType.JAKARTA_HANDLER.equals(shellType))) {
-            builder = JettyHandlerGenerator.build(builder, shellConfig.getServerVersion());
-        }
-
-        if (shellConfig.isJakarta()) {
-            builder = builder.visit(ServletRenameVisitorWrapper.INSTANCE);
-        }
-
-        if (shellConfig.isDebugOff()) {
-            builder = LogRemoveMethodVisitor.extend(builder);
-        }
-
-        builder = builder
+        builder = ProcessorRegistry.applyBuilderProcessors(builder, shellConfig, shellToolConfig)
                .name(shellClassName)
                .visit(new TargetJreVersionVisitorWrapper(shellConfig.getTargetJreVersion()));

        try (DynamicType.Unloaded<?> unloaded = builder.make()) {
-            return ClassBytesShrink.shrink(unloaded.getBytes(), shellConfig.isShrink());
+            return ProcessorRegistry.applyByteProcessors(unloaded.getBytes(), shellConfig, shellToolConfig);
        }
    }
 }
@@ -1,85 +0,0 @@
-package com.reajason.javaweb.memshell.generator;
-
-import com.reajason.javaweb.GenerationException;
-import com.reajason.javaweb.buddy.ClassRenameVisitorWrapper;
-import net.bytebuddy.asm.AsmVisitorWrapper;
-import net.bytebuddy.description.field.FieldDescription;
-import net.bytebuddy.description.field.FieldList;
-import net.bytebuddy.description.method.MethodList;
-import net.bytebuddy.description.type.TypeDescription;
-import net.bytebuddy.dynamic.DynamicType;
-import net.bytebuddy.implementation.Implementation;
-import net.bytebuddy.jar.asm.ClassVisitor;
-import net.bytebuddy.jar.asm.MethodVisitor;
-import net.bytebuddy.jar.asm.Opcodes;
-import net.bytebuddy.pool.TypePool;
-
-/**
- * @author ReaJason
- * @since 2025/12/2
- */
-public class JettyHandlerGenerator {
-    public static DynamicType.Builder<?> build(DynamicType.Builder<?> builder, String serverVersion) {
-        String superClassName = null;
-        DynamicType.Builder<?> newBuilder = builder;
-        if (serverVersion != null) {
-            switch (serverVersion) {
-                case "6":
-                    superClassName = "org/mortbay/jetty/handler/AbstractHandler";
-                    newBuilder = newBuilder.visit(new ClassRenameVisitorWrapper("org/eclipse/jetty/server", "org/mortbay/jetty"));
-                    break;
-                case "7+":
-                    superClassName = "org/eclipse/jetty/server/handler/AbstractHandler";
-                    break;
-                case "12":
-                    superClassName = "org/eclipse/jetty/server/Handler$Abstract";
-                    break;
-            }
-        }
-        if (superClassName == null) {
-            throw new GenerationException("serverVersion is needed for Jetty Handler or unknow serverVersion: [" + serverVersion + "], please use one of ['6', '7+', '12'] for shellConfig.serverVersion");
-        }
-        String finalSuperClassName = superClassName;
-        return newBuilder.visit(new AsmVisitorWrapper.ForDeclaredMethods() {
-            @Override
-            public ClassVisitor wrap(TypeDescription instrumentedType,
-                                     ClassVisitor classVisitor,
-                                     Implementation.Context implementationContext,
-                                     TypePool typePool,
-                                     FieldList<FieldDescription.InDefinedShape> fields,
-                                     MethodList<?> methods,
-                                     int writerFlags,
-                                     int readerFlags) {
-                return new ClassVisitor(Opcodes.ASM9, classVisitor) {
-                    @Override
-                    public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
-                        super.visit(version, access, name, signature, finalSuperClassName, interfaces);
-                    }
-
-                    @Override
-                    public MethodVisitor visitMethod(int access, String name, String descriptor,
-                                                     String signature, String[] exceptions) {
-
-                        MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
-                        if (!name.equals("<init>")) {
-                            return mv;
-                        }
-                        return new MethodVisitor(Opcodes.ASM9, mv) {
-                            @Override
-                            public void visitMethodInsn(int opcode, String owner, String name, String descriptor, boolean isInterface) {
-                                if (opcode == org.objectweb.asm.Opcodes.INVOKESPECIAL &&
-                                        "java/lang/Object".equals(owner) &&
-                                        "<init>".equals(name)) {
-                                    super.visitMethodInsn(opcode, finalSuperClassName, name, descriptor, isInterface);
-                                } else {
-                                    super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
-                                }
-                            }
-                        };
-                    }
-                };
-
-            }
-        });
-    }
-}
@@ -0,0 +1,12 @@
+package com.reajason.javaweb.memshell.generator;
+
+import com.reajason.javaweb.memshell.config.ShellConfig;
+import com.reajason.javaweb.memshell.config.ShellToolConfig;
+
+/**
+ * @author ReaJason
+ * @since 2025/12/7
+ */
+public interface Processor<T> {
+    T process(T input, ShellConfig shellConfig, ShellToolConfig shellToolConfig);
+}
@@ -0,0 +1,52 @@
+package com.reajason.javaweb.memshell.generator;
+
+import com.reajason.javaweb.memshell.config.ShellConfig;
+import com.reajason.javaweb.memshell.config.ShellToolConfig;
+import com.reajason.javaweb.memshell.generator.processors.*;
+import net.bytebuddy.dynamic.DynamicType;
+
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * @author ReaJason
+ * @since 2025/12/7
+ */
+public final class ProcessorRegistry {
+
+    private static final List<Processor<DynamicType.Builder<?>>> BUILDER_PROCESSORS = Arrays.asList(
+            new ListenerBuilderModifier(),
+            new ValveBuilderModifier(),
+            new JakartaBuilderModifier(),
+            new DebugOffBuilderModifier()
+    );
+
+    private static final List<Processor<byte[]>> BYTE_PROCESSORS = Arrays.asList(
+            new ShrinkPostProcessor(),
+            new JettyHandlerPostProcessor()
+    );
+
+    private ProcessorRegistry() {
+        // Prevent instantiation
+    }
+
+    public static DynamicType.Builder<?> applyBuilderProcessors(
+            DynamicType.Builder<?> builder,
+            ShellConfig shellConfig,
+            ShellToolConfig shellToolConfig) {
+        for (Processor<DynamicType.Builder<?>> processor : BUILDER_PROCESSORS) {
+            builder = processor.process(builder, shellConfig, shellToolConfig);
+        }
+        return builder;
+    }
+
+    public static byte[] applyByteProcessors(
+            byte[] bytes,
+            ShellConfig shellConfig,
+            ShellToolConfig shellToolConfig) {
+        for (Processor<byte[]> processor : BYTE_PROCESSORS) {
+            bytes = processor.process(bytes, shellConfig, shellToolConfig);
+        }
+        return bytes;
+    }
+}
@@ -0,0 +1,22 @@
+package com.reajason.javaweb.memshell.generator.processors;
+
+import com.reajason.javaweb.buddy.LogRemoveMethodVisitor;
+import com.reajason.javaweb.memshell.config.ShellConfig;
+import com.reajason.javaweb.memshell.config.ShellToolConfig;
+import com.reajason.javaweb.memshell.generator.Processor;
+import net.bytebuddy.dynamic.DynamicType;
+
+/**
+ * @author ReaJason
+ * @since 2025/12/7
+ */
+public class DebugOffBuilderModifier implements Processor<DynamicType.Builder<?>> {
+
+    @Override
+    public DynamicType.Builder<?> process(DynamicType.Builder<?> builder, ShellConfig shellConfig, ShellToolConfig shellToolConfig) {
+        if (shellConfig.isDebugOff()) {
+            builder = LogRemoveMethodVisitor.extend(builder);
+        }
+        return builder;
+    }
+}
@@ -0,0 +1,22 @@
+package com.reajason.javaweb.memshell.generator.processors;
+
+import com.reajason.javaweb.buddy.ServletRenameVisitorWrapper;
+import com.reajason.javaweb.memshell.config.ShellConfig;
+import com.reajason.javaweb.memshell.config.ShellToolConfig;
+import com.reajason.javaweb.memshell.generator.Processor;
+import net.bytebuddy.dynamic.DynamicType;
+
+/**
+ * @author ReaJason
+ * @since 2025/12/7
+ */
+public class JakartaBuilderModifier implements Processor<DynamicType.Builder<?>> {
+
+    @Override
+    public DynamicType.Builder<?> process(DynamicType.Builder<?> builder, ShellConfig shellConfig, ShellToolConfig shellToolConfig) {
+        if (shellConfig.isJakarta()) {
+            builder = builder.visit(ServletRenameVisitorWrapper.INSTANCE);
+        }
+        return builder;
+    }
+}
@@ -0,0 +1,51 @@
+package com.reajason.javaweb.memshell.generator.processors;
+
+import com.reajason.javaweb.GenerationException;
+import com.reajason.javaweb.asm.ClassRenameUtils;
+import com.reajason.javaweb.asm.ClassSuperClassUtils;
+import com.reajason.javaweb.memshell.ServerFactory;
+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.generator.Processor;
+import com.reajason.javaweb.memshell.server.AbstractServer;
+import com.reajason.javaweb.memshell.server.Jetty;
+
+/**
+ * @author ReaJason
+ * @since 2025/12/7
+ */
+public class JettyHandlerPostProcessor implements Processor<byte[]> {
+
+    @Override
+    public byte[] process(byte[] bytes, ShellConfig shellConfig, ShellToolConfig shellToolConfig) {
+        AbstractServer server = ServerFactory.getServer(shellConfig.getServer());
+        String shellType = shellConfig.getShellType();
+        if (server instanceof Jetty
+                && (ShellType.HANDLER.equals(shellType)
+                || ShellType.JAKARTA_HANDLER.equals(shellType))
+        ) {
+            String superClassName = null;
+            String serverVersion = shellConfig.getServerVersion();
+            if (serverVersion != null) {
+                switch (serverVersion) {
+                    case "6":
+                        superClassName = "org/mortbay/jetty/handler/AbstractHandler";
+                        bytes = ClassRenameUtils.relocateClass(bytes, "org/eclipse/jetty/server", "org/mortbay/jetty");
+                        break;
+                    case "7+":
+                        superClassName = "org/eclipse/jetty/server/handler/AbstractHandler";
+                        break;
+                    case "12":
+                        superClassName = "org/eclipse/jetty/server/Handler$Abstract";
+                        break;
+                }
+            }
+            if (superClassName == null) {
+                throw new GenerationException("serverVersion is needed for Jetty Handler or unknow serverVersion: [" + serverVersion + "], please use one of ['6', '7+', '12'] for shellConfig.serverVersion");
+            }
+            return ClassSuperClassUtils.addSuperClass(bytes, superClassName);
+        }
+        return bytes;
+    }
+}
@@ -1,7 +1,13 @@
-package com.reajason.javaweb.memshell.generator;
+package com.reajason.javaweb.memshell.generator.processors;

 import com.reajason.javaweb.GenerationException;
 import com.reajason.javaweb.buddy.MethodCallReplaceVisitorWrapper;
+import com.reajason.javaweb.memshell.ServerFactory;
+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.generator.Processor;
+import com.reajason.javaweb.memshell.server.AbstractServer;
 import com.reajason.javaweb.utils.ShellCommonUtil;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
@@ -18,12 +24,26 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments;

 /**
 * @author ReaJason
- * @since 2025/2/22
+ * @since 2025/12/7
 */
-public class ListenerGenerator {
+public class ListenerBuilderModifier implements Processor<DynamicType.Builder<?>> {

-    public static DynamicType.Builder<?> build(DynamicType.Builder<?> builder, Class<?> implInterceptor,
-                                               TypeDescription typeDefinition, String newClassName) {
+    @Override
+    public DynamicType.Builder<?> process(DynamicType.Builder<?> builder, ShellConfig shellConfig, ShellToolConfig shellToolConfig) {
+        String shellType = shellConfig.getShellType();
+        if (ShellType.LISTENER.equals(shellType) || ShellType.JAKARTA_LISTENER.equals(shellType)) {
+            AbstractServer server = ServerFactory.getServer(shellConfig.getServer());
+            String shellClassName = shellToolConfig.getShellClassName();
+            builder = modifier(builder,
+                    server.getListenerInterceptor(),
+                    shellToolConfig.getShellTypeDescription(),
+                    shellClassName);
+        }
+        return builder;
+    }
+
+    public static DynamicType.Builder<?> modifier(DynamicType.Builder<?> builder, Class<?> implInterceptor,
+                                                  TypeDescription typeDefinition, String newClassName) {
        MethodList<MethodDescription.InDefinedShape> methods = typeDefinition.getDeclaredMethods();

        if (methods.filter(ElementMatchers.named("getResponseFromRequest")
@@ -0,0 +1,18 @@
+package com.reajason.javaweb.memshell.generator.processors;
+
+import com.reajason.javaweb.ClassBytesShrink;
+import com.reajason.javaweb.memshell.config.ShellConfig;
+import com.reajason.javaweb.memshell.config.ShellToolConfig;
+import com.reajason.javaweb.memshell.generator.Processor;
+
+/**
+ * @author ReaJason
+ * @since 2025/12/7
+ */
+public class ShrinkPostProcessor implements Processor<byte[]> {
+
+    @Override
+    public byte[] process(byte[] bytes, ShellConfig shellConfig, ShellToolConfig shellToolConfig) {
+        return ClassBytesShrink.shrink(bytes, shellConfig.isShrink());
+    }
+}
@@ -1,6 +1,11 @@
-package com.reajason.javaweb.memshell.generator;
+package com.reajason.javaweb.memshell.generator.processors;

 import com.reajason.javaweb.GenerationException;
+import com.reajason.javaweb.memshell.ServerFactory;
+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.generator.Processor;
 import com.reajason.javaweb.memshell.server.AbstractServer;
 import com.reajason.javaweb.memshell.server.Bes;
 import com.reajason.javaweb.memshell.server.TongWeb;
@@ -12,6 +17,7 @@ import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.DynamicType;
 import net.bytebuddy.implementation.Implementation;
 import net.bytebuddy.jar.asm.ClassVisitor;
+import net.bytebuddy.jar.asm.Opcodes;
 import net.bytebuddy.jar.asm.commons.ClassRemapper;
 import net.bytebuddy.jar.asm.commons.Remapper;
 import net.bytebuddy.pool.TypePool;
@@ -20,9 +26,19 @@ import org.jetbrains.annotations.NotNull;

 /**
 * @author ReaJason
- * @since 2025/2/22
+ * @since 2025/12/7
 */
-public class ValveGenerator {
+public class ValveBuilderModifier implements Processor<DynamicType.Builder<?>> {
+
+    @Override
+    public DynamicType.Builder<?> process(DynamicType.Builder<?> builder, ShellConfig shellConfig, ShellToolConfig shellToolConfig) {
+        String shellType = shellConfig.getShellType();
+        AbstractServer server = ServerFactory.getServer(shellConfig.getServer());
+        if (ShellType.VALVE.equals(shellType) || ShellType.JAKARTA_VALVE.equals(shellType)) {
+            builder = modifier(builder, server, shellConfig.getServerVersion());
+        }
+        return builder;
+    }

    public static final String CATALINA_VALVE_PACKAGE = "org.apache.catalina";
    public static final String BES_VALVE_PACKAGE = "com.bes.enterprise.webtier";
@@ -30,7 +46,7 @@ public class ValveGenerator {
    public static final String TONGWEB7_VALVE_PACKAGE = "com.tongweb.catalina";
    public static final String TONGWEB8_VALVE_PACKAGE = "com.tongweb.server";

-    public static DynamicType.Builder<?> build(DynamicType.Builder<?> builder, AbstractServer shell, String serverVersion) {
+    public static DynamicType.Builder<?> modifier(DynamicType.Builder<?> builder, AbstractServer shell, String serverVersion) {
        String packageName = null;
        if (shell instanceof Bes) {
            packageName = BES_VALVE_PACKAGE;
@@ -88,7 +104,7 @@ public class ValveGenerator {
                                 int readerFlags) {
            return new ClassRemapper(
                    classVisitor,
-                    new Remapper() {
+                    new Remapper(Opcodes.ASM9) {
                        @Override
                        public String map(String typeName) {
                            String packageName = CATALINA_VALVE_PACKAGE.replace(".", "/");
@@ -1,6 +1,7 @@
 package com.reajason.javaweb.memshell.generator;

 import com.reajason.javaweb.GenerationException;
+import com.reajason.javaweb.memshell.generator.processors.ListenerBuilderModifier;
 import com.reajason.javaweb.memshell.server.Tomcat;
 import lombok.SneakyThrows;
 import net.bytebuddy.ByteBuddy;
@@ -42,20 +43,20 @@ class ListenerGeneratorTest {
    @Test
    void testNoGetResponseFromRequest() {
        DynamicType.Builder<?> builder = new ByteBuddy().redefine(Object.class);
-        Assertions.assertThrows(GenerationException.class, () -> ListenerGenerator.build(builder, Tomcat.ListenerInterceptor.class, TypeDescription.ForLoadedType.of(Object.class), "hello.world"));
+        Assertions.assertThrows(GenerationException.class, () -> ListenerBuilderModifier.modifier(builder, Tomcat.ListenerInterceptor.class, TypeDescription.ForLoadedType.of(Object.class), "hello.world"));
    }

    @Test
    void testGetResponseFromRequestSignatureError() {
        DynamicType.Builder<?> builder = new ByteBuddy().redefine(J.class);
-        Assertions.assertThrows(GenerationException.class, () -> ListenerGenerator.build(builder, Tomcat.ListenerInterceptor.class, TypeDescription.ForLoadedType.of(J.class), "hello.world"));
+        Assertions.assertThrows(GenerationException.class, () -> ListenerBuilderModifier.modifier(builder, Tomcat.ListenerInterceptor.class, TypeDescription.ForLoadedType.of(J.class), "hello.world"));
    }

    @Test
    @SneakyThrows
    void test() {
        String className = "hello.world";
-        DynamicType.Builder<?> build = ListenerGenerator.build(new ByteBuddy().redefine(L.class).name(className), Tomcat.ListenerInterceptor.class, TypeDescription.ForLoadedType.of(L.class), className);
+        DynamicType.Builder<?> build = ListenerBuilderModifier.modifier(new ByteBuddy().redefine(L.class).name(className), Tomcat.ListenerInterceptor.class, TypeDescription.ForLoadedType.of(L.class), className);
        Class<?> clazz = build.make().load(getClass().getClassLoader()).getLoaded();
        Object obj = clazz.newInstance();
        Method getResponseFromRequest = clazz.getDeclaredMethod("getResponseFromRequest", Object.class);