From 264a1552de1cecf9bc11952ee6af4dcf7697832b Mon Sep 17 00:00:00 2001 From: Nikaidou Haruki Date: Sat, 6 Sep 2025 19:03:40 +0900 Subject: [PATCH] implement half integrated udp sward --- ubw-sward/src/http/counter.rs | 19 ++++++++-- ubw-sward/src/http/integrated.rs | 36 +++++++++++++------ ubw-sward/src/udp/counter.rs | 22 ++++++++++++ ubw-sward/src/udp/generator.rs | 11 +++--- ubw-sward/src/udp/integrated.rs | 61 ++++++++++++++++++++++++++++++++ ubw-sward/src/udp/mod.rs | 8 ++++- 6 files changed, 139 insertions(+), 18 deletions(-) create mode 100644 ubw-sward/src/udp/counter.rs create mode 100644 ubw-sward/src/udp/integrated.rs diff --git a/ubw-sward/src/http/counter.rs b/ubw-sward/src/http/counter.rs index 901578e..ae9599b 100644 --- a/ubw-sward/src/http/counter.rs +++ b/ubw-sward/src/http/counter.rs @@ -2,13 +2,28 @@ use std::sync::atomic::AtomicUsize; #[derive(Default)] pub struct AtomicHttpCounter { + pub sent: AtomicUsize, pub http_2xx: AtomicUsize, pub http_3xx: AtomicUsize, pub http_4xx: AtomicUsize, pub http_5xx: AtomicUsize, + pub error: AtomicUsize, } -#[derive(Default)] +impl AtomicHttpCounter { + pub fn read(&self) -> HttpCounter { + HttpCounter { + sent: self.sent.load(std::sync::atomic::Ordering::Relaxed), + http_2xx: self.http_2xx.load(std::sync::atomic::Ordering::Relaxed), + http_3xx: self.http_3xx.load(std::sync::atomic::Ordering::Relaxed), + http_4xx: self.http_4xx.load(std::sync::atomic::Ordering::Relaxed), + http_5xx: self.http_5xx.load(std::sync::atomic::Ordering::Relaxed), + error: self.error.load(std::sync::atomic::Ordering::Relaxed), + } + } +} + +#[derive(Default, Clone, Copy)] pub struct HttpCounter { pub sent: usize, pub http_2xx: usize, @@ -16,4 +31,4 @@ pub struct HttpCounter { pub http_4xx: usize, pub http_5xx: usize, pub error: usize, -} \ No newline at end of file +} diff --git a/ubw-sward/src/http/integrated.rs b/ubw-sward/src/http/integrated.rs index c8ad622..22d6977 100644 --- a/ubw-sward/src/http/integrated.rs +++ b/ubw-sward/src/http/integrated.rs @@ -1,4 +1,4 @@ -use crate::http::counter::HttpCounter; +use crate::http::counter::{AtomicHttpCounter, HttpCounter}; use crate::http::header_config::HeadersConfig; use crate::http::simple::HttpSwardArray; use crate::http::{RandomUrlGenerator, SimpleHttpRequest, SimpleHttpSward}; @@ -6,7 +6,7 @@ use crate::utils::multiplexed::{MultiplexedSward, MultiplexedSwardError}; use rand::Rng; use reqwest::Method; use std::net::{IpAddr, SocketAddr}; -use std::sync::Arc; +use std::sync::{Arc, atomic}; use thiserror::Error; use tower::ServiceExt; use tower::{BoxError, Service}; @@ -22,12 +22,14 @@ enum AttackTarget { pub struct IntegratedHttpSward { attack_target: AttackTarget, request_sender: MultiplexedSward>>, - result_counter: HttpCounter, + result_counter: Arc, } impl IntegratedHttpSward { pub async fn oneshot(&mut self) -> Result { - self.result_counter.sent += 1; + self.result_counter + .sent + .fetch_add(1, atomic::Ordering::Relaxed); let url = match &mut self.attack_target { AttackTarget::Random(random) => random.generate_url()?, AttackTarget::Fixed(url) => url.clone(), @@ -44,16 +46,24 @@ impl IntegratedHttpSward { match res.status().as_u16() { 200..=299 => { - self.result_counter.http_2xx += 1; + self.result_counter + .http_2xx + .fetch_add(1, atomic::Ordering::Relaxed); } 300..=399 => { - self.result_counter.http_3xx += 1; + self.result_counter + .http_3xx + .fetch_add(1, atomic::Ordering::Relaxed); } 400..=499 => { - self.result_counter.http_4xx += 1; + self.result_counter + .http_4xx + .fetch_add(1, atomic::Ordering::Relaxed); } 500..=599 => { - self.result_counter.http_5xx += 1; + self.result_counter + .http_5xx + .fetch_add(1, atomic::Ordering::Relaxed); } _ => {} } @@ -73,12 +83,16 @@ impl IntegratedHttpSward { } oneshot_result = self.oneshot() => { if oneshot_result.is_err() { - self.result_counter.error += 1; + self.result_counter.error.fetch_and(1, atomic::Ordering::Relaxed); } } } } } + + pub fn get_counter(&self) -> Arc { + self.result_counter.clone() + } } impl IntegratedHttpSward { @@ -192,7 +206,7 @@ impl IntegratedHttpSwardBuilder { Ok(IntegratedHttpSward { attack_target: AttackTarget::Fixed(url), request_sender: MultiplexedSward::new(SimpleHttpSward::array(swards), capacity), - result_counter: HttpCounter::default(), + result_counter: Arc::new(AtomicHttpCounter::default()), }) } @@ -223,7 +237,7 @@ impl IntegratedHttpSwardBuilder { Ok(IntegratedHttpSward { attack_target: AttackTarget::Random(random_url_generator), request_sender: MultiplexedSward::new(SimpleHttpSward::array(swards), capacity), - result_counter: HttpCounter::default(), + result_counter: Arc::new(AtomicHttpCounter::default()), }) } } diff --git a/ubw-sward/src/udp/counter.rs b/ubw-sward/src/udp/counter.rs new file mode 100644 index 0000000..fd61ee7 --- /dev/null +++ b/ubw-sward/src/udp/counter.rs @@ -0,0 +1,22 @@ +use std::sync::atomic::AtomicU64; + +#[derive(Default, Clone, Copy)] +pub struct UdpFloodCounter { + pub sent_bytes: u64, + pub sent_packets: u64, +} + +#[derive(Default)] +pub struct AtomicUdpFloodCounter { + pub sent_bytes: AtomicU64, + pub sent_packets: AtomicU64, +} + +impl AtomicUdpFloodCounter { + pub fn read(&self) -> UdpFloodCounter { + UdpFloodCounter { + sent_bytes: self.sent_bytes.load(std::sync::atomic::Ordering::Relaxed), + sent_packets: self.sent_packets.load(std::sync::atomic::Ordering::Relaxed), + } + } +} diff --git a/ubw-sward/src/udp/generator.rs b/ubw-sward/src/udp/generator.rs index 7201574..2d1d101 100644 --- a/ubw-sward/src/udp/generator.rs +++ b/ubw-sward/src/udp/generator.rs @@ -1,4 +1,5 @@ use rand::RngCore; +use crate::udp::{BoxedUdpRequest, SizedUdpRequest}; pub struct RandomUdpWorkloadGenerator { rng: Rng, @@ -8,14 +9,16 @@ impl RandomUdpWorkloadGenerator { pub fn new(rng: Rng) -> Self { Self { rng } } - pub fn generate_sized(&mut self) -> [u8; SIZE] { + pub fn generate_sized(&mut self) -> SizedUdpRequest { let mut buf = [0u8; SIZE]; self.rng.fill_bytes(&mut buf); - buf + SizedUdpRequest { + bytes: buf + } } - pub fn generate_boxed(&mut self, size: usize) -> Box<[u8]> { + pub fn generate_boxed(&mut self, size: usize) -> BoxedUdpRequest { let mut buf = vec![0u8; size].into_boxed_slice(); self.rng.fill_bytes(&mut buf); - buf + BoxedUdpRequest(buf) } } \ No newline at end of file diff --git a/ubw-sward/src/udp/integrated.rs b/ubw-sward/src/udp/integrated.rs new file mode 100644 index 0000000..65ab1a7 --- /dev/null +++ b/ubw-sward/src/udp/integrated.rs @@ -0,0 +1,61 @@ +use crate::udp::counter::AtomicUdpFloodCounter; +use crate::udp::generator::RandomUdpWorkloadGenerator; +use crate::udp::random_flood::UdpSward; +use crate::udp::{BoxedUdpRequest, SizedUdpRequest}; +use rand::Rng; +use std::net::SocketAddr; +use std::sync::{atomic, Arc}; +use tokio::net::UdpSocket; +use tower::{Service, ServiceExt}; +use wyrand::WyRand; + +pub struct IntegratedUdpSward { + udp_sward: UdpSward, + random_udp_workload_generator: RandomUdpWorkloadGenerator, + counter: Arc, +} + +impl IntegratedUdpSward { + pub async fn new(target: SocketAddr) -> Result { + let random_core = WyRand::new(rand::rng().random()); + let generator = RandomUdpWorkloadGenerator::new(random_core); + let udp_socket = UdpSocket::bind("0.0.0.0:0").await?; + udp_socket.connect(target).await?; + let udp_socket = Arc::new(udp_socket); + let sward = UdpSward::new(udp_socket); + Ok(Self { + random_udp_workload_generator: generator, + udp_sward: sward, + counter: Arc::new(AtomicUdpFloodCounter::default()), + }) + } + pub async fn oneshot_array(&mut self) -> Result { + self.counter + .sent_bytes + .fetch_add(N as u64, atomic::Ordering::Relaxed); + self.counter + .sent_packets + .fetch_add(1, atomic::Ordering::Relaxed); + let content = self.random_udp_workload_generator.generate_sized::(); + >>::ready(&mut self.udp_sward) + .await? + .call(content) + .await + } + pub async fn oneshot_dynamic(&mut self, size: usize) -> Result { + self.counter + .sent_bytes + .fetch_add(size as u64, atomic::Ordering::Relaxed); + self.counter + .sent_packets + .fetch_add(1, atomic::Ordering::Relaxed); + let content = self.random_udp_workload_generator.generate_boxed(size); + >::ready(&mut self.udp_sward) + .await? + .call(content) + .await + } + pub fn get_counter(&self) -> Arc { + self.counter.clone() + } +} diff --git a/ubw-sward/src/udp/mod.rs b/ubw-sward/src/udp/mod.rs index 5fc6ce0..5fccf93 100644 --- a/ubw-sward/src/udp/mod.rs +++ b/ubw-sward/src/udp/mod.rs @@ -1,2 +1,8 @@ +pub mod generator; +pub mod integrated; pub mod random_flood; -pub mod generator; \ No newline at end of file +mod counter; + +pub use generator::RandomUdpWorkloadGenerator; +pub use integrated::IntegratedUdpSward; +pub use random_flood::{BoxedUdpRequest, SizedUdpRequest, UdpSward};