diff --git a/src/fuzzer.rs b/src/fuzzer.rs index d4e147d..f1f571e 100644 --- a/src/fuzzer.rs +++ b/src/fuzzer.rs @@ -7,10 +7,11 @@ use std::{ path::{Path, PathBuf}, process::ExitCode, rc::Rc, - time::Instant, + thread, + time::{Duration, Instant}, }; -use anyhow::{Context, Error, Result}; +use anyhow::{anyhow, Context, Error, Result}; use indexmap::IndexMap; use openapi_utils::ReferenceOrExt; use openapiv3::{OpenAPI, ReferenceOr, Response, StatusCode}; @@ -27,6 +28,8 @@ use crate::{ stats::Stats, }; +const BACKOFF_STATUS_CODES: [u16; 2] = [429, 503]; + #[derive(Debug, Deserialize, Serialize)] pub struct FuzzResult<'a> { pub payload: Payload, @@ -126,17 +129,11 @@ impl Fuzzer { &any_with::(Rc::new(ArbitraryParameters::new(operation))), |payload| { let now = Instant::now(); - let response = Fuzzer::send_request( - &self.url, - path_with_params.to_owned(), - method, - &payload, - &self.extra_headers, - &self.agent, - ) - .map_err(|e| { - TestCaseError::Fail(format!("unable to send request: {e}").into()) - })?; + let response = self + .send_request_with_backoff(path_with_params, method, &payload) + .map_err(|e| { + TestCaseError::Fail(format!("unable to send request: {e}").into()) + })?; let is_expected_response = self.is_expected_response(&response, &responses); stats.borrow_mut().times.push(now.elapsed().as_micros()); @@ -170,14 +167,55 @@ impl Fuzzer { } } + fn send_request_with_backoff( + &self, + path_with_params: &str, + method: &str, + payload: &Payload, + ) -> Result { + let max_backoff = 10; + + for backoff in 0..max_backoff { + let response = self.send_request_(path_with_params, method, payload)?; + if !BACKOFF_STATUS_CODES.contains(&response.status()) { + return Ok(response); + } + + let wait_seconds = response + .header("Retry-After") + .and_then(|s| s.parse::().ok()) + .unwrap_or(1 << backoff); + thread::sleep(Duration::from_millis(wait_seconds * 1000)); + } + + Err(anyhow!("max backoff threshold reached")) + } + + fn send_request_( + &self, + path_with_params: &str, + method: &str, + payload: &Payload, + ) -> Result { + Fuzzer::send_request( + &self.url, + path_with_params, + method, + payload, + &self.extra_headers, + &self.agent, + ) + } + pub fn send_request( url: &Url, - mut path_with_params: String, + path_with_params: &str, method: &str, payload: &Payload, extra_headers: &HashMap, agent: &Agent, ) -> Result { + let mut path_with_params = path_with_params.to_owned(); for (name, value) in payload.path_params().iter() { path_with_params = path_with_params.replace(&format!("{{{name}}}"), value); } diff --git a/src/main.rs b/src/main.rs index 9107618..031ced1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -172,7 +172,7 @@ fn main() -> Result { let response = Fuzzer::send_request( &args.url.into(), - result.path.to_owned(), + result.path, result.method, &result.payload, &args.header.into_iter().map(Into::into).collect(),