From b3ba882ac1f6cb5f8a5389e8513a9019b4aff4f9 Mon Sep 17 00:00:00 2001 From: Reboot-Codes Date: Fri, 4 Apr 2025 14:50:03 -0700 Subject: [PATCH] Clean up stupid errors --- src/arbiter/mod.rs | 165 +++++++++++++++++++-------------------- src/arbiter/models.rs | 22 +++--- src/server/listener.rs | 42 +++++++--- src/server/websockets.rs | 16 ++-- src/utils.rs | 18 ++--- 5 files changed, 135 insertions(+), 128 deletions(-) diff --git a/src/arbiter/mod.rs b/src/arbiter/mod.rs index 5fe00dd..17353f5 100644 --- a/src/arbiter/mod.rs +++ b/src/arbiter/mod.rs @@ -1,6 +1,5 @@ pub mod models; -use crate::utils::send_ipc_message; use log::{ debug, info, @@ -20,85 +19,85 @@ use crate::server::models::{ NexusStore, }; -pub async fn arbiter_main( - ipc_tx: UnboundedSender, - mut ipc_rx: UnboundedReceiver, - store: Arc, - user_config: Arc, - cancellation_tokens: (CancellationToken, CancellationToken), -) { - info!("Starting Arbiter..."); - - let init_store = Arc::new(store.clone()); - let init_user = Arc::new(user_config.clone()); - let (init_from_tx, mut init_from_rx) = unbounded_channel::(); - cancellation_tokens - .0 - .run_until_cancelled(async move { - let _ = send_ipc_message( - &init_store, - &init_user, - Arc::new(init_from_tx), - "clover://arbiter.clover.reboot-codes.com/status".to_string(), - "finished-init".to_string(), - ) - .await; - }) - .await; - - let ipc_recv_token = cancellation_tokens.0.clone(); - let ipc_recv_handle = tokio::task::spawn(async move { - tokio::select! { - _ = ipc_recv_token.cancelled() => { - debug!("ipc_recv exited"); - }, - _ = async move { - while let Some(msg) = ipc_rx.recv().await { - let kind = Url::parse(&msg.kind.clone()).unwrap(); - - // Verify that we care about this event. - if kind.host().unwrap() == url::Host::Domain("arbiter.clover.reboot-codes.com") { - debug!("Processing: {}", msg.kind.clone()); - } - } - } => {} - } - }); - - let ipc_trans_token = cancellation_tokens.0.clone(); - let ipc_trans_tx = Arc::new(ipc_tx.clone()); - let ipc_trans_handle = tokio::task::spawn(async move { - tokio::select! { - _ = async move { - while let Some(msg) = init_from_rx.recv().await { - match ipc_trans_tx.send(msg) { - Ok(_) => {}, - Err(_) => { - debug!("Failed to send message to IPC bus!"); - } - } - } - } => {}, - _ = ipc_trans_token.cancelled() => { - debug!("ipc_trans exited"); - } - } - }); - - let cleanup_token = cancellation_tokens.0.clone(); - tokio::select! { - _ = cleanup_token.cancelled() => { - ipc_recv_handle.abort(); - ipc_trans_handle.abort(); - - info!("Cleaning up users..."); - // TODO: Clean up registered users when server is shutting down. - - std::mem::drop(store); - - cancellation_tokens.1.cancel(); - } - } - - info!("Arbiter has stopped!"); -} +// pub async fn arbiter_main( +// ipc_tx: UnboundedSender, +// mut ipc_rx: UnboundedReceiver, +// store: Arc, +// user_config: Arc, +// cancellation_tokens: (CancellationToken, CancellationToken), +// ) { +// info!("Starting Arbiter..."); +// +// let init_store = Arc::new(store.clone()); +// let init_user = Arc::new(user_config.clone()); +// let (init_from_tx, mut init_from_rx) = unbounded_channel::(); +// cancellation_tokens +// .0 +// .run_until_cancelled(async move { +// let _ = send_ipc_message( +// &init_store, +// &init_user, +// Arc::new(init_from_tx), +// "clover://arbiter.clover.reboot-codes.com/status".to_string(), +// "finished-init".to_string(), +// ) +// .await; +// }) +// .await; +// +// let ipc_recv_token = cancellation_tokens.0.clone(); +// let ipc_recv_handle = tokio::task::spawn(async move { +// tokio::select! { +// _ = ipc_recv_token.cancelled() => { +// debug!("ipc_recv exited"); +// }, +// _ = async move { +// while let Some(msg) = ipc_rx.recv().await { +// let kind = Url::parse(&msg.kind.clone()).unwrap(); +// +// // Verify that we care about this event. +// if kind.host().unwrap() == url::Host::Domain("arbiter.clover.reboot-codes.com") { +// debug!("Processing: {}", msg.kind.clone()); +// } +// } +// } => {} +// } +// }); +// +// let ipc_trans_token = cancellation_tokens.0.clone(); +// let ipc_trans_tx = Arc::new(ipc_tx.clone()); +// let ipc_trans_handle = tokio::task::spawn(async move { +// tokio::select! { +// _ = async move { +// while let Some(msg) = init_from_rx.recv().await { +// match ipc_trans_tx.send(msg) { +// Ok(_) => {}, +// Err(_) => { +// debug!("Failed to send message to IPC bus!"); +// } +// } +// } +// } => {}, +// _ = ipc_trans_token.cancelled() => { +// debug!("ipc_trans exited"); +// } +// } +// }); +// +// let cleanup_token = cancellation_tokens.0.clone(); +// tokio::select! { +// _ = cleanup_token.cancelled() => { +// ipc_recv_handle.abort(); +// ipc_trans_handle.abort(); +// +// info!("Cleaning up users..."); +// // TODO: Clean up registered users when server is shutting down. +// +// std::mem::drop(store); +// +// cancellation_tokens.1.cancel(); +// } +// } +// +// info!("Arbiter has stopped!"); +// } diff --git a/src/arbiter/models.rs b/src/arbiter/models.rs index 2e6bfcb..d140ac1 100644 --- a/src/arbiter/models.rs +++ b/src/arbiter/models.rs @@ -1,4 +1,4 @@ -use crate::server::models::{Session, UserConfigWithId}; +use crate::server::models::Session; use serde::{ Deserialize, Serialize, @@ -31,7 +31,7 @@ impl User { user_type: self.user_type, pretty_name: self.pretty_name, parent_id: self.parent_id, - children: self.children + children: self.children, } } } @@ -55,7 +55,7 @@ impl Into for UserWithId { user_type: self.user_type, pretty_name: self.pretty_name, parent_id: self.parent_id, - children: self.children + children: self.children, } } } @@ -66,7 +66,7 @@ pub struct ApiKey { pub allowed_events_from: Vec, pub user_id: String, pub echo: bool, - pub proxy: bool + pub proxy: bool, } impl ApiKey { @@ -77,7 +77,7 @@ impl ApiKey { allowed_events_from: self.allowed_events_from.clone(), user_id: self.user_id.clone(), echo: self.echo, - proxy: self.proxy + proxy: self.proxy, } } } @@ -89,7 +89,7 @@ pub struct ApiKeyWithKey { pub allowed_events_from: Vec, pub user_id: String, pub echo: bool, - pub proxy: bool + pub proxy: bool, } impl Into for ApiKeyWithKey { @@ -99,7 +99,7 @@ impl Into for ApiKeyWithKey { allowed_events_from: self.allowed_events_from, user_id: self.user_id, echo: self.echo, - proxy: self.proxy + proxy: self.proxy, } } } @@ -111,7 +111,7 @@ impl Into for ApiKeyWithKey { allowed_events_to: self.allowed_events_to, allowed_events_from: self.allowed_events_from, echo: self.echo, - proxy: self.proxy + proxy: self.proxy, } } } @@ -122,7 +122,7 @@ pub struct ApiKeyWithKeyWithoutUID { pub allowed_events_to: Vec, pub allowed_events_from: Vec, pub echo: bool, - pub proxy: bool + pub proxy: bool, } impl Into for ApiKey { @@ -131,7 +131,7 @@ impl Into for ApiKey { allowed_events_to: self.allowed_events_to, allowed_events_from: self.allowed_events_from, echo: self.echo, - proxy: self.proxy + proxy: self.proxy, } } } @@ -141,5 +141,5 @@ pub struct ApiKeyWithoutUID { pub allowed_events_to: Vec, pub allowed_events_from: Vec, pub echo: bool, - pub proxy: bool + pub proxy: bool, } diff --git a/src/server/listener.rs b/src/server/listener.rs index e4498a0..48484a6 100644 --- a/src/server/listener.rs +++ b/src/server/listener.rs @@ -10,9 +10,14 @@ use crate::server::models::{ Session, }; use crate::server::websockets::handle_ws_client; -use crate::server::{AUTH_HEADER, DEAUTH_EVENT}; +use crate::server::{ + AUTH_HEADER, + DEAUTH_EVENT, +}; use crate::utils::{ - gen_cid_with_check, gen_message_id_with_check, iso8601 + gen_cid_with_check, + gen_message_id_with_check, + iso8601, }; use log::{ debug, @@ -35,8 +40,11 @@ use std::net::{ use std::sync::Arc; use std::time::SystemTime; use thiserror::Error; -use tokio::sync::{broadcast, Mutex}; use tokio::sync::mpsc; +use tokio::sync::{ + Mutex, + broadcast, +}; use tokio_util::sync::CancellationToken; use url::Url; use warp::{ @@ -44,6 +52,8 @@ use warp::{ http::StatusCode, }; +use super::websockets::WsIn; + // example error response #[derive(Serialize, Debug)] struct ApiErrorResult { @@ -54,7 +64,7 @@ struct ApiErrorResult { pub struct IPCWithKey { kind: String, message: String, - api_key: String + api_key: String, } // errors thrown by handlers and custom filters, @@ -262,12 +272,8 @@ pub struct ServerHealth { pub async fn nexus_listener( port: u16, store: Arc, - internal_client_senders: Vec<( - &UserWithId, - usize, - broadcast::Sender, - )>, - internal_client_recivers: Vec>, + internal_client_senders: Vec<(Arc, usize, broadcast::Sender)>, + internal_client_recivers: Vec<(Arc, usize, broadcast::Sender)>, cancellation_tokens: (CancellationToken, CancellationToken), ) { info!("Starting nexus on port: {}...", port); @@ -467,7 +473,7 @@ pub async fn nexus_listener( let mut internal_ipc_handles = vec![]; - for mut client in internal_client_recivers { + for client in internal_client_recivers { // Internal IPC Handles let internal_ipc_store = Arc::new(store.clone()); let internal_ipc_tx = Arc::new(from_client_tx.clone()); @@ -478,8 +484,18 @@ pub async fn nexus_listener( debug!("Internal IPC Handle Exited!"); }, _ = async move { - while let Ok(msg) = client.recv().await { - match internal_ipc_store.api_keys.lock().await.get(&msg.api_key.to_string()) { + let mut client_rx = client.2.subscribe(); + while let Ok(msg) = client_rx.recv().await { + let msg_api_key = match msg.api_key { + Some(api_key) => { + &api_key.clone() + }, + None => { + client.0.api_keys.get(client.1).unwrap() + } + }; + + match internal_ipc_store.api_keys.lock().await.get(msg_api_key) { Some(api_key) => { match internal_ipc_store.users.lock().await.get(&api_key.user_id) { Some(_user) => { diff --git a/src/server/websockets.rs b/src/server/websockets.rs index 809518c..d4d53ef 100644 --- a/src/server/websockets.rs +++ b/src/server/websockets.rs @@ -2,6 +2,7 @@ use crate::arbiter::models::{ ApiKeyWithKey, UserWithId, }; +use crate::server::DEAUTH_EVENT; use crate::server::models::{ Client, ClientWithId, @@ -9,7 +10,6 @@ use crate::server::models::{ NexusStore, Session, }; -use crate::server::DEAUTH_EVENT; use crate::utils::iso8601; use futures::{ SinkExt, @@ -29,10 +29,10 @@ use serde::{ use std::collections::HashMap; use std::sync::Arc; use std::time::SystemTime; -use tokio::sync::{broadcast, Mutex}; -use tokio::sync::mpsc::{ - self, - UnboundedSender, +use tokio::sync::mpsc::UnboundedSender; +use tokio::sync::{ + Mutex, + broadcast, }; use url::Url; use uuid::Uuid; @@ -45,7 +45,7 @@ use warp::filters::ws::{ pub struct WsIn { pub kind: String, pub message: String, - pub api_key: Option + pub api_key: Option, } pub async fn handle_ws_client( @@ -70,7 +70,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) = mpsc::unbounded_channel::(); + let (to_client_tx, mut to_client_rx) = broadcast::channel::(usize::MAX / 2); let mut deauthed = false; to_clients_tx.lock().await.insert(ws_client.id.clone(), to_client_tx); @@ -210,7 +210,7 @@ pub async fn handle_ws_client( let send_api_key = Arc::new(api_key.clone()); let send_client = Arc::new(client.clone()); let send_handle = tokio::task::spawn(async move { - while let Some(msg) = to_client_rx.recv().await { + while let Ok(msg) = to_client_rx.recv().await { if msg.kind == Url::parse(DEAUTH_EVENT) .unwrap() .query_pairs_mut() diff --git a/src/utils.rs b/src/utils.rs index b770aa8..0cf1e96 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,8 +1,4 @@ -use crate::server::models::{ - CoreUserConfig, - IPCMessageWithId, - NexusStore, -}; +use crate::server::models::NexusStore; use api_key::types::{ ApiKeyResults, Default, @@ -14,15 +10,11 @@ use chrono::prelude::{ }; use os_path::OsPath; use simple_error::SimpleError; -use std::{ - hash::{ - DefaultHasher, - Hash, - Hasher, - }, - sync::Arc, +use std::hash::{ + DefaultHasher, + Hash, + Hasher, }; -use tokio::sync::mpsc::UnboundedSender; use tokio::{ fs, io::AsyncReadExt, -- 2.51.2