From 28f9b75e8fafa2f96ee3c360d7f767dc854f9263 Mon Sep 17 00:00:00 2001
From: Nikaidou Haruki <nharuki@pm.me>
Date: Fri, 5 Sep 2025 01:24:11 +0900
Subject: [PATCH] implement full layer7 functions

---
 ubw-sward/src/http/integrated.rs | 28 ++++++++++++++++++++++++----
 ubw-sward/src/http/mod.rs        |  2 +-
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/ubw-sward/src/http/integrated.rs b/ubw-sward/src/http/integrated.rs
index 5b3932b..1b28bde 100644
--- a/ubw-sward/src/http/integrated.rs
+++ b/ubw-sward/src/http/integrated.rs
@@ -1,3 +1,4 @@
+use crate::http::counter::HttpCounter;
 use crate::http::header_config::HeadersConfig;
 use crate::http::simple::HttpSwardArray;
 use crate::http::{RandomUrlGenerator, SimpleHttpRequest, SimpleHttpSward};
@@ -11,7 +12,6 @@ use tower::ServiceExt;
 use tower::{BoxError, Service};
 use url::Url;
 use wyrand::WyRand;
-use crate::http::counter::{HttpCounter};
 
 #[derive(Clone)]
 enum AttackTarget {
@@ -22,7 +22,7 @@ enum AttackTarget {
 pub struct IntegratedHttpSward {
     attack_target: AttackTarget,
     request_sender: MultiplexedSward<HttpSwardArray<Box<[SimpleHttpSward]>>>,
-    result_counter: HttpCounter
+    result_counter: HttpCounter,
 }
 
 impl IntegratedHttpSward {
@@ -59,6 +59,23 @@ impl IntegratedHttpSward {
 
         Ok(res)
     }
+
+    pub async fn run_with_signal(
+        mut self,
+        mut signal: tokio::sync::oneshot::Receiver<()>,
+    ) -> Result<(), HttpSwardError> {
+        loop {
+            tokio::select! {
+                receive_result = &mut signal => {
+                    receive_result?;
+                    return Ok(());
+                }
+                oneshop_result = self.oneshot() => {
+                    oneshop_result?;
+                }
+            }
+        }
+    }
 }
 
 impl IntegratedHttpSward {
@@ -172,7 +189,7 @@ impl IntegratedHttpSwardBuilder {
         Ok(IntegratedHttpSward {
             attack_target: AttackTarget::Fixed(url),
             request_sender: MultiplexedSward::new(SimpleHttpSward::array(swards), capacity),
-            result_counter: HttpCounter::default()
+            result_counter: HttpCounter::default(),
         })
     }
 
@@ -203,7 +220,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: HttpCounter::default(),
         })
     }
 }
@@ -230,4 +247,7 @@ pub enum HttpSwardError {
 
     #[error(transparent)]
     BadUrl(#[from] url::ParseError),
+
+    #[error(transparent)]
+    ReceiverError(#[from] tokio::sync::oneshot::error::RecvError),
 }
diff --git a/ubw-sward/src/http/mod.rs b/ubw-sward/src/http/mod.rs
index c107376..b2d8746 100644
--- a/ubw-sward/src/http/mod.rs
+++ b/ubw-sward/src/http/mod.rs
@@ -2,7 +2,7 @@ pub mod random;
 pub mod simple;
 pub mod integrated;
 pub mod header_config;
-mod counter;
+pub mod counter;
 
 pub use random::RandomUrlGenerator;
 pub use simple::{SimpleHttpRequest, SimpleHttpSward};