fix: method body write failed with jakarta renamed

ByteBuddy Advice cache the class method list at first, the method signature will be changed when use jakarta rename, if follow advice need enhance the javax method, it will not find method signature. because ByteBuddy cached javax method, but follow advice need find the jakarta method, boom~

Author: ReaJason

generator/src/main/java/com/reajason/javaweb/memshell/generator/ProcessorRegistry.java

@@ -17,11 +17,11 @@ 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 JakartaPostProcessor(),
             new JettyHandlerPostProcessor(),
             new ShrinkPostProcessor()
     );

generator/src/main/java/com/reajason/javaweb/memshell/generator/processors/JakartaBuilderModifier.java

@@ -0,0 +1,16 @@
+package com.reajason.javaweb.memshell.generator.processors;
+
+import com.reajason.javaweb.asm.ClassRenameUtils;
+import com.reajason.javaweb.memshell.config.ShellConfig;
+import com.reajason.javaweb.memshell.config.ShellToolConfig;
+import com.reajason.javaweb.memshell.generator.Processor;
+
+public class JakartaPostProcessor implements Processor<byte[]> {
+    @Override
+    public byte[] process(byte[] input, ShellConfig shellConfig, ShellToolConfig shellToolConfig) {
+        if (shellConfig.isJakarta()) {
+            return ClassRenameUtils.relocateJakarta(input);
+        }
+        return input;
+    }
+}

generator/src/main/java/com/reajason/javaweb/memshell/generator/processors/JakartaPostProcessor.java

@@ -17,6 +17,7 @@
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.DynamicType;
 import net.bytebuddy.implementation.FixedValue;
+import net.bytebuddy.implementation.StubMethod;
 import net.bytebuddy.matcher.ElementMatchers;
 
 import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -46,26 +47,21 @@ public static DynamicType.Builder<?> modifier(DynamicType.Builder<?> builder, Cl
                                                   TypeDescription typeDefinition, String newClassName) {
         MethodList<MethodDescription.InDefinedShape> methods = typeDefinition.getDeclaredMethods();
 
-        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");
+        if (methods.filter(named("getResponseFromRequest").and(takesArguments(1))).isEmpty()) {
+            throw new GenerationException("please add [getResponseFromRequest(Object request)] method," +
+                    " the method body will be auto adapted for multi server");
         } 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 (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)
-                    .intercept(FixedValue.nullValue())
+                    .intercept(StubMethod.INSTANCE)
                     .visit(Advice.to(ShellCommonUtil.GetFieldValueInterceptor.class).on(named("getFieldValue")));
         }
         return builder;

generator/src/main/java/com/reajason/javaweb/memshell/generator/processors/ListenerBuilderModifier.java

@@ -28,12 +28,6 @@ public Object getResponseFromRequest(Object request) {
         }
     }
 
-    public static class J {
-        public HttpServletResponse getResponseFromRequest(HttpServletRequest request) {
-            return null;
-        }
-    }
-
     public static class FakeRequest {
         public Object response = "i'm a good boy";
     }
@@ -44,12 +38,6 @@ void testNoGetResponseFromRequest() {
         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, () -> ListenerBuilderModifier.modifier(builder, Tomcat.ListenerInterceptor.class, TypeDescription.ForLoadedType.of(J.class), "hello.world"));
-    }
-
     @Test
     @SneakyThrows
     void test() {

generator/src/test/java/com/reajason/javaweb/memshell/generator/ListenerGeneratorTest.java

@@ -1,5 +1,6 @@
 package com.reajason.javaweb.asm;
 
+import org.jetbrains.annotations.NotNull;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.ClassWriter;
 import org.objectweb.asm.Opcodes;
@@ -14,12 +15,7 @@
 public class ClassRenameUtils {
 
     public static byte[] renameClass(byte[] classBytes, String newName) {
-        ClassReader reader = null;
-        try {
-            reader = new ClassReader(classBytes);
-        } catch (Exception e) {
-            throw new RuntimeException("invalid class bytes");
-        }
+        ClassReader reader = getClassReader(classBytes);
         String oldClassName = reader.getClassName();
         String newClassName = newName.replace('.', '/');
         ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES);
@@ -29,12 +25,7 @@ public static byte[] renameClass(byte[] classBytes, String newName) {
     }
 
     public static byte[] relocateClass(byte[] classBytes, String relocateClassPackage, String relocatePrefix) {
-        ClassReader reader = null;
-        try {
-            reader = new ClassReader(classBytes);
-        } catch (Exception e) {
-            throw new RuntimeException("invalid class bytes");
-        }
+        ClassReader reader = getClassReader(classBytes);
         String oldClassName = relocateClassPackage.replace('.', '/');
         String newClassName = relocatePrefix.replace('.', '/');
         ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_MAXS);
@@ -51,4 +42,32 @@ public String map(String typeName) {
         reader.accept(adapter, 0);
         return writer.toByteArray();
     }
+
+    public static byte[] relocateJakarta(byte[] classBytes) {
+        ClassReader reader = getClassReader(classBytes);
+        ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_MAXS);
+        ClassRemapper adapter = new ClassRemapper(writer, new Remapper(Opcodes.ASM9) {
+            @Override
+            public String map(String typeName) {
+                if (typeName.startsWith("javax/servlet/")
+                        || typeName.startsWith("javax/websocket/")) {
+                    return typeName.replaceFirst("javax", "jakarta");
+                } else {
+                    return typeName;
+                }
+            }
+        });
+        reader.accept(adapter, 0);
+        return writer.toByteArray();
+    }
+
+    private static @NotNull ClassReader getClassReader(byte[] classBytes) {
+        ClassReader reader = null;
+        try {
+            reader = new ClassReader(classBytes);
+        } catch (Exception e) {
+            throw new RuntimeException("invalid class bytes");
+        }
+        return reader;
+    }
 }