feat: support probe get payload from header by default

This commit is contained in:
ReaJason
2025-11-20 00:36:25 +08:00
parent 3838d3d27a
commit fb5b39a4c6
9 changed files with 9 additions and 37 deletions
@@ -21,7 +21,6 @@ public class ProbeShellGenerateRequest {
private String server; private String server;
private String sleepServer; private String sleepServer;
private String reqParamName; private String reqParamName;
private String reqHeaderName;
} }
public ProbeContentConfig parseProbeContentConfig() { public ProbeContentConfig parseProbeContentConfig() {
@@ -35,7 +34,6 @@ public class ProbeShellGenerateRequest {
.build(); .build();
case ResponseBody -> ResponseBodyConfig.builder() case ResponseBody -> ResponseBodyConfig.builder()
.reqParamName(probeContentConfig.reqParamName) .reqParamName(probeContentConfig.reqParamName)
.reqHeaderName(probeContentConfig.reqHeaderName)
.server(probeContentConfig.server) .server(probeContentConfig.server)
.build(); .build();
default -> throw new UnsupportedOperationException("unknown probe method: " + probeConfig.getProbeMethod()); default -> throw new UnsupportedOperationException("unknown probe method: " + probeConfig.getProbeMethod());
@@ -14,5 +14,4 @@ import lombok.experimental.SuperBuilder;
public class ResponseBodyConfig extends ProbeContentConfig { public class ResponseBodyConfig extends ProbeContentConfig {
private String server; private String server;
private String reqParamName; private String reqParamName;
private String reqHeaderName;
} }
@@ -31,17 +31,8 @@ public class ResponseBodyGenerator extends ByteBuddyShellGenerator<ResponseBodyC
@Override @Override
protected DynamicType.Builder<?> build(ByteBuddy buddy) { protected DynamicType.Builder<?> build(ByteBuddy buddy) {
String name; String name = probeContentConfig.getReqParamName();
Class<?> getDataFromReqInterceptor; Class<?> getDataFromReqInterceptor = getDataFromReqInterceptor.class;
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");
}
Class<?> writerClass = getWriterClass(); Class<?> writerClass = getWriterClass();
Class<?> runnerClass = getRunnerClass(); Class<?> runnerClass = getRunnerClass();
return buddy.redefine(writerClass) return buddy.redefine(writerClass)
@@ -95,26 +86,17 @@ public class ResponseBodyGenerator extends ByteBuddyShellGenerator<ResponseBodyC
} }
} }
static class getDataFromReqHeaderInterceptor { static class getDataFromReqInterceptor {
@Advice.OnMethodExit @Advice.OnMethodExit
public static void enter(@Advice.Argument(value = 0) Object request, public static void enter(@Advice.Argument(value = 0) Object request,
@NameAnnotation String name, @NameAnnotation String name,
@Advice.Return(readOnly = false) String ret) throws Exception { @Advice.Return(readOnly = false) String ret) throws Exception {
try { try {
ret = ((String) ShellCommonUtil.invokeMethod(request, "getHeader", new Class[]{String.class}, new Object[]{name})); String p = (String) ShellCommonUtil.invokeMethod(request, "getParameter", new Class[]{String.class}, new Object[]{name});
} catch (Exception e) { if (p == null || p.isEmpty()) {
ret = null; p = (String) ShellCommonUtil.invokeMethod(request, "getHeader", new Class[]{String.class}, new Object[]{name});
} }
} ret = p;
}
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}));
} catch (Exception e) { } catch (Exception e) {
ret = null; ret = null;
} }
@@ -92,7 +92,7 @@ public class ProbeAssertion {
String headerName = "X-Header"; String headerName = "X-Header";
ResponseBodyConfig responseBodyConfig = ResponseBodyConfig.builder() ResponseBodyConfig responseBodyConfig = ResponseBodyConfig.builder()
.server(server) .server(server)
.reqHeaderName(headerName) .reqParamName(headerName)
.build(); .build();
ProbeShellResult probeResult = ProbeShellGenerator.generate(probeConfig, responseBodyConfig); ProbeShellResult probeResult = ProbeShellGenerator.generate(probeConfig, responseBodyConfig);
String content = probeResult.getShellBytesBase64Str(); String content = probeResult.getShellBytesBase64Str();
@@ -57,7 +57,6 @@ const PROBE_METHOD_OPTIONS = [
const DEFAULT_FORM_VALUES = { const DEFAULT_FORM_VALUES = {
reqParamName: "payload", reqParamName: "payload",
reqHeaderName: "X-PAYLOAD",
sleepServer: "Tomcat", sleepServer: "Tomcat",
seconds: 5, seconds: 5,
} as const; } as const;
-1
View File
@@ -53,7 +53,6 @@ export default function ProbeShellGenerator() {
host: "", host: "",
server: "Tomcat", server: "Tomcat",
reqParamName: "payload", reqParamName: "payload",
reqHeaderName: "X-PAYLOAD",
seconds: 5, seconds: 5,
sleepServer: "Tomcat", sleepServer: "Tomcat",
shrink: true, shrink: true,
-3
View File
@@ -25,7 +25,6 @@ export interface ProbeContentConfig {
sleepServer?: string; sleepServer?: string;
server?: string; server?: string;
reqParamName?: string; reqParamName?: string;
reqHeaderName?: string;
} }
export interface DNSLogConfig { export interface DNSLogConfig {
@@ -40,7 +39,6 @@ export interface SleepConfig {
export interface ResponseBodyConfig { export interface ResponseBodyConfig {
server: string; server: string;
reqParamName: string; reqParamName: string;
reqHeaderName: string;
} }
export interface PayloadFormValues { export interface PayloadFormValues {
@@ -52,7 +50,6 @@ export interface PayloadFormValues {
host?: string; host?: string;
server?: string; server?: string;
reqParamName?: string; reqParamName?: string;
reqHeaderName?: string;
sleepServer?: string; sleepServer?: string;
seconds?: number; seconds?: number;
packingMethod: string; packingMethod: string;
-1
View File
@@ -145,7 +145,6 @@ export const probeShellFormSchema = yup.object().shape({
host: yup.string().optional(), host: yup.string().optional(),
server: yup.string().optional(), server: yup.string().optional(),
reqParamName: yup.string().optional(), reqParamName: yup.string().optional(),
reqHeaderName: yup.string().optional(),
seconds: yup.number().optional(), seconds: yup.number().optional(),
sleepServer: yup.string().optional(), sleepServer: yup.string().optional(),
packingMethod: yup.string().required(), packingMethod: yup.string().required(),
-1
View File
@@ -63,7 +63,6 @@ export function transformToProbePostData(formValue: ProbeShellFormSchema) {
sleepServer: formValue.sleepServer, sleepServer: formValue.sleepServer,
server: formValue.server, server: formValue.server,
reqParamName: formValue.reqParamName, reqParamName: formValue.reqParamName,
reqHeaderName: formValue.reqHeaderName,
}; };
return { return {