bless

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;

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;
 

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),

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}"));

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;