From 79bc2958efddcbe680f7fcdf953deea2b413e5f4 Mon Sep 17 00:00:00 2001
From: victor <victorvobis@web.de>
Date: Sun, 14 Sep 2025 19:42:21 +0200
Subject: [PATCH] bless

---
 node/src/lib.rs                              |  5 ++--
 node/src/{protocol => network}/connection.rs |  2 +-
 node/src/{protocol => network}/connector.rs  |  2 +-
 node/src/{protocol => network}/message.rs    |  0
 node/src/{node => network}/ws_server.rs      | 26 ++++++++++++--------
 node/src/node/node.rs                        |  6 ++---
 6 files changed, 24 insertions(+), 17 deletions(-)
 rename node/src/{protocol => network}/connection.rs (98%)
 rename node/src/{protocol => network}/connector.rs (99%)
 rename node/src/{protocol => network}/message.rs (100%)
 rename node/src/{node => network}/ws_server.rs (85%)

diff --git a/node/src/lib.rs b/node/src/lib.rs
index e34c788..e4ac026 100644
--- a/node/src/lib.rs
+++ b/node/src/lib.rs
@@ -8,7 +8,6 @@ pub mod node {
   pub mod blockchain;
   pub use blockchain::*;
 
-  pub mod ws_server;
 }
 
 pub mod constants;
@@ -65,7 +64,7 @@ pub mod watcher {
   pub use command::*;
 }
 
-pub mod protocol {
+pub mod network {
   pub mod message;
   pub use message::*;
 
@@ -74,6 +73,8 @@ pub mod protocol {
 
   pub mod connector;
   pub use connector::*;
+
+  pub mod ws_server;
 }
 
 pub mod seeds_constants;
diff --git a/node/src/protocol/connection.rs b/node/src/network/connection.rs
similarity index 98%
rename from node/src/protocol/connection.rs
rename to node/src/network/connection.rs
index 3fe02a7..1f8ad23 100644
--- a/node/src/protocol/connection.rs
+++ b/node/src/network/connection.rs
@@ -1,7 +1,7 @@
 use crate::executor::ExecutorCommand;
 use crate::log;
 use crate::node::node;
-use crate::protocol::ProtocolMessage;
+use super::ProtocolMessage;
 use tokio::net;
 use tokio::sync::mpsc;
 
diff --git a/node/src/protocol/connector.rs b/node/src/network/connector.rs
similarity index 99%
rename from node/src/protocol/connector.rs
rename to node/src/network/connector.rs
index 486dc6c..e9a22d1 100644
--- a/node/src/protocol/connector.rs
+++ b/node/src/network/connector.rs
@@ -13,7 +13,7 @@ use crate::bus::*;
 use crate::executor::ExecutorCommand;
 use crate::node::node;
 use crate::node::{NetworkError, error};
-use crate::protocol::ProtocolMessage;
+use super::ProtocolMessage;
 
 pub enum ConnectorCommand {
   ConnectToTcpPeer(SocketAddr),
diff --git a/node/src/protocol/message.rs b/node/src/network/message.rs
similarity index 100%
rename from node/src/protocol/message.rs
rename to node/src/network/message.rs
diff --git a/node/src/node/ws_server.rs b/node/src/network/ws_server.rs
similarity index 85%
rename from node/src/node/ws_server.rs
rename to node/src/network/ws_server.rs
index d789e57..6892dac 100644
--- a/node/src/node/ws_server.rs
+++ b/node/src/network/ws_server.rs
@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 use std::net::SocketAddr;
 
+use tokio::io::AsyncWriteExt;
 use tokio::net::TcpStream;
 use futures::{SinkExt, StreamExt};
 use tokio::sync::mpsc::{self, Receiver, Sender};
@@ -116,19 +117,24 @@ impl WsServer {
             None => {}
           }
         }
+
         con_res =  listener.accept() => {
           match con_res {
-            Ok((stream, addr)) => {
+            Ok((mut stream, addr)) => {
               log(msg!(DEBUG, "Received Connection Attempt from {}", addr));
-              let (tx, rx) = mpsc::channel::<WsClientResponse>(100);
-              self.clients.insert(addr, tx);
-              let tx = self.tx.clone();
-              let task_handle = tokio::spawn(async move {
-                if let Err(e) = ws_connection(stream, rx, tx).await {
-                  log(msg!(ERROR, "{e}"));
-                }
-              });
-              tasks.push(task_handle);
+              if self.clients.len() < 5 {
+                let (tx, rx) = mpsc::channel::<WsClientResponse>(100);
+                self.clients.insert(addr, tx);
+                let tx = self.tx.clone();
+                let task_handle = tokio::spawn(async move {
+                  if let Err(e) = ws_connection(stream, rx, tx).await {
+                    log(msg!(ERROR, "{e}"));
+                  }
+                });
+                tasks.push(task_handle);
+              } else {
+                stream.shutdown().await?;
+              }
             }
             Err(e) => {
               log(msg!(ERROR, "{e}"));
diff --git a/node/src/node/node.rs b/node/src/node/node.rs
index 2caff51..5d55dc6 100644
--- a/node/src/node/node.rs
+++ b/node/src/node/node.rs
@@ -3,11 +3,11 @@ use shared::core::{self, ChainData};
 use shared::print_error_chain;
 use crate::executor::ExecutorCommand;
 use crate::log;
-use crate::protocol::ProtocolMessage;
-use crate::protocol::{Connector, ConnectorCommand};
+use crate::network::ProtocolMessage;
+use crate::network::{Connector, ConnectorCommand};
 use crate::seeds_constants::SEED_NODES;
 use crate::watcher::{WatcherCommand, WatcherMode};
-use super::ws_server::{WsCommand, WsServer};
+use crate::network::ws_server::{WsCommand, WsServer};
 use super::{ Blockchain, BlockchainError, ValidationError };
 
 use std::collections::HashMap;