implement integrated http sward

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"

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<String>,
+    pub gzip: bool,
+    pub deflate: bool,
+    pub cookie: Option<String>,
+    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<Jar> {
+        if let Some(cookie) = &self.cookie {
+            let jar = Jar::default();
+            jar.add_cookie_str(cookie, url);
+            Some(jar)
+        } else {
+            None
+        }
+    }
+}

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<WyRand>),
+    Fixed(Url),
+}
+
+pub struct IntegratedHttpSward {
+    attack_target: AttackTarget,
+    request_sender: MultiplexedSward<HttpSwardArray<Box<[SimpleHttpSward]>>>,
+}
+
+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<Item = Result<reqwest::Client, HttpSwardBuildError>>,
+        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<IntegratedHttpSward, HttpSwardBuildError> {
+        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::<Result<_, HttpSwardBuildError>>()?;
+
+        Ok(IntegratedHttpSward {
+            attack_target: AttackTarget::Fixed(url),
+            request_sender: MultiplexedSward::new(SimpleHttpSward::array(swards), capacity),
+        })
+    }
+
+    pub fn template(
+        self,
+        template: impl AsRef<str>,
+    ) -> Result<IntegratedHttpSward, HttpSwardBuildError> {
+        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::<Result<_, HttpSwardBuildError>>()?;
+
+        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),
+}

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;

ubw-sward/src/http/random.rs

@@ -5,12 +5,12 @@ use url::Url;
 
 #[derive(Clone)]
 /// A Regex-based URL generator.
-pub struct RandomUrlGenerator<Rng: rand::Rng> {
+pub struct RandomUrlGenerator<Rng: rand::RngCore> {
     random_engine: Rng,
     random_parts: Box<[RegexPart]>,
 }
 
-impl<Rng: rand::Rng> RandomUrlGenerator<Rng> {
+impl<Rng: rand::RngCore> RandomUrlGenerator<Rng> {
     /// Create a new RandomUrlGenerator.
     ///
     /// # Arguments

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<T> = Balance<ServiceList<T>, 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<T>(swards: T) -> HttpSwardArray<T>
+    where
+        T: IntoIterator<Item = Self>,
+    {
+        Balance::new(ServiceList::new(swards))
+    }
 }
 
 #[derive(Clone)]
@@ -55,3 +66,10 @@ impl Service<SimpleHttpRequest> for SimpleHttpSward {
         )
     }
 }
+
+impl tower::load::Load for SimpleHttpSward {
+    type Metric = usize;
+    fn load(&self) -> Self::Metric {
+        self.sent_count
+    }
+}

ubw-sward/src/utils/mod.rs

@@ -1 +1 @@
-pub mod multiplexed;
+pub mod multiplexed;