diff --git a/Cargo.lock b/Cargo.lock index 9f4cd61..2b39830 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1497,11 +1497,14 @@ checksum = "d039ad9159c98b70ecfd540b2573b97f7f52c3e8d9f8ad57a24b916a536975f9" dependencies = [ "futures-core", "futures-util", + "indexmap", "pin-project-lite", + "slab", "sync_wrapper", "tokio", "tower-layer", "tower-service", + "tracing", ] [[package]] @@ -1573,6 +1576,7 @@ dependencies = [ "tokio-util", "tower", "url", + "wyrand", ] [[package]] @@ -1864,6 +1868,15 @@ version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb" +[[package]] +name = "wyrand" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15e0359b0b8d9cdef235a1fd4a8c5d02e4c9204e9fac861c14c229a8e803d1a6" +dependencies = [ + "rand_core", +] + [[package]] name = "yoke" version = "0.8.0" diff --git a/Cargo.toml b/Cargo.toml index f4d99db..b3f5276 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,7 +4,7 @@ members = ["ubw-sward"] [workspace.dependencies] tokio = { version = "1", features = ["full"] } -tower = "0.5" +tower = { version = "0.5", features = ["balance", "load", "discover"] } reqwest = { version = "0.12", features = [ "json", "deflate", diff --git a/ubw-sward/Cargo.toml b/ubw-sward/Cargo.toml index b07eab4..e87c0b7 100644 --- a/ubw-sward/Cargo.toml +++ b/ubw-sward/Cargo.toml @@ -12,5 +12,6 @@ rand = {workspace = true} thiserror = {workspace = true} compact_str = {workspace = true} bytes = {workspace = true} +wyrand = "0.3" regex = "1.11" tokio-util = "0.7" \ No newline at end of file diff --git a/ubw-sward/src/http/header_config.rs b/ubw-sward/src/http/header_config.rs new file mode 100644 index 0000000..6e215ff --- /dev/null +++ b/ubw-sward/src/http/header_config.rs @@ -0,0 +1,47 @@ +use reqwest::cookie::Jar; +use reqwest::header::HeaderMap; +use url::Url; + +#[derive(Debug, Clone)] +pub struct HeadersConfig { + pub user_agent: Option, + pub gzip: bool, + pub deflate: bool, + pub cookie: Option, + pub other_headers: HeaderMap, +} + +impl Default for HeadersConfig { + fn default() -> Self { + Self { + user_agent: None, + gzip: true, + deflate: true, + cookie: None, + other_headers: HeaderMap::new(), + } + } +} + +impl HeadersConfig { + pub fn set_client_header( + &self, + client_builder: reqwest::ClientBuilder, + ) -> reqwest::ClientBuilder { + let mut client_builder = client_builder.gzip(self.gzip).deflate(self.deflate); + if let Some(user_agent) = &self.user_agent { + client_builder = client_builder.user_agent(user_agent) + } + client_builder + } + + pub fn get_cookie_jar(&self, url: &Url) -> Option { + if let Some(cookie) = &self.cookie { + let jar = Jar::default(); + jar.add_cookie_str(cookie, url); + Some(jar) + } else { + None + } + } +} diff --git a/ubw-sward/src/http/integrated.rs b/ubw-sward/src/http/integrated.rs new file mode 100644 index 0000000..642dda8 --- /dev/null +++ b/ubw-sward/src/http/integrated.rs @@ -0,0 +1,182 @@ +use crate::http::header_config::HeadersConfig; +use crate::http::simple::HttpSwardArray; +use crate::http::{RandomUrlGenerator, SimpleHttpSward}; +use crate::utils::multiplexed::MultiplexedSward; +use rand::Rng; +use reqwest::Method; +use std::net::{IpAddr, SocketAddr}; +use std::sync::Arc; +use thiserror::Error; +use url::Url; +use wyrand::WyRand; + +#[derive(Clone)] +enum AttackTarget { + Random(RandomUrlGenerator), + Fixed(Url), +} + +pub struct IntegratedHttpSward { + attack_target: AttackTarget, + request_sender: MultiplexedSward>>, +} + +impl IntegratedHttpSward { + pub fn builder() -> IntegratedHttpSwardBuilder { + IntegratedHttpSwardBuilder::new() + } +} + +pub struct IntegratedHttpSwardBuilder { + capacity: usize, + headers_config: HeadersConfig, + method: Method, + ip_list: Box<[IpAddr]>, +} + +impl Default for IntegratedHttpSwardBuilder { + fn default() -> Self { + Self { + capacity: 32, + method: Method::GET, + headers_config: HeadersConfig::default(), + ip_list: Box::new([]), + } + } +} + +impl IntegratedHttpSwardBuilder { + pub fn new() -> Self { + Self::default() + } + pub fn set_capacity(mut self, cap: usize) -> Self { + self.capacity = cap; + self + } + pub fn set_method(mut self, method: Method) -> Self { + self.method = method; + self + } + pub fn set_headers_config(mut self, headers: HeadersConfig) -> Self { + self.headers_config = headers; + self + } + pub fn set_ip_list(mut self, ip: &[IpAddr]) -> Self { + let cloned: Box<[_]> = ip.iter().copied().collect(); + self.ip_list = cloned; + self + } + + fn build_clients( + self, + url_example: &Url, + ) -> Result< + impl Iterator>, + HttpSwardBuildError, + > { + // check the ip list + if self.ip_list.is_empty() { + return Err(HttpSwardBuildError::EmptyIpResolve); + } + + // resolve the IP + let domain = url_example + .host_str() + .ok_or(HttpSwardBuildError::MissingHost)?; + let port = url_example.port_or_known_default().unwrap_or(443); + let socket_addrs = self + .ip_list + .into_iter() + .map(move |ip| SocketAddr::new(ip, port)); + + // get the cookie + let maybe_cookie = self + .headers_config + .get_cookie_jar(&url_example) + .map(Arc::new); + + // build the clients + let clients = socket_addrs + .map(|addr| { + reqwest::Client::builder() + .resolve(domain, addr) + .use_native_tls() + }) + .map(move |builder| self.headers_config.set_client_header(builder)) + .map(move |builder| match &maybe_cookie { + Some(jar) => builder.cookie_provider(jar.clone()).cookie_store(true), + None => builder, + }) + .map(|builder| { + builder + .build() + .map_err(HttpSwardBuildError::ReqwestBuildError) + }); + + Ok(clients) + } + + pub fn url(self, url: Url) -> Result { + let request_time_headers = self.headers_config.other_headers.clone(); + let method = self.method.clone(); + let capacity = self.capacity; + let clients = self.build_clients(&url)?; + let swards: Box<[_]> = clients + .map(|client_result| { + client_result.map(|client| { + SimpleHttpSward::new(client, method.clone(), request_time_headers.clone()) + }) + }) + .collect::>()?; + + Ok(IntegratedHttpSward { + attack_target: AttackTarget::Fixed(url), + request_sender: MultiplexedSward::new(SimpleHttpSward::array(swards), capacity), + }) + } + + pub fn template( + self, + template: impl AsRef, + ) -> Result { + let request_time_headers = self.headers_config.other_headers.clone(); + let method = self.method.clone(); + let capacity = self.capacity; + let random_core = WyRand::new(rand::rng().random()); + + let mut random_url_generator = RandomUrlGenerator::new(random_core, template.as_ref()) + .map_err(|_| HttpSwardBuildError::BadTemplate)?; + let example_url = random_url_generator + .generate_url() + .map_err(|_| HttpSwardBuildError::BadTemplate)?; + + let clients = self.build_clients(&example_url)?; + let swards: Box<[_]> = clients + .map(|client_result| { + client_result.map(|client| { + SimpleHttpSward::new(client, method.clone(), request_time_headers.clone()) + }) + }) + .collect::>()?; + + Ok(IntegratedHttpSward { + attack_target: AttackTarget::Random(random_url_generator), + request_sender: MultiplexedSward::new(SimpleHttpSward::array(swards), capacity), + }) + } +} + +#[derive(Debug, Error)] +pub enum HttpSwardBuildError { + #[error("Cannot parse domain or address in the url.")] + MissingHost, + + #[error("The IP list is empty")] + EmptyIpResolve, + + #[error("The template cannot generate a correct URL")] + BadTemplate, + + #[error("Failed to build reqwest client {0}")] + ReqwestBuildError(#[from] reqwest::Error), +} diff --git a/ubw-sward/src/http/mod.rs b/ubw-sward/src/http/mod.rs index b78dd22..404abec 100644 --- a/ubw-sward/src/http/mod.rs +++ b/ubw-sward/src/http/mod.rs @@ -1,5 +1,8 @@ pub mod random; pub mod simple; +pub mod integrated; +pub mod header_config; pub use random::RandomUrlGenerator; pub use simple::{SimpleHttpRequest, SimpleHttpSward}; +pub use integrated::IntegratedHttpSward; \ No newline at end of file diff --git a/ubw-sward/src/http/random.rs b/ubw-sward/src/http/random.rs index f1a87af..a9255bd 100644 --- a/ubw-sward/src/http/random.rs +++ b/ubw-sward/src/http/random.rs @@ -5,12 +5,12 @@ use url::Url; #[derive(Clone)] /// A Regex-based URL generator. -pub struct RandomUrlGenerator { +pub struct RandomUrlGenerator { random_engine: Rng, random_parts: Box<[RegexPart]>, } -impl RandomUrlGenerator { +impl RandomUrlGenerator { /// Create a new RandomUrlGenerator. /// /// # Arguments diff --git a/ubw-sward/src/http/simple.rs b/ubw-sward/src/http/simple.rs index 8da0372..f2c5ae3 100644 --- a/ubw-sward/src/http/simple.rs +++ b/ubw-sward/src/http/simple.rs @@ -4,6 +4,8 @@ use reqwest::{Client, Method}; use std::pin::Pin; use std::task::{Context, Poll}; use tower::Service; +use tower::balance::p2c::Balance; +use tower::discover::ServiceList; use url::Url; #[derive(Clone)] @@ -15,6 +17,8 @@ pub struct SimpleHttpSward { sent_count: usize, } +pub type HttpSwardArray = Balance, SimpleHttpRequest>; + impl SimpleHttpSward { /// Create a new simple http sward. pub fn new(client: Client, method: Method, headers: HeaderMap) -> Self { @@ -30,6 +34,13 @@ impl SimpleHttpSward { pub fn sent_count(&self) -> usize { self.sent_count } + + pub fn array(swards: T) -> HttpSwardArray + where + T: IntoIterator, + { + Balance::new(ServiceList::new(swards)) + } } #[derive(Clone)] @@ -55,3 +66,10 @@ impl Service for SimpleHttpSward { ) } } + +impl tower::load::Load for SimpleHttpSward { + type Metric = usize; + fn load(&self) -> Self::Metric { + self.sent_count + } +} diff --git a/ubw-sward/src/utils/mod.rs b/ubw-sward/src/utils/mod.rs index 0f75dc6..8ef693a 100644 --- a/ubw-sward/src/utils/mod.rs +++ b/ubw-sward/src/utils/mod.rs @@ -1 +1 @@ -pub mod multiplexed; \ No newline at end of file +pub mod multiplexed;