haruki/ubw-network
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

implement half integrated udp sward

Browse Source
This commit is contained in:
Nikaidou Haruki
2025-09-06 19:03:40 +09:00
parent ff4feb5faa
commit 264a1552de
6 changed files with 139 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
ubw-sward/src/http/counter.rs
View File

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

36
ubw-sward/src/http/integrated.rs
View File

@@ -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()),
})
}
}

22
ubw-sward/src/udp/counter.rs Normal file
View File

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

11
ubw-sward/src/udp/generator.rs
View File

@@ -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)
}
}

61
ubw-sward/src/udp/integrated.rs Normal file
View File

@@ -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()
}
}

8
ubw-sward/src/udp/mod.rs
View File

@@ -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};