diff --git a/src/banner/container.rs b/src/banner/container.rs index 9b34263..3d36d66 100644 --- a/src/banner/container.rs +++ b/src/banner/container.rs @@ -1,8 +1,8 @@ use super::entry::BannerEntry; -use crate::event_handlers::Handles; use crate::{ config::Configuration, - utils::{make_request, status_colorizer}, + event_handlers::Handles, + utils::{logged_request, status_colorizer}, VERSION, }; use anyhow::{bail, Result}; @@ -364,15 +364,8 @@ by Ben "epi" Risher {} ver: {}"#, let api_url = Url::parse(url)?; - let response = make_request( - &handles.config.client, - &api_url, - handles.config.output_level, - handles.stats.tx.clone(), - ) - .await?; - - let body = response.text().await?; + let result = logged_request(&api_url, handles.clone()).await?; + let body = result.text().await?; let json_response: Value = serde_json::from_str(&body)?; diff --git a/src/event_handlers/outputs.rs b/src/event_handlers/outputs.rs index ea19efc..8df1cfa 100644 --- a/src/event_handlers/outputs.rs +++ b/src/event_handlers/outputs.rs @@ -210,7 +210,7 @@ impl TermOutHandler { if self.config.replay_client.is_some() && should_process_response { // replay proxy specified/client created and this response's status code is one that - // should be replayed + // should be replayed; not using logged_request due to replay proxy client make_request( self.config.replay_client.as_ref().unwrap(), &resp.url(), diff --git a/src/event_handlers/statistics.rs b/src/event_handlers/statistics.rs index 98f7b16..82b805f 100644 --- a/src/event_handlers/statistics.rs +++ b/src/event_handlers/statistics.rs @@ -89,6 +89,7 @@ impl StatsHandler { } Command::AddStatus(status) => { self.stats.add_status_code(status); + self.increment_bar(); } Command::AddRequest => { diff --git a/src/extractor/container.rs b/src/extractor/container.rs index a5263a8..d158b63 100644 --- a/src/extractor/container.rs +++ b/src/extractor/container.rs @@ -12,7 +12,7 @@ use crate::{ StatField::{LinksExtracted, TotalExpected}, }, url::FeroxUrl, - utils::make_request, + utils::{logged_request, make_request}, }; use anyhow::{bail, Context, Result}; use reqwest::{StatusCode, Url}; @@ -303,13 +303,7 @@ impl<'a> Extractor<'a> { } // make the request and store the response - let new_response = make_request( - &self.handles.config.client, - &new_url, - self.handles.config.output_level, - self.handles.stats.tx.clone(), - ) - .await?; + let new_response = logged_request(&new_url, self.handles.clone()).await?; let new_ferox_response = FeroxResponse::from(new_response, true, self.handles.config.output_level).await; @@ -384,6 +378,7 @@ impl<'a> Extractor<'a> { let mut url = Url::parse(&self.url)?; url.set_path("/robots.txt"); // overwrite existing path with /robots.txt + // purposefully not using logged_request here due to using the special client let response = make_request( &client, &url, @@ -391,6 +386,7 @@ impl<'a> Extractor<'a> { self.handles.stats.tx.clone(), ) .await?; + let ferox_response = FeroxResponse::from(response, true, self.handles.config.output_level).await; diff --git a/src/filters/init.rs b/src/filters/init.rs index 9b0cf0c..35440b0 100644 --- a/src/filters/init.rs +++ b/src/filters/init.rs @@ -5,7 +5,7 @@ use crate::{ event_handlers::Handles, response::FeroxResponse, skip_fail, - utils::{fmt_err, make_request}, + utils::{fmt_err, logged_request}, Command::AddFilter, SIMILARITY_THRESHOLD, }; @@ -72,15 +72,7 @@ pub async fn initialize(handles: Arc) -> Result<()> { let url = skip_fail!(Url::parse(&similarity_filter)); // attempt to request the given url - let resp = skip_fail!( - make_request( - &handles.config.client, - &url, - handles.config.output_level, - handles.stats.tx.clone() - ) - .await - ); + let resp = skip_fail!(logged_request(&url, handles.clone()).await); // if successful, create a filter based on the response's body let fr = FeroxResponse::from(resp, true, handles.config.output_level).await; diff --git a/src/heuristics.rs b/src/heuristics.rs index 6f1e4da..e64d58a 100644 --- a/src/heuristics.rs +++ b/src/heuristics.rs @@ -12,7 +12,7 @@ use crate::{ response::FeroxResponse, skip_fail, url::FeroxUrl, - utils::{ferox_print, fmt_err, make_request, status_colorizer}, + utils::{ferox_print, fmt_err, logged_request, status_colorizer}, }; /// length of a standard UUID, used when determining wildcard responses @@ -158,13 +158,7 @@ impl HeuristicTests { let unique_str = self.unique_string(length); let nonexistent_url = target.format(&unique_str, None)?; - let response = make_request( - &self.handles.config.client, - &nonexistent_url.to_owned(), - self.handles.config.output_level, - self.handles.stats.tx.clone(), - ) - .await?; + let response = logged_request(&nonexistent_url.to_owned(), self.handles.clone()).await?; if self .handles @@ -215,13 +209,8 @@ impl HeuristicTests { for target_url in target_urls { let url = FeroxUrl::from_string(&target_url, self.handles.clone()); let request = skip_fail!(url.format("", None)); - let result = make_request( - &self.handles.config.client, - &request, - self.handles.config.output_level, - self.handles.stats.tx.clone(), - ) - .await; + + let result = logged_request(&request, self.handles.clone()).await; match result { Ok(_) => { diff --git a/src/lib.rs b/src/lib.rs index 2424ae4..b8ee387 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -57,7 +57,10 @@ pub const DEFAULT_WORDLIST: &str = "/usr/share/seclists/Discovery/Web-Content/raft-medium-directories.txt"; /// Number of milliseconds to wait between polls of `PAUSE_SCAN` when user pauses a scan -pub(crate) static SLEEP_DURATION: u64 = 500; +pub(crate) const SLEEP_DURATION: u64 = 500; + +/// The percentage of requests as errors it takes to be deemed too high +pub const HIGH_ERROR_RATIO: f64 = 0.90; /// Default list of status codes to report /// diff --git a/src/scan_manager/scan.rs b/src/scan_manager/scan.rs index ca624c1..c427e67 100644 --- a/src/scan_manager/scan.rs +++ b/src/scan_manager/scan.rs @@ -2,6 +2,7 @@ use super::*; use crate::{ config::OutputLevel, progress::{add_bar, BarType}, + scanner::PolicyTrigger, }; use anyhow::Result; use console::style; @@ -14,6 +15,7 @@ use std::{ sync::{Arc, Mutex}, }; +use std::sync::atomic::{AtomicUsize, Ordering}; use tokio::{sync, task::JoinHandle}; use uuid::Uuid; @@ -49,6 +51,15 @@ pub struct FeroxScan { /// whether or not the user passed --silent|--quiet on the command line pub(super) output_level: OutputLevel, + + /// todo + pub(super) status_403s: AtomicUsize, + + /// todo + pub(super) status_429s: AtomicUsize, + + /// todo + pub(super) errors: AtomicUsize, } /// Default implementation for FeroxScan @@ -67,6 +78,9 @@ impl Default for FeroxScan { progress_bar: Mutex::new(None), scan_type: ScanType::File, output_level: Default::default(), + errors: Default::default(), + status_429s: Default::default(), + status_403s: Default::default(), } } } @@ -75,8 +89,10 @@ impl Default for FeroxScan { impl FeroxScan { /// Stop a currently running scan pub async fn abort(&self) -> Result<()> { + println!("IN ABORT: {:?}", self); let mut guard = self.task.lock().await; + println!("IN ABORT: {:?}", self); if guard.is_some() { if let Some(task) = std::mem::replace(&mut *guard, None) { task.abort(); @@ -217,6 +233,42 @@ impl FeroxScan { log::trace!("exit join({:?})", self); } + /// increment the value in question by 1 + pub(super) fn add_403(&self) { + self.status_403s.fetch_add(1, Ordering::Relaxed); + } + + /// increment the value in question by 1 + pub(super) fn add_429(&self) { + self.status_429s.fetch_add(1, Ordering::Relaxed); + } + + /// increment the value in question by 1 + pub(super) fn add_error(&self) { + self.errors.fetch_add(1, Ordering::SeqCst); + } + + /// simple wrapper to call the appropriate getter based on the given PolicyTrigger + pub fn num_errors(&self, trigger: PolicyTrigger) -> usize { + return match trigger { + PolicyTrigger::Status403 => self.status_403s(), + PolicyTrigger::Status429 => self.status_429s(), + PolicyTrigger::Errors => self.errors(), + }; + } + + /// return the number of errors seen by this scan + fn errors(&self) -> usize { + self.errors.load(Ordering::Relaxed) + } + /// return the number of 403s seen by this scan + fn status_403s(&self) -> usize { + self.status_403s.load(Ordering::Relaxed) + } + /// return the number of 429s seen by this scan + fn status_429s(&self) -> usize { + self.status_429s.load(Ordering::Relaxed) + } } /// Display implementation @@ -360,3 +412,34 @@ impl Default for ScanStatus { Self::NotStarted } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + /// ensure that num_errors returns the correct values for the given PolicyTrigger + /// + /// covers tests for add_[403,429,error] and the related getters in addition to num_errors + fn num_errors_returns_correct_values() { + let scan = FeroxScan::new( + "http://localhost", + ScanType::Directory, + ScanOrder::Latest, + 1000, + OutputLevel::Default, + None, + ); + + scan.add_error(); + scan.add_403(); + scan.add_403(); + scan.add_429(); + scan.add_429(); + scan.add_429(); + + assert_eq!(scan.num_errors(PolicyTrigger::Errors), 1); + assert_eq!(scan.num_errors(PolicyTrigger::Status403), 2); + assert_eq!(scan.num_errors(PolicyTrigger::Status429), 3); + } +} diff --git a/src/scan_manager/scan_container.rs b/src/scan_manager/scan_container.rs index 6a50cf3..755ad2c 100644 --- a/src/scan_manager/scan_container.rs +++ b/src/scan_manager/scan_container.rs @@ -9,6 +9,7 @@ use crate::{ SLEEP_DURATION, }; use anyhow::Result; +use reqwest::StatusCode; use serde::{ser::SerializeSeq, Serialize, Serializer}; use std::{ convert::TryInto, @@ -161,6 +162,28 @@ impl FeroxScans { None } + /// add one to either 403 or 429 tracker in the scan related to the given url + pub fn increment_status_code(&self, url: &str, code: StatusCode) { + if let Some(scan) = self.get_scan_by_url(url) { + match code { + StatusCode::TOO_MANY_REQUESTS => { + scan.add_429(); + } + StatusCode::FORBIDDEN => { + scan.add_403(); + } + _ => {} + } + } + } + + /// add one to either 403 or 429 tracker in the scan related to the given url + pub fn increment_error(&self, url: &str) { + if let Some(scan) = self.get_scan_by_url(url) { + scan.add_error(); + } + } + /// Print all FeroxScans of type Directory /// /// Example: diff --git a/src/scan_manager/tests.rs b/src/scan_manager/tests.rs index 2a284d8..6d56484 100644 --- a/src/scan_manager/tests.rs +++ b/src/scan_manager/tests.rs @@ -438,9 +438,12 @@ fn feroxscan_display() { scan_type: Default::default(), num_requests: 0, output_level: OutputLevel::Default, + status_403s: Default::default(), + status_429s: Default::default(), status: Default::default(), task: tokio::sync::Mutex::new(None), progress_bar: std::sync::Mutex::new(None), + errors: Default::default(), }; let not_started = format!("{}", scan); @@ -478,11 +481,14 @@ async fn ferox_scan_abort() { scan_type: Default::default(), num_requests: 0, output_level: OutputLevel::Default, + status_403s: Default::default(), + status_429s: Default::default(), status: std::sync::Mutex::new(ScanStatus::Running), task: tokio::sync::Mutex::new(Some(tokio::spawn(async move { sleep(Duration::from_millis(SLEEP_DURATION * 2)); }))), progress_bar: std::sync::Mutex::new(None), + errors: Default::default(), }; scan.abort().await.unwrap(); diff --git a/src/scanner/mod.rs b/src/scanner/mod.rs index e28e0bf..619f1ea 100644 --- a/src/scanner/mod.rs +++ b/src/scanner/mod.rs @@ -6,3 +6,4 @@ mod tests; pub use self::container::{FeroxScanner, RESPONSES}; pub use self::init::initialize; +pub use self::utils::PolicyTrigger; diff --git a/src/scanner/utils.rs b/src/scanner/utils.rs index 6952e3b..0e612d1 100644 --- a/src/scanner/utils.rs +++ b/src/scanner/utils.rs @@ -6,16 +6,33 @@ use crate::{ Handles, }, extractor::{ExtractionTarget::ResponseBody, ExtractorBuilder}, + progress::PROGRESS_PRINTER, response::FeroxResponse, statistics::StatError::Other, url::FeroxUrl, - utils::make_request, + utils::logged_request, + HIGH_ERROR_RATIO, }; use anyhow::Result; use leaky_bucket::LeakyBucket; +use std::ops::Index; +use std::sync::atomic::Ordering; use std::{cmp::max, sync::Arc}; use tokio::{sync::oneshot, time::Duration}; +#[derive(Copy, Clone, PartialEq, Debug)] +/// represents different situations where different criteria can trigger auto-tune/bail behavior +pub enum PolicyTrigger { + /// excessive 403 trigger + Status403, + + /// excessive 429 trigger + Status429, + + /// excessive general errors + Errors, +} + /// Makes multiple requests based on the presence of extensions pub(super) struct Requester { /// handles to handlers and config @@ -52,10 +69,10 @@ impl Requester { None }; - // let policy = scanner.handles.config.config.policy; todo + let policy = scanner.handles.config.requester_policy; Ok(Self { - policy: RequesterPolicy::Default, // todo replace with dynamic from config + policy, rate_limiter, handles: scanner.handles.clone(), target_url: scanner.target_url.to_owned(), @@ -68,6 +85,98 @@ impl Requester { Ok(()) } + /// determine whether or not a policy needs to be enforce + /// + /// criteria: + /// - threads * 2 for general errors (timeouts etc) + /// - 90% of requests are 403 + /// - 30% of requests are 429 + fn should_enforce_policy(&self) -> Option { + let requests = self.handles.stats.data.requests.load(Ordering::Relaxed); + + if requests < max(self.handles.config.threads, 50) { + // check whether at least a full round of threads has made requests or 50 (default # of + // threads), whichever is higher + return None; + } + + let errors = self.handles.stats.data.errors.load(Ordering::Relaxed); + let s403s = self.handles.stats.data.status_403s.load(Ordering::Relaxed); + let s429s = self.handles.stats.data.status_429s.load(Ordering::Relaxed); + + let threshold = self.handles.config.threads * 2; + if errors >= threshold { + // general errors should not exceed the given threshold + return Some(PolicyTrigger::Errors); + } + + let ratio_403s = s403s as f64 / requests as f64; + if ratio_403s >= HIGH_ERROR_RATIO { + // almost exclusively 403 + return Some(PolicyTrigger::Status403); + } + + let ratio_429s = s429s as f64 / requests as f64; + if ratio_429s >= HIGH_ERROR_RATIO / 3.0 { + // high # of 429 responses + return Some(PolicyTrigger::Status429); + } + + None + } + + /// enforce auto-tune policy + fn tune(&self, _trigger: PolicyTrigger) {} + + /// enforce auto-bail policy + async fn bail(&self, trigger: PolicyTrigger) -> Result<()> { + let scans = self.handles.ferox_scans()?; + + let mut scan_tuples = vec![]; + + { + if let Ok(guard) = scans.scans.read() { + for (i, scan) in guard.iter().enumerate() { + PROGRESS_PRINTER.println(format!( + "{} {}", + scan.is_active(), + scan.num_errors(trigger) + )); + if scan.is_active() && scan.num_errors(trigger) > 0 { + // only active scans that have at least 1 error + + scan_tuples.push((i, scan.num_errors(trigger))); + } + } + } + } + + if scan_tuples.len() == 0 { + return Ok(()); + } + + // sort by number of errors + scan_tuples.sort_unstable_by(|x, y| y.1.cmp(&x.1)); + + for (idx, _errors) in scan_tuples { + let scan = if let Ok(guard) = scans.scans.read() { + guard.index(idx).clone() + } else { + // todo think about logging + continue; + }; + + if scan.is_active() { + scan.abort() + .await + .unwrap_or_else(|e| log::warn!("Could not bail on scan: {}", e)); + break; + } + } + + Ok(()) + } + /// Wrapper for make_request /// /// Attempts recursion when appropriate and sends Responses to the output handler for processing @@ -86,16 +195,21 @@ impl Requester { } } - let response = make_request( - &self.handles.config.client, - &url, - self.handles.config.output_level, - self.handles.stats.tx.clone(), - ) - .await?; + let response = logged_request(&url, self.handles.clone()).await?; - // todo this is where bail should go, tune can probably just set a limiter if one isn't - // already present + match self.policy { + RequesterPolicy::AutoTune => { + if let Some(trigger) = self.should_enforce_policy() { + self.tune(trigger); + } + } + RequesterPolicy::AutoBail => { + if let Some(trigger) = self.should_enforce_policy() { + self.bail(trigger).await?; // todo may or may not be right to bubble up + } + } + RequesterPolicy::Default => {} + } // response came back without error, convert it to FeroxResponse let ferox_response = @@ -141,3 +255,272 @@ impl Requester { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::OutputLevel; + use crate::scan_manager::ScanStatus; + use crate::statistics::StatError; + use crate::{ + config::Configuration, + event_handlers::{FiltersHandler, ScanHandler, StatsHandler, Tasks, TermOutHandler}, + filters, + }; + use crate::{ + scan_manager::FeroxScan, + scan_manager::{ScanOrder, ScanType}, + }; + use reqwest::StatusCode; + + /// helper to setup a realistic requester test + async fn setup_requester_test(config: Option>) -> (Arc, Tasks) { + // basically C&P from main::wrapped_main, can look there for comments etc if needed + let configuration = config.unwrap_or_else(|| Arc::new(Configuration::new().unwrap())); + + let (stats_task, stats_handle) = StatsHandler::initialize(configuration.clone()); + let (filters_task, filters_handle) = FiltersHandler::initialize(); + let (out_task, out_handle) = + TermOutHandler::initialize(configuration.clone(), stats_handle.tx.clone()); + + let handles = Arc::new(Handles::new( + stats_handle, + filters_handle, + out_handle, + configuration.clone(), + )); + + let (scan_task, scan_handle) = ScanHandler::initialize(handles.clone()); + + handles.set_scan_handle(scan_handle); + filters::initialize(handles.clone()).await.unwrap(); + + let tasks = Tasks::new(out_task, stats_task, filters_task, scan_task); + + (handles, tasks) + } + + /// helper to stay DRY + async fn increment_errors(handles: Arc, num_errors: usize) { + for _ in 0..num_errors { + handles + .stats + .send(Command::AddError(StatError::Other)) + .unwrap(); + } + + handles.stats.sync().await.unwrap(); + } + + /// helper to stay DRY + async fn increment_scan_errors(handles: Arc, url: &str, num_errors: usize) { + let scans = handles.ferox_scans().unwrap(); + for _ in 0..num_errors { + scans.increment_error(url); + } + } + + /// helper to stay DRY + async fn increment_scan_status_codes( + handles: Arc, + url: &str, + code: StatusCode, + num_errors: usize, + ) { + let scans = handles.ferox_scans().unwrap(); + for _ in 0..num_errors { + scans.increment_status_code(url, code); + } + } + + /// helper to stay DRY + async fn increment_status_codes(handles: Arc, num_codes: usize, code: StatusCode) { + for _ in 0..num_codes { + handles.stats.send(Command::AddStatus(code)).unwrap(); + } + + handles.stats.sync().await.unwrap(); + } + + /// helper to stay DRY + fn get_requests(handles: Arc) -> usize { + handles.stats.data.requests.load(Ordering::Relaxed) + } + + async fn create_scan( + handles: Arc, + url: &str, + num_errors: usize, + trigger: PolicyTrigger, + ) -> Arc { + let scan = FeroxScan::new( + url, + ScanType::Directory, + ScanOrder::Initial, + 1000, + OutputLevel::Default, + None, + ); + + scan.set_status(ScanStatus::Running).unwrap(); + scan.progress_bar(); // create a new pb + + let scans = handles.ferox_scans().unwrap(); + scans.insert(scan.clone()); + + match trigger { + PolicyTrigger::Status403 => { + increment_scan_status_codes( + handles.clone(), + url, + StatusCode::FORBIDDEN, + num_errors, + ) + .await; + } + PolicyTrigger::Status429 => { + increment_scan_status_codes( + handles.clone(), + url, + StatusCode::TOO_MANY_REQUESTS, + num_errors, + ) + .await; + } + PolicyTrigger::Errors => { + increment_scan_errors(handles.clone(), url, num_errors).await; + } + } + + assert_eq!(scan.num_errors(trigger), num_errors); + + scan + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + /// should_enforce_policy should return false when # of requests is < threads; also when < 50 + async fn should_enforce_policy_returns_false_on_not_enough_requests_seen() { + let (handles, _) = setup_requester_test(None).await; + + let requester = Requester { + handles, + target_url: "http://localhost".to_string(), + rate_limiter: None, + policy: Default::default(), + }; + + increment_errors(requester.handles.clone(), 49).await; + // 49 errors is false because we haven't hit the min threshold + assert_eq!(get_requests(requester.handles.clone()), 49); + assert_eq!(requester.should_enforce_policy(), None); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + /// should_enforce_policy should return true when # of requests is >= 50 and errors >= threads * 2 + async fn should_enforce_policy_returns_true_on_error_times_threads() { + let mut config = Configuration::new().unwrap_or_default(); + config.threads = 50; + + let (handles, _) = setup_requester_test(Some(Arc::new(config))).await; + + let requester = Requester { + handles, + target_url: "http://localhost".to_string(), + rate_limiter: None, + policy: Default::default(), + }; + + increment_errors(requester.handles.clone(), 50).await; + assert_eq!(requester.should_enforce_policy(), None); + increment_errors(requester.handles.clone(), 50).await; + assert_eq!( + requester.should_enforce_policy(), + Some(PolicyTrigger::Errors) + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + /// should_enforce_policy should return true when # of requests is >= 50 and 403s >= 45 (90%) + async fn should_enforce_policy_returns_true_on_excessive_403s() { + let (handles, _) = setup_requester_test(None).await; + + let requester = Requester { + handles, + target_url: "http://localhost".to_string(), + rate_limiter: None, + policy: Default::default(), + }; + + increment_status_codes(requester.handles.clone(), 45, StatusCode::FORBIDDEN).await; + assert_eq!(requester.should_enforce_policy(), None); + increment_status_codes(requester.handles.clone(), 5, StatusCode::OK).await; + assert_eq!( + requester.should_enforce_policy(), + Some(PolicyTrigger::Status403) + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + /// should_enforce_policy should return true when # of requests is >= 50 and errors >= 45 (90%) + async fn should_enforce_policy_returns_true_on_excessive_429s() { + let mut config = Configuration::new().unwrap_or_default(); + config.threads = 50; + + let (handles, _) = setup_requester_test(Some(Arc::new(config))).await; + + let requester = Requester { + handles, + target_url: "http://localhost".to_string(), + rate_limiter: None, + policy: Default::default(), + }; + + increment_status_codes(requester.handles.clone(), 15, StatusCode::TOO_MANY_REQUESTS).await; + assert_eq!(requester.should_enforce_policy(), None); + increment_status_codes(requester.handles.clone(), 35, StatusCode::OK).await; + assert_eq!( + requester.should_enforce_policy(), + Some(PolicyTrigger::Status429) + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + /// bail should return call abort on the scan with the most errors + async fn bail_calls_abort_on_highest_errored_feroxscan() { + let url = "http://one"; + + let (handles, _) = setup_requester_test(None).await; + + let scan_one = create_scan(handles.clone(), url, 10, PolicyTrigger::Errors).await; + let scan_two = create_scan(handles.clone(), "http://two", 14, PolicyTrigger::Errors).await; + let scan_three = + create_scan(handles.clone(), "http://three", 4, PolicyTrigger::Errors).await; + let scan_four = create_scan(handles.clone(), "http://four", 7, PolicyTrigger::Errors).await; + + // set up a fake JoinHandle for the scan that's expected to have .abort called on it + // the reason being if there's no task, the status is never updated, so can't be checked + let dummy_task = + tokio::spawn(async move { tokio::time::sleep(Duration::new(15, 0)).await }); + scan_two.set_task(dummy_task).await.unwrap(); + + assert!(scan_one.is_active()); + assert!(scan_two.is_active()); + + let scans = handles.ferox_scans().unwrap(); + assert_eq!(scans.get_active_scans().len(), 4); + + let requester = Requester { + handles, + target_url: url.to_string(), + rate_limiter: None, + policy: Default::default(), + }; + + requester.bail(PolicyTrigger::Errors).await.unwrap(); + assert_eq!(scans.get_active_scans().len(), 3); + assert!(scan_one.is_active()); + assert!(scan_three.is_active()); + assert!(scan_four.is_active()); + assert!(!scan_two.is_active()); + } +} diff --git a/src/statistics/container.rs b/src/statistics/container.rs index cf8469a..f372666 100644 --- a/src/statistics/container.rs +++ b/src/statistics/container.rs @@ -29,7 +29,7 @@ pub struct Stats { timeouts: AtomicUsize, /// tracker for total number of requests sent by the client - requests: AtomicUsize, + pub(crate) requests: AtomicUsize, /// tracker for total number of requests expected to send if the scan runs to completion /// @@ -42,7 +42,7 @@ pub struct Stats { total_expected: AtomicUsize, /// tracker for total number of errors encountered by the client - errors: AtomicUsize, + pub(crate) errors: AtomicUsize, /// tracker for overall number of 2xx status codes seen by the client successes: AtomicUsize, @@ -58,7 +58,7 @@ pub struct Stats { /// tracker for number of scans performed, this directly equates to number of directories /// recursed into and affects the total number of expected requests - total_scans: AtomicUsize, + pub(crate) total_scans: AtomicUsize, /// tracker for initial number of requested targets initial_targets: AtomicUsize, @@ -80,10 +80,10 @@ pub struct Stats { status_401s: AtomicUsize, /// tracker for overall number of 403s seen by the client - status_403s: AtomicUsize, + pub(crate) status_403s: AtomicUsize, /// tracker for overall number of 429s seen by the client - status_429s: AtomicUsize, + pub(crate) status_429s: AtomicUsize, /// tracker for overall number of 500s seen by the client status_500s: AtomicUsize, @@ -222,10 +222,6 @@ impl Stats { StatError::Timeout => { atomic_increment!(self.timeouts); } - StatError::Status403 => { - atomic_increment!(self.status_403s); - atomic_increment!(self.client_errors); - } StatError::UrlFormat => { atomic_increment!(self.url_format_errors); } @@ -238,9 +234,7 @@ impl Stats { StatError::Request => { atomic_increment!(self.request_errors); } - StatError::Other => { - atomic_increment!(self.errors); - } + _ => {} // no need to hit Other as we always increment self.errors anyway } } @@ -248,7 +242,7 @@ impl Stats { /// /// Implies incrementing: /// - requests - /// - status_403s (when code is 403) + /// - appropriate status_* codes /// - errors (when code is [45]xx) pub fn add_status_code(&self, status: StatusCode) { self.add_request(); @@ -264,9 +258,6 @@ impl Stats { } match status { - StatusCode::FORBIDDEN => { - atomic_increment!(self.status_403s); - } StatusCode::OK => { atomic_increment!(self.status_200s); } @@ -279,6 +270,9 @@ impl Stats { StatusCode::UNAUTHORIZED => { atomic_increment!(self.status_401s); } + StatusCode::FORBIDDEN => { + atomic_increment!(self.status_403s); + } StatusCode::TOO_MANY_REQUESTS => { atomic_increment!(self.status_429s); } @@ -435,30 +429,6 @@ mod tests { Ok(()) } - #[tokio::test(flavor = "multi_thread", worker_threads = 1)] - /// when sent StatCommand::AddRequest, stats object should reflect the change - /// - /// incrementing a 403 (tracked in status_403s) should also increment: - /// - errors - /// - requests - /// - client_errors - async fn statistics_handler_increments_403() { - let (task, handle) = setup_stats_test(); - - let err = Command::AddError(StatError::Status403); - let err2 = Command::AddError(StatError::Status403); - - handle.tx.send(err).unwrap_or_default(); - handle.tx.send(err2).unwrap_or_default(); - - teardown_stats_test(handle.tx.clone(), task).await; - - assert_eq!(handle.data.errors.load(Ordering::Relaxed), 2); - assert_eq!(handle.data.requests.load(Ordering::Relaxed), 2); - assert_eq!(handle.data.status_403s.load(Ordering::Relaxed), 2); - assert_eq!(handle.data.client_errors.load(Ordering::Relaxed), 2); - } - #[tokio::test(flavor = "multi_thread", worker_threads = 1)] /// when sent StatCommand::AddRequest, stats object should reflect the change /// diff --git a/src/statistics/error.rs b/src/statistics/error.rs index 374e119..fbeca51 100644 --- a/src/statistics/error.rs +++ b/src/statistics/error.rs @@ -1,9 +1,6 @@ #[derive(Debug, Copy, Clone)] /// Enum variants used to inform the `StatCommand` protocol what `Stats` fields should be updated pub enum StatError { - /// Represents a 403 response code - Status403, - /// Represents a timeout error Timeout, diff --git a/src/utils.rs b/src/utils.rs index caed60c..a8b74ae 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -7,12 +7,16 @@ use rlimit::{getrlimit, setrlimit, Resource, Rlim}; use std::{ fs, io::{self, BufWriter, Write}, + sync::Arc, }; use tokio::sync::mpsc::UnboundedSender; use crate::{ config::OutputLevel, - event_handlers::Command::{self, AddError, AddStatus}, + event_handlers::{ + Command::{self, AddError, AddStatus}, + Handles, + }, progress::PROGRESS_PRINTER, send_command, statistics::StatError::{Connection, Other, Redirection, Request, Timeout}, @@ -81,6 +85,29 @@ pub fn ferox_print(msg: &str, bar: &ProgressBar) { } } +/// wrapper for make_request used to pass error/response codes to FeroxScans for per-scan stats +/// tracking of information related to auto-tune/bail +pub async fn logged_request(url: &Url, handles: Arc) -> Result { + let client = &handles.config.client; + let level = handles.config.output_level; + let tx_stats = handles.stats.tx.clone(); + + let response = make_request(client, url, level, tx_stats).await; + + let scans = handles.ferox_scans()?; + + match response { + Ok(resp) => { + scans.increment_status_code(url.as_str(), resp.status()); + Ok(resp) + } + Err(e) => { + scans.increment_error(url.as_str()); + bail!(e) + } + } +} + /// Initiate request to the given `Url` using `Client` pub async fn make_request( client: &Client,