From 3d54a2b57d4d984eac27fdc840c88c8b74d1a939 Mon Sep 17 00:00:00 2001
From: syntaxjak <syntaxjak@users.noreply.github.com>
Date: Fri, 19 Jun 2026 11:15:48 -0500
Subject: [PATCH] Refactor stratum server leak handling

---
 .gitignore                          |   1 +
 api/src/types.rs                    |   3 +-
 servers/src/mining/stratumserver.rs | 139 ++++++++++++++++------------
 3 files changed, 82 insertions(+), 61 deletions(-)

diff --git a/.gitignore b/.gitignore
index 4e438431b7..2bcaf370c4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,5 +13,6 @@ wallet_data
 wallet/db
 .idea/
 .vscode/
+.venv/
 /node*
 result
diff --git a/api/src/types.rs b/api/src/types.rs
index 040e1eec5d..06918262eb 100644
--- a/api/src/types.rs
+++ b/api/src/types.rs
@@ -465,7 +465,8 @@ impl<'de> serde::de::Deserialize<'de> for OutputPrintable {
 				}
 
 				if output_type.is_none()
-					|| commit.is_none() || spent.is_none()
+					|| commit.is_none()
+					|| spent.is_none()
 					|| proof_hash.is_none()
 					|| mmr_index.is_none()
 				{
diff --git a/servers/src/mining/stratumserver.rs b/servers/src/mining/stratumserver.rs
index d39312a560..0baddf40fd 100644
--- a/servers/src/mining/stratumserver.rs
+++ b/servers/src/mining/stratumserver.rs
@@ -15,9 +15,8 @@
 //! Mining Stratum Server
 
 use futures::channel::mpsc;
-use futures::pin_mut;
-use futures::{SinkExt, StreamExt, TryStreamExt};
-use tokio::net::TcpListener;
+use futures::{SinkExt, StreamExt};
+use tokio::net::{TcpListener, TcpStream};
 use tokio::runtime::Runtime;
 use tokio_util::codec::{Framed, LinesCodec};
 
@@ -599,72 +598,92 @@ impl Handler {
 
 // ----------------------------------------
 // Worker Factory Thread Function
+
+struct WorkerCleanup {
+	worker_id: usize,
+	workers: Arc<WorkersList>,
+}
+
+impl Drop for WorkerCleanup {
+	fn drop(&mut self) {
+		self.workers.remove_worker(self.worker_id);
+		info!("Worker {} disconnected", self.worker_id);
+	}
+}
+
+async fn handle_connection(socket: TcpStream, handler: Arc<Handler>) {
+	let (tx, mut rx) = mpsc::unbounded();
+	let worker_id = handler.workers.add_worker(tx);
+	let _cleanup = WorkerCleanup {
+		worker_id,
+		workers: handler.workers.clone(),
+	};
+
+	if let Ok(addr) = socket.peer_addr() {
+		info!("Worker {} connected from {}", worker_id, addr);
+	} else {
+		info!("Worker {} connected", worker_id);
+	}
+
+	let mut framed = Framed::new(socket, LinesCodec::new());
+
+	loop {
+		tokio::select! {
+			incoming = framed.next() => {
+				match incoming {
+					Some(Ok(line)) => {
+						let request: RpcRequest = match serde_json::from_str(&line) {
+							Ok(req) => req,
+							Err(e) => {
+								error!("error serializing line: {}", e);
+								break;
+							}
+						};
+						let resp = handler.handle_rpc_requests(request, worker_id);
+						handler.workers.send_to(worker_id, resp);
+					}
+					Some(Err(e)) => {
+						error!("error reading line: {}", e);
+						break;
+					}
+					None => break,
+				}
+			}
+			outgoing = rx.next() => {
+				match outgoing {
+					Some(line) => {
+						if let Err(e) = framed.send(line).await {
+							error!("error writing line: {}", e);
+							break;
+						}
+					}
+					None => break,
+				}
+			}
+		}
+	}
+}
+
 fn accept_connections(listen_addr: SocketAddr, handler: Arc<Handler>) {
 	info!("Start tokio stratum server");
 	let task = async move {
 		let listener = TcpListener::bind(&listen_addr).await.unwrap_or_else(|_| {
 			panic!("Stratum: Failed to bind to listen address {}", listen_addr)
 		});
-		let server = async_stream::stream! {
-			loop {
-				match listener.accept().await {
-					Ok((socket, _)) => yield socket,
-					Err(e) => {
-						error!("accept error = {:?}", e);
-						continue;
-					}
+		loop {
+			match listener.accept().await {
+				Ok((socket, _)) => {
+					let handler = handler.clone();
+					tokio::spawn(async move {
+						handle_connection(socket, handler).await;
+					});
+				}
+				Err(e) => {
+					error!("accept error = {:?}", e);
+					continue;
 				}
 			}
 		}
-		.for_each(move |socket| {
-			let handler = handler.clone();
-			async move {
-				// Spawn a task to process the connection
-				let (tx, mut rx) = mpsc::unbounded();
-
-				let worker_id = handler.workers.add_worker(tx);
-				info!("Worker {} connected", worker_id);
-
-				let framed = Framed::new(socket, LinesCodec::new());
-				let (mut writer, mut reader) = framed.split();
-
-				let h = handler.clone();
-				let read = async move {
-					while let Some(line) = reader
-						.try_next()
-						.await
-						.map_err(|e| error!("error reading line: {}", e))?
-					{
-						let request = serde_json::from_str(&line)
-							.map_err(|e| error!("error serializing line: {}", e))?;
-						let resp = h.handle_rpc_requests(request, worker_id);
-						h.workers.send_to(worker_id, resp);
-					}
-
-					Result::<_, ()>::Ok(())
-				};
-
-				let write = async move {
-					while let Some(line) = rx.next().await {
-						writer
-							.send(line)
-							.await
-							.map_err(|e| error!("error writing line: {}", e))?;
-					}
-
-					Result::<_, ()>::Ok(())
-				};
-
-				let task = async move {
-					pin_mut!(read, write);
-					futures::future::select(read, write).await;
-					handler.workers.remove_worker(worker_id);
-					info!("Worker {} disconnected", worker_id);
-				};
-				tokio::spawn(task);
-			}
-		});
-		server.await
 	};
 
 	let rt = Runtime::new().unwrap();