From 699b0308b1da5dc664d2689593f0bbcd08659533 Mon Sep 17 00:00:00 2001 From: Reboot-Codes Date: Fri, 25 Apr 2025 10:34:05 -0700 Subject: [PATCH] Fix max client overflow --- src/client/mod.rs | 94 +++++++++++++++++++++------------------- src/server/mod.rs | 1 + src/server/models.rs | 82 +++++++++++++++++++++++------------ src/server/websockets.rs | 7 ++- 4 files changed, 110 insertions(+), 74 deletions(-) diff --git a/src/client/mod.rs b/src/client/mod.rs index 75a9714..da94a99 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -1,4 +1,13 @@ -use crate::{arbiter::models::ApiKey, server::{models::IPCMessageWithId, websockets::WsIn, AUTH_HEADER}, user::NexusUser}; +use crate::{ + arbiter::models::ApiKey, + server::{ + AUTH_HEADER, + MAX_SIZE, + models::IPCMessageWithId, + websockets::WsIn, + }, + user::NexusUser, +}; use fastwebsockets::{ FragmentCollector, Frame, @@ -14,12 +23,16 @@ use hyper::{ UPGRADE, }, }; +use log::error; use log::info; +use std::collections::HashMap; use std::future::Future; use std::sync::Arc; use std::time::Duration; -use std::collections::HashMap; -use tokio::{net::TcpStream, sync::broadcast}; +use tokio::{ + net::TcpStream, + sync::broadcast, +}; use tokio::{ sync::Mutex, task::{ @@ -28,18 +41,15 @@ use tokio::{ }, }; use tokio_util::sync::CancellationToken; -use log::error; #[derive(Debug, Clone)] pub struct ClientStatus { - pub connected: bool + pub connected: bool, } impl ClientStatus { pub fn new(connected: bool) -> Self { - ClientStatus { - connected - } + ClientStatus { connected } } pub fn set(&mut self, connected: bool) { @@ -65,8 +75,7 @@ pub struct NexusClient { from_server: broadcast::Sender, handles: Vec>, api_keys: Arc>>, - cancellation_token: CancellationToken - // TODO: Add user registry to see if a user is connected via this client to route back instead of sending to server. + cancellation_token: CancellationToken, // TODO: Add user registry to see if a user is connected via this client to route back instead of sending to server. } struct AddSend(T); @@ -87,15 +96,9 @@ where } impl NexusClient { - pub fn new( - secure: bool, - url: &String, - port: &u16, - api_key: &String, - keep_trying: bool, - ) -> Self { - let (from_server, _) = broadcast::channel::(usize::MAX / 2); - let (to_server_tx, _) = broadcast::channel::(usize::MAX / 2); + pub fn new(secure: bool, url: &String, port: &u16, api_key: &String, keep_trying: bool) -> Self { + let (from_server, _) = broadcast::channel::(MAX_SIZE); + let (to_server_tx, _) = broadcast::channel::(MAX_SIZE); NexusClient { secure, @@ -108,27 +111,18 @@ impl NexusClient { from_server, handles: Vec::new(), api_keys: Arc::new(Mutex::new(HashMap::new())), - cancellation_token: CancellationToken::new() + cancellation_token: CancellationToken::new(), } } // TODO: Send as connected API key. (ensure routing is supported) // TODO: Send as proxied API key. (ensure routing is supported) - pub async fn connect_to_ws( - &mut self, - ) -> Result< - JoinHandle<()>, - anyhow::Error, - > { + pub async fn connect_to_ws(&mut self) -> Result, anyhow::Error> { let mut keep_trying = true; let mut ws_opt = None; let mut error: Option = None; - let host = format!( - "{}:{}", - self.url.clone(), - self.port.clone().to_string() - ); + let host = format!("{}:{}", self.url.clone(), self.port.clone().to_string()); let uri = format!( "{}://{}:{}/ws", (if self.secure { "https" } else { "http" }), @@ -142,10 +136,7 @@ impl NexusClient { match Request::builder() .method("GET") .uri(uri.clone()) - .header( - "Host", - host.clone(), - ) + .header("Host", host.clone()) .header(UPGRADE, "websocket") .header(CONNECTION, "upgrade") .header( @@ -161,7 +152,11 @@ impl NexusClient { ws_opt = Some(the_socket); } Err(e) => { - error!("Failed to perform WebSocket Handshake with \"{}\":\n{}", uri.clone(), e); + error!( + "Failed to perform WebSocket Handshake with \"{}\":\n{}", + uri.clone(), + e + ); if self.keep_trying { error!("Retrying WebSocket Handshake with \"{}\"...", uri.clone()); tokio::time::sleep(Duration::from_secs(1)).await; @@ -172,9 +167,16 @@ impl NexusClient { } }, Err(e) => { - error!("Failed to send WebSocket Upgrade request to \"{}\":\n{}", uri.clone(), e); + error!( + "Failed to send WebSocket Upgrade request to \"{}\":\n{}", + uri.clone(), + e + ); if self.keep_trying { - info!("Retrying WebSocket Upgrade request for \"{}\"...", uri.clone()); + info!( + "Retrying WebSocket Upgrade request for \"{}\"...", + uri.clone() + ); tokio::time::sleep(Duration::from_secs(1)).await; } else { error = Some(e.into()); @@ -184,7 +186,11 @@ impl NexusClient { } } Err(e) => { - error!("Failed to open TCP connection to \"{}\":\n{}", host.clone(), e); + error!( + "Failed to open TCP connection to \"{}\":\n{}", + host.clone(), + e + ); if self.keep_trying { info!("Retrying TCP connection to \"{}\"...", host.clone()); tokio::time::sleep(Duration::from_secs(1)).await; @@ -202,7 +208,7 @@ impl NexusClient { let (mut reader, og_writer) = ws.split(tokio::io::split); let writer = Arc::new(Mutex::new(AddSend(og_writer))); - let (from_tx, _) = broadcast::channel::(usize::MAX / 2); + let (from_tx, _) = broadcast::channel::(MAX_SIZE); let senders_writer = writer.clone(); let mut sender = move |frame| { @@ -322,12 +328,12 @@ impl NexusClient { cancellation_token, api_key.to_api_key_with_key(&api_key_str.clone()), to_server, - from_server_tx + from_server_tx, )) - }, - None => { - Err(anyhow::anyhow!("Client's API key does not exist in the mini-store... sure ya have the right one?")) } + None => Err(anyhow::anyhow!( + "Client's API key does not exist in the mini-store... sure ya have the right one?" + )), } } else { Err(anyhow::anyhow!("Client is not connected yet!")) diff --git a/src/server/mod.rs b/src/server/mod.rs index fa464d5..46bfbd9 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -6,3 +6,4 @@ pub mod websockets; pub const DEAUTH_EVENT: &str = "nexus://com.reboot-codes.nexus/websockets/deauthorize"; pub const AUTH_HEADER: &str = "Authorization"; +pub const MAX_SIZE: usize = 8388608; diff --git a/src/server/models.rs b/src/server/models.rs index dc6df8b..f192287 100644 --- a/src/server/models.rs +++ b/src/server/models.rs @@ -1,23 +1,37 @@ use crate::{ arbiter::models::{ - ApiKey, ApiKeyWithKeyWithoutUID, ApiKeyWithoutUID, User, UserWithId - }, client::ClientStatus, user::NexusUser, utils::{ + ApiKey, + ApiKeyWithKeyWithoutUID, + ApiKeyWithoutUID, + User, + UserWithId, + }, + client::ClientStatus, + user::NexusUser, + utils::{ gen_api_key_with_check, gen_uid_with_check, - } + }, }; +use log::info; use serde::{ Deserialize, Serialize, }; -use tokio_util::sync::CancellationToken; use std::{ collections::HashMap, sync::Arc, }; -use tokio::sync::{broadcast, Mutex}; +use tokio::sync::{ + Mutex, + broadcast, +}; +use tokio_util::sync::CancellationToken; -use super::websockets::WsIn; +use super::{ + MAX_SIZE, + websockets::WsIn, +}; // TODO: Define defaults via `Default` trait impl. @@ -127,26 +141,28 @@ impl NexusStore { (ret, master_user) } - pub async fn add_user(&mut self, user_config: UserConfig, parent: Option) -> Result { + pub async fn add_user( + &mut self, + user_config: UserConfig, + parent: Option, + ) -> Result { let mut parent_id = None; let mut error = None; match parent.clone() { - Some(target_parent_id) => { - match self.users.lock().await.get(&target_parent_id) { - Some(_parent) => { - parent_id = parent.clone(); - }, - None => { - error = Some(anyhow::anyhow!("Parent ID does not exist in store!")); - } + Some(target_parent_id) => match self.users.lock().await.get(&target_parent_id) { + Some(_parent) => { + parent_id = parent.clone(); + } + None => { + error = Some(anyhow::anyhow!("Parent ID does not exist in store!")); } }, None => {} } match error { - Some(e) => { Err(e) }, + Some(e) => Err(e), None => { let mut key_ids = vec![]; let mut key_configs = vec![]; @@ -166,7 +182,7 @@ impl NexusStore { api_keys: key_ids.clone(), sessions: Arc::new(Mutex::new(HashMap::new())), parent_id: parent.clone(), - children: Vec::new() + children: Vec::new(), }, ); @@ -178,7 +194,7 @@ impl NexusStore { allowed_events_from: key_config.allowed_events_from.clone(), user_id: id.clone(), echo: key_config.echo, - proxy: key_config.proxy + proxy: key_config.proxy, }, ); } @@ -190,7 +206,7 @@ impl NexusStore { sessions: Arc::new(Mutex::new(HashMap::new())), parent_id: parent, children: Vec::new(), - id: id.clone() + id: id.clone(), }) } } @@ -204,20 +220,30 @@ impl NexusStore { allowed_events_to: vec![".*".to_string()], allowed_events_from: vec![".*".to_string()], echo: true, - proxy: true + proxy: true, }], }; self.add_user(ret.clone(), None).await.unwrap() } - pub async fn connect_user(&mut self, api_key_str: &String) -> Result<(NexusUser, broadcast::Sender, broadcast::Sender), anyhow::Error> { + pub async fn connect_user( + &mut self, + api_key_str: &String, + ) -> Result< + ( + NexusUser, + broadcast::Sender, + broadcast::Sender, + ), + anyhow::Error, + > { match self.api_keys.lock().await.get(&api_key_str.clone()) { Some(api_key) => { let status = Arc::new(Mutex::new(ClientStatus::new(true))); let cancellation_token = CancellationToken::new(); - let (to_server, _) = broadcast::channel(usize::MAX / 2); - let (from_server_tx, _) = broadcast::channel(usize::MAX / 2); + let (to_server, _) = broadcast::channel(MAX_SIZE); + let (from_server_tx, _) = broadcast::channel(MAX_SIZE); let user_to_server = to_server.clone(); let user_from_server_tx = from_server_tx.clone(); @@ -228,15 +254,15 @@ impl NexusStore { cancellation_token, api_key.to_api_key_with_key(&api_key_str.clone()), user_to_server, - user_from_server_tx + user_from_server_tx, ), to_server, - from_server_tx + from_server_tx, )) - }, - None => { - Err(anyhow::anyhow!("Client's API key does not exist in the store... sure ya have the right one?")) } + None => Err(anyhow::anyhow!( + "Client's API key does not exist in the store... sure ya have the right one?" + )), } } } diff --git a/src/server/websockets.rs b/src/server/websockets.rs index d4d53ef..8f4f89a 100644 --- a/src/server/websockets.rs +++ b/src/server/websockets.rs @@ -2,7 +2,6 @@ use crate::arbiter::models::{ ApiKeyWithKey, UserWithId, }; -use crate::server::DEAUTH_EVENT; use crate::server::models::{ Client, ClientWithId, @@ -10,6 +9,10 @@ use crate::server::models::{ NexusStore, Session, }; +use crate::server::{ + DEAUTH_EVENT, + MAX_SIZE, +}; use crate::utils::iso8601; use futures::{ SinkExt, @@ -70,7 +73,7 @@ pub async fn handle_ws_client( info!("Upgraded client: {}, to websocket connection!", ws_client.id.clone()); let (mut sender, mut receiver) = websocket.split(); - let (to_client_tx, mut to_client_rx) = broadcast::channel::(usize::MAX / 2); + let (to_client_tx, mut to_client_rx) = broadcast::channel::(MAX_SIZE); let mut deauthed = false; to_clients_tx.lock().await.insert(ws_client.id.clone(), to_client_tx); -- 2.51.2