From fb5b39a4c64cace42f82a19f8f02412792f16470 Mon Sep 17 00:00:00 2001
From: ReaJason <reajason1225@gmail.com>
Date: Tue, 18 Nov 2025 19:37:03 +0800
Subject: [PATCH] feat: support probe get payload from header by default

---
 .../boot/dto/ProbeShellGenerateRequest.java   |  2 --
 .../probe/config/ResponseBodyConfig.java      |  1 -
 .../response/ResponseBodyGenerator.java       | 34 +++++--------------
 .../javaweb/integration/ProbeAssertion.java   |  2 +-
 .../probeshell/main-config-card.tsx           |  1 -
 web/src/pages/probeshell.tsx                  |  1 -
 web/src/types/probeshell.ts                   |  3 --
 web/src/types/schema.ts                       |  1 -
 web/src/utils/transformer.ts                  |  1 -
 9 files changed, 9 insertions(+), 37 deletions(-)

diff --git a/boot/src/main/java/com/reajason/javaweb/boot/dto/ProbeShellGenerateRequest.java b/boot/src/main/java/com/reajason/javaweb/boot/dto/ProbeShellGenerateRequest.java
index cb4920a0..78cca343 100644
--- a/boot/src/main/java/com/reajason/javaweb/boot/dto/ProbeShellGenerateRequest.java
+++ b/boot/src/main/java/com/reajason/javaweb/boot/dto/ProbeShellGenerateRequest.java
@@ -21,7 +21,6 @@ public class ProbeShellGenerateRequest {
         private String server;
         private String sleepServer;
         private String reqParamName;
-        private String reqHeaderName;
     }
 
     public ProbeContentConfig parseProbeContentConfig() {
@@ -35,7 +34,6 @@ public class ProbeShellGenerateRequest {
                     .build();
             case ResponseBody -> ResponseBodyConfig.builder()
                     .reqParamName(probeContentConfig.reqParamName)
-                    .reqHeaderName(probeContentConfig.reqHeaderName)
                     .server(probeContentConfig.server)
                     .build();
             default -> throw new UnsupportedOperationException("unknown probe method: " + probeConfig.getProbeMethod());
diff --git a/generator/src/main/java/com/reajason/javaweb/probe/config/ResponseBodyConfig.java b/generator/src/main/java/com/reajason/javaweb/probe/config/ResponseBodyConfig.java
index 99fae0ac..f3e59f55 100644
--- a/generator/src/main/java/com/reajason/javaweb/probe/config/ResponseBodyConfig.java
+++ b/generator/src/main/java/com/reajason/javaweb/probe/config/ResponseBodyConfig.java
@@ -14,5 +14,4 @@ import lombok.experimental.SuperBuilder;
 public class ResponseBodyConfig extends ProbeContentConfig {
     private String server;
     private String reqParamName;
-    private String reqHeaderName;
 }
diff --git a/generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java b/generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java
index 1073b530..a63a875f 100644
--- a/generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java
+++ b/generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java
@@ -31,17 +31,8 @@ public class ResponseBodyGenerator extends ByteBuddyShellGenerator<ResponseBodyC
 
     @Override
     protected DynamicType.Builder<?> build(ByteBuddy buddy) {
-        String name;
-        Class<?> getDataFromReqInterceptor;
-        if (probeContentConfig.getReqParamName() != null) {
-            name = probeContentConfig.getReqParamName();
-            getDataFromReqInterceptor = getDataFromReqParamInterceptor.class;
-        } else if(probeContentConfig.getReqHeaderName() != null) {
-            name = probeContentConfig.getReqHeaderName();
-            getDataFromReqInterceptor = getDataFromReqHeaderInterceptor.class;
-        }else{
-            throw new GenerationException("responseBody probeShell must set headerName or paramName");
-        }
+        String name = probeContentConfig.getReqParamName();
+        Class<?> getDataFromReqInterceptor = getDataFromReqInterceptor.class;
         Class<?> writerClass = getWriterClass();
         Class<?> runnerClass = getRunnerClass();
         return buddy.redefine(writerClass)
@@ -95,26 +86,17 @@ public class ResponseBodyGenerator extends ByteBuddyShellGenerator<ResponseBodyC
         }
     }
 
-    static class getDataFromReqHeaderInterceptor {
+    static class getDataFromReqInterceptor {
         @Advice.OnMethodExit
         public static void enter(@Advice.Argument(value = 0) Object request,
                                  @NameAnnotation String name,
                                  @Advice.Return(readOnly = false) String ret) throws Exception {
             try {
-                ret = ((String) ShellCommonUtil.invokeMethod(request, "getHeader", new Class[]{String.class}, new Object[]{name}));
-            } catch (Exception e) {
-                ret = null;
-            }
-        }
-    }
-
-    static class getDataFromReqParamInterceptor {
-        @Advice.OnMethodExit
-        public static void enter(@Advice.Argument(value = 0) Object request,
-                                 @NameAnnotation String name,
-                                 @Advice.Return(readOnly = false) String ret) throws Exception {
-            try {
-                ret = ((String) ShellCommonUtil.invokeMethod(request, "getParameter", new Class[]{String.class}, new Object[]{name}));
+                String p = (String) ShellCommonUtil.invokeMethod(request, "getParameter", new Class[]{String.class}, new Object[]{name});
+                if (p == null || p.isEmpty()) {
+                    p = (String) ShellCommonUtil.invokeMethod(request, "getHeader", new Class[]{String.class}, new Object[]{name});
+                }
+                ret = p;
             } catch (Exception e) {
                 ret = null;
             }
diff --git a/integration-test/src/test/java/com/reajason/javaweb/integration/ProbeAssertion.java b/integration-test/src/test/java/com/reajason/javaweb/integration/ProbeAssertion.java
index 8f6036bb..d72112b0 100644
--- a/integration-test/src/test/java/com/reajason/javaweb/integration/ProbeAssertion.java
+++ b/integration-test/src/test/java/com/reajason/javaweb/integration/ProbeAssertion.java
@@ -92,7 +92,7 @@ public class ProbeAssertion {
         String headerName = "X-Header";
         ResponseBodyConfig responseBodyConfig = ResponseBodyConfig.builder()
                 .server(server)
-                .reqHeaderName(headerName)
+                .reqParamName(headerName)
                 .build();
         ProbeShellResult probeResult = ProbeShellGenerator.generate(probeConfig, responseBodyConfig);
         String content = probeResult.getShellBytesBase64Str();
diff --git a/web/src/components/probeshell/main-config-card.tsx b/web/src/components/probeshell/main-config-card.tsx
index b7f98545..2766d88a 100644
--- a/web/src/components/probeshell/main-config-card.tsx
+++ b/web/src/components/probeshell/main-config-card.tsx
@@ -57,7 +57,6 @@ const PROBE_METHOD_OPTIONS = [
 
 const DEFAULT_FORM_VALUES = {
   reqParamName: "payload",
-  reqHeaderName: "X-PAYLOAD",
   sleepServer: "Tomcat",
   seconds: 5,
 } as const;
diff --git a/web/src/pages/probeshell.tsx b/web/src/pages/probeshell.tsx
index 98dc8c3d..23e1a087 100644
--- a/web/src/pages/probeshell.tsx
+++ b/web/src/pages/probeshell.tsx
@@ -53,7 +53,6 @@ export default function ProbeShellGenerator() {
       host: "",
       server: "Tomcat",
       reqParamName: "payload",
-      reqHeaderName: "X-PAYLOAD",
       seconds: 5,
       sleepServer: "Tomcat",
       shrink: true,
diff --git a/web/src/types/probeshell.ts b/web/src/types/probeshell.ts
index 94a37f5b..a8979620 100644
--- a/web/src/types/probeshell.ts
+++ b/web/src/types/probeshell.ts
@@ -25,7 +25,6 @@ export interface ProbeContentConfig {
   sleepServer?: string;
   server?: string;
   reqParamName?: string;
-  reqHeaderName?: string;
 }
 
 export interface DNSLogConfig {
@@ -40,7 +39,6 @@ export interface SleepConfig {
 export interface ResponseBodyConfig {
   server: string;
   reqParamName: string;
-  reqHeaderName: string;
 }
 
 export interface PayloadFormValues {
@@ -52,7 +50,6 @@ export interface PayloadFormValues {
   host?: string;
   server?: string;
   reqParamName?: string;
-  reqHeaderName?: string;
   sleepServer?: string;
   seconds?: number;
   packingMethod: string;
diff --git a/web/src/types/schema.ts b/web/src/types/schema.ts
index dca9c3d7..579331e6 100644
--- a/web/src/types/schema.ts
+++ b/web/src/types/schema.ts
@@ -145,7 +145,6 @@ export const probeShellFormSchema = yup.object().shape({
   host: yup.string().optional(),
   server: yup.string().optional(),
   reqParamName: yup.string().optional(),
-  reqHeaderName: yup.string().optional(),
   seconds: yup.number().optional(),
   sleepServer: yup.string().optional(),
   packingMethod: yup.string().required(),
diff --git a/web/src/utils/transformer.ts b/web/src/utils/transformer.ts
index 817cc7bb..9a0c1a24 100644
--- a/web/src/utils/transformer.ts
+++ b/web/src/utils/transformer.ts
@@ -63,7 +63,6 @@ export function transformToProbePostData(formValue: ProbeShellFormSchema) {
     sleepServer: formValue.sleepServer,
     server: formValue.server,
     reqParamName: formValue.reqParamName,
-    reqHeaderName: formValue.reqHeaderName,
   };
 
   return {