implement half integrated udp sward

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,
-}
+}

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<HttpSwardArray<Box<[SimpleHttpSward]>>>,
-    result_counter: HttpCounter,
+    result_counter: Arc<AtomicHttpCounter>,
 }
 
 impl IntegratedHttpSward {
     pub async fn oneshot(&mut self) -> Result<reqwest::Response, HttpSwardError> {
-        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<AtomicHttpCounter> {
+        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()),
         })
     }
 }

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),
+        }
+    }
+}

ubw-sward/src/udp/generator.rs

@@ -1,4 +1,5 @@
 use rand::RngCore;
+use crate::udp::{BoxedUdpRequest, SizedUdpRequest};
 
 pub struct RandomUdpWorkloadGenerator<Rng: RngCore> {
     rng: Rng,
@@ -8,14 +9,16 @@ impl<Rng: RngCore> RandomUdpWorkloadGenerator<Rng> {
     pub fn new(rng: Rng) -> Self {
         Self { rng }
     }
-    pub fn generate_sized<const SIZE: usize>(&mut self) -> [u8; SIZE] {
+    pub fn generate_sized<const SIZE: usize>(&mut self) -> SizedUdpRequest<SIZE> {
         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)
     }
 }

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<WyRand>,
+    counter: Arc<AtomicUdpFloodCounter>,
+}
+
+impl IntegratedUdpSward {
+    pub async fn new(target: SocketAddr) -> Result<Self, std::io::Error> {
+        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<const N: usize>(&mut self) -> Result<usize, std::io::Error> {
+        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::<N>();
+        <UdpSward as ServiceExt<SizedUdpRequest<N>>>::ready(&mut self.udp_sward)
+            .await?
+            .call(content)
+            .await
+    }
+    pub async fn oneshot_dynamic(&mut self, size: usize) -> Result<usize, std::io::Error> {
+        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);
+        <UdpSward as ServiceExt<BoxedUdpRequest>>::ready(&mut self.udp_sward)
+            .await?
+            .call(content)
+            .await
+    }
+    pub fn get_counter(&self) -> Arc<AtomicUdpFloodCounter> {
+        self.counter.clone()
+    }
+}

ubw-sward/src/udp/mod.rs

@@ -1,2 +1,8 @@
+pub mod generator;
+pub mod integrated;
 pub mod random_flood;
-pub mod generator;
+mod counter;
+
+pub use generator::RandomUdpWorkloadGenerator;
+pub use integrated::IntegratedUdpSward;
+pub use random_flood::{BoxedUdpRequest, SizedUdpRequest, UdpSward};