fixes buffer overflow problem

Author: Shourya742

roles/pool/src/lib/mining_pool/mod.rs

@@ -248,8 +248,9 @@ impl Downstream {
                     return;
                 }
             };
+            let mut receiver = receiver.subscribe();
             loop {
-                match receiver.subscribe().recv().await {
+                match receiver.recv().await {
                     Ok(received) => {
                         let received: Result<StdFrame, _> = received
                             .try_into()

roles/pool/src/lib/template_receiver/mod.rs

@@ -1,3 +1,5 @@
+#![allow(dead_code)]
+#![allow(clippy::result_large_err)]
 use super::{
     error::{PoolError, PoolResult},
     mining_pool::{EitherFrame, StdFrame},
@@ -23,6 +25,7 @@ mod message_handler;
 mod setup_connection;
 use setup_connection::SetupConnectionHandler;
 
+#[derive(Clone)]
 pub struct TemplateRx {
     receiver: tokio::sync::broadcast::Sender<EitherFrame>,
     sender: tokio::sync::broadcast::Sender<EitherFrame>,
@@ -68,15 +71,14 @@ impl TemplateRx {
 
         SetupConnectionHandler::setup(&mut receiver, &mut sender, address).await?;
 
-        let self_ = Arc::new(Mutex::new(Self {
+        let self_ = Self {
             receiver,
-            sender,
+            sender: sender.clone(),
             new_template_sender: templ_sender,
             new_prev_hash_sender: prev_h_sender,
             message_received_signal,
-            status_tx,
-        }));
-        let cloned = self_.clone();
+            status_tx: status_tx.clone(),
+        };
 
         let c_additional_size = CoinbaseOutputDataSize {
             coinbase_output_max_additional_size: coinbase_out_len,
@@ -86,29 +88,26 @@ impl TemplateRx {
         )
         .try_into()?;
 
-        Self::send(self_.clone(), frame).await?;
+        Self::send(sender.clone(), frame)?;
 
-        task::spawn(async { Self::start(cloned).await });
-        task::spawn(async { Self::on_new_solution(self_, solution_receiver).await });
+        task::spawn(async { Self::start(self_).await });
+        task::spawn(async { Self::on_new_solution(sender, status_tx, solution_receiver).await });
 
         Ok(())
     }
 
-    pub async fn start(self_: Arc<Mutex<Self>>) {
-        let (mut recv_msg_signal, receiver, new_template_sender, new_prev_hash_sender, status_tx) =
-            self_
-                .safe_lock(|s| {
-                    (
-                        s.message_received_signal.subscribe(),
-                        s.receiver.clone(),
-                        s.new_template_sender.clone(),
-                        s.new_prev_hash_sender.clone(),
-                        s.status_tx.clone(),
-                    )
-                })
-                .unwrap();
+    pub async fn start(self_: Self) {
+        let (mut recv_msg_signal, receiver, new_template_sender, new_prev_hash_sender, status_tx) = (
+            self_.message_received_signal.subscribe(),
+            self_.receiver.clone(),
+            self_.new_template_sender.clone(),
+            self_.new_prev_hash_sender.clone(),
+            self_.status_tx.clone(),
+        );
+        let mut receiver = receiver.subscribe();
         loop {
-            let message_from_tp = handle_result!(status_tx, receiver.subscribe().recv().await);
+            let message_from_tp = handle_result!(status_tx, receiver.recv().await);
+            info!("Message: {:?}", message_from_tp);
             let mut message_from_tp: StdFrame = handle_result!(
                 status_tx,
                 message_from_tp
@@ -123,7 +122,7 @@ impl TemplateRx {
             let msg = handle_result!(
                 status_tx,
                 ParseServerTemplateDistributionMessages::handle_message_template_distribution(
-                    self_.clone(),
+                    Arc::new(Mutex::new(self_.clone())),
                     message_type,
                     payload,
                 )
@@ -155,34 +154,35 @@ impl TemplateRx {
         }
     }
 
-    pub async fn send(self_: Arc<Mutex<Self>>, sv2_frame: StdFrame) -> PoolResult<()> {
+    pub fn send(
+        sender: tokio::sync::broadcast::Sender<EitherFrame>,
+        sv2_frame: StdFrame,
+    ) -> PoolResult<()> {
         let either_frame = sv2_frame.into();
-        let sender = self_
-            .safe_lock(|self_| self_.sender.clone())
-            .map_err(|e| PoolError::PoisonLock(e.to_string()))?;
         sender.send(either_frame)?;
         Ok(())
     }
 
     async fn on_new_solution(
-        self_: Arc<Mutex<Self>>,
+        sender: tokio::sync::broadcast::Sender<EitherFrame>,
+        status_tx: status::Sender,
         mut rx: tokio::sync::mpsc::Receiver<SubmitSolution<'static>>,
-    ) {
-        let status_tx = self_.safe_lock(|s| s.status_tx.clone()).unwrap();
+    ) -> PoolResult<()> {
         while let Some(solution) = rx.recv().await {
             info!("Sending Solution to TP: {:?}", &solution);
             let sv2_frame_res: Result<StdFrame, _> =
                 PoolMessages::TemplateDistribution(TemplateDistribution::SubmitSolution(solution))
                     .try_into();
             match sv2_frame_res {
                 Ok(frame) => {
-                    handle_result!(status_tx, Self::send(self_.clone(), frame).await);
+                    handle_result!(status_tx, Self::send(sender.clone(), frame));
                 }
                 Err(_e) => {
                     // return submit error
                     todo!()
                 }
             };
         }
+        Ok(())
     }
 }

roles/pool/src/main.rs

@@ -68,19 +68,20 @@ mod args {
     }
 }
 
-use tracing_subscriber::prelude::*;
-
-fn init_tracing() {
-    let console_layer = console_subscriber::spawn();
-    tracing_subscriber::registry()
-        .with(console_layer)
-        .with(tracing_subscriber::fmt::layer())
-        .init();
-}
+// use tracing_subscriber::prelude::*;
+
+// fn init_tracing() {
+//     let console_layer = console_subscriber::spawn();
+//     tracing_subscriber::registry()
+//         .with(console_layer)
+//         .with(tracing_subscriber::fmt::layer())
+//         .init();
+// }
 
 #[tokio::main]
 async fn main() {
-    init_tracing();
+    // init_tracing();
+    tracing_subscriber::fmt::init();
 
     let args = match args::Args::from_args() {
         Ok(cfg) => cfg,