From 20eb7e472652658df1a61065560ea5a7a837b2b9 Mon Sep 17 00:00:00 2001
From: baerwang <baerwang@apache.org>
Date: Wed, 17 Jan 2024 19:29:13 +0800
Subject: [PATCH] style: add rx store field

---
 src/channel/mod.rs | 41 +++++++++++++++++++++++++++++++++++++++--
 src/cli/cmd.rs     |  3 +--
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/src/channel/mod.rs b/src/channel/mod.rs
index 4ec9c53..fc100e3 100644
--- a/src/channel/mod.rs
+++ b/src/channel/mod.rs
@@ -1,5 +1,6 @@
 use std::collections::HashSet;
 
+use headless_chrome::protocol::cdp::Network::{Request, Response};
 use headless_chrome::Browser;
 use tokio::sync::mpsc::Sender;
 
@@ -14,6 +15,7 @@ pub struct GlobalState {
     pub scan: Box<dyn ScanPolicy>,
     pub repeat: Box<dyn Duplicate>,
     pub store: HashSet<String>,
+    pub rx_store: HashSet<String>,
 
     pub sender: Option<Sender<String>>,
 }
@@ -31,10 +33,11 @@ impl GlobalState {
             domain,
             browser,
             config,
-            store: HashSet::new(),
-            sender: Some(tx),
             scan,
             repeat,
+            store: HashSet::new(),
+            rx_store: HashSet::new(),
+            sender: Some(tx),
         }
     }
 
@@ -45,4 +48,38 @@ impl GlobalState {
             }
         }
     }
+
+    pub async fn send_req(&mut self, req: Request) {
+        let handle = self.repeat.handle(req.clone());
+        if self.rx_store.insert(handle) {
+            self.send_message(req.url.as_str()).await
+        }
+    }
+
+    pub async fn send_rsp(&mut self, rsp: Response) {
+        if self.rx_store.insert(rsp.url.clone()) {
+            self.send_message(rsp.url.as_str()).await
+        }
+    }
+
+    pub fn send_block_message(&self, message: &str) {
+        if let Some(ref sender) = self.sender {
+            if sender.blocking_send(message.to_owned()).is_err() {
+                log::error!("Failed to send URL through channel");
+            }
+        }
+    }
+
+    pub fn send_block_req(&mut self, req: Request) {
+        let handle = self.repeat.handle(req.clone());
+        if self.rx_store.insert(handle) {
+            self.send_block_message(req.url.as_str())
+        }
+    }
+
+    pub fn send_block_rsp(&mut self, rsp: Response) {
+        if self.rx_store.insert(rsp.url.clone()) {
+            self.send_block_message(rsp.url.as_str())
+        }
+    }
 }
diff --git a/src/cli/cmd.rs b/src/cli/cmd.rs
index e0b181b..aae7aff 100644
--- a/src/cli/cmd.rs
+++ b/src/cli/cmd.rs
@@ -120,11 +120,10 @@ pub async fn cli() -> Result<(), Box<dyn std::error::Error>> {
                 config,
             );
             while let Some(url) = rx.recv().await {
-                if state.store.insert(url.clone()) {
+                if state.rx_store.insert(url.clone()) {
                     _ = crawler::tasks(url.clone().as_str(), tx.clone(), &mut state).await;
                 }
             }
-
             Ok(())
         }
     }