From 2196a4aefba8f23e3ce8c8916822e0c7bc08a5e9 Mon Sep 17 00:00:00 2001 From: ReaJason Date: Tue, 16 Sep 2025 21:05:46 +0800 Subject: [PATCH] fix: custom listener shell generate failed --- .../memshell/config/ShellToolConfig.java | 2 + .../generator/ByteBuddyShellGenerator.java | 12 +++- .../generator/CustomShellGenerator.java | 1 + .../memshell/generator/ListenerGenerator.java | 33 ++++++---- .../generator/CustomShellGeneratorTest.java | 25 ++++++- .../generator/ListenerGeneratorTest.java | 66 +++++++++++++++++++ 6 files changed, 125 insertions(+), 14 deletions(-) create mode 100644 generator/src/test/java/com/reajason/javaweb/memshell/generator/ListenerGeneratorTest.java 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 index 767ee922..ce3a4e08 100644 --- 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.AllArgsConstructor; 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 index a6393ff3..f1d8eddb 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 @@ -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.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.description.type.TypeDescription; import net.bytebuddy.dynamic.DynamicType; /** @@ -29,15 +31,21 @@ public abstract class ByteBuddyShellGenerator impleme @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 index a8ebfe01..d0f57a93 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 @@ -32,6 +32,7 @@ public class CustomShellGenerator extends ByteBuddyShellGenerator 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 index 3a689062..88fb0de5 100644 --- 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 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; */ 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 methods = typeDefinition.getDeclaredMethods(); - boolean methodNotFound = targetClass != null && TypeDescription.ForLoadedType.of(targetClass) - .getDeclaredMethods() - .filter(named("getFieldValue") + 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 (methods.filter(named("getFieldValue") .and(takesArguments(Object.class, String.class))) - .isEmpty(); - - if (methodNotFound) { + .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 index de9a4910..9909e716 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 @@ -5,10 +5,12 @@ 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.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 @@ import static org.junit.jupiter.api.Assertions.*; * @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 new file mode 100644 index 00000000..3112a139 --- /dev/null +++ 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); + } +} \ No newline at end of file