diff --git a/src/arbiter/models.rs b/src/arbiter/models.rs index c156221..2e6bfcb 100644 --- a/src/arbiter/models.rs +++ b/src/arbiter/models.rs @@ -1,4 +1,4 @@ -use crate::server::models::Session; +use crate::server::models::{Session, UserConfigWithId}; use serde::{ Deserialize, Serialize, @@ -124,3 +124,22 @@ pub struct ApiKeyWithKeyWithoutUID { pub echo: bool, pub proxy: bool } + +impl Into for ApiKey { + fn into(self) -> ApiKeyWithoutUID { + ApiKeyWithoutUID { + allowed_events_to: self.allowed_events_to, + allowed_events_from: self.allowed_events_from, + echo: self.echo, + proxy: self.proxy + } + } +} + +#[derive(Debug, Deserialize, Serialize, Clone)] +pub struct ApiKeyWithoutUID { + pub allowed_events_to: Vec, + pub allowed_events_from: Vec, + pub echo: bool, + pub proxy: bool +} diff --git a/src/server/listener.rs b/src/server/listener.rs index c7c4ec7..e4498a0 100644 --- a/src/server/listener.rs +++ b/src/server/listener.rs @@ -35,19 +35,14 @@ use std::net::{ use std::sync::Arc; use std::time::SystemTime; use thiserror::Error; -use tokio::sync::Mutex; -use tokio::sync::mpsc::{ - self, - UnboundedReceiver, - UnboundedSender, -}; +use tokio::sync::{broadcast, Mutex}; +use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; use url::Url; use warp::{ Filter, http::StatusCode, }; -use super::models::UserConfig; // example error response #[derive(Serialize, Debug)] @@ -268,20 +263,21 @@ pub async fn nexus_listener( port: u16, store: Arc, internal_client_senders: Vec<( - &UserConfig, - UnboundedSender, + &UserWithId, + usize, + broadcast::Sender, )>, - internal_client_recivers: Vec>, + internal_client_recivers: Vec>, cancellation_tokens: (CancellationToken, CancellationToken), ) { info!("Starting nexus on port: {}...", port); - let clients_tx: Arc>>> = + let clients_tx: Arc>>> = Arc::new(Mutex::new(HashMap::new())); // For everything connected via thread IPC, we still need a client or else everything freaks out. for client in internal_client_senders { let client_obj = Client { - api_key: client.0.api_keys[0].key.clone(), + api_key: client.0.api_keys[client.1].clone(), user_id: client.0.id.clone(), active: true, }; @@ -298,7 +294,7 @@ pub async fn nexus_listener( .lock() .await .insert(cid.clone(), client_obj.clone()); - clients_tx.lock().await.insert(cid.clone(), client.1); + clients_tx.lock().await.insert(cid.clone(), client.2); } let (from_client_tx, mut from_client_rx) = mpsc::unbounded_channel::(); @@ -482,7 +478,7 @@ pub async fn nexus_listener( debug!("Internal IPC Handle Exited!"); }, _ = async move { - while let Some(msg) = client.recv().await { + while let Ok(msg) = client.recv().await { match internal_ipc_store.api_keys.lock().await.get(&msg.api_key.to_string()) { Some(api_key) => { match internal_ipc_store.users.lock().await.get(&api_key.user_id) { diff --git a/src/server/models.rs b/src/server/models.rs index 9097f1e..dc6df8b 100644 --- a/src/server/models.rs +++ b/src/server/models.rs @@ -1,23 +1,23 @@ use crate::{ arbiter::models::{ - ApiKey, - ApiKeyWithKeyWithoutUID, - User, UserWithId, - }, - utils::{ + ApiKey, ApiKeyWithKeyWithoutUID, ApiKeyWithoutUID, User, UserWithId + }, client::ClientStatus, user::NexusUser, utils::{ gen_api_key_with_check, gen_uid_with_check, - }, + } }; use serde::{ Deserialize, Serialize, }; +use tokio_util::sync::CancellationToken; use std::{ collections::HashMap, sync::Arc, }; -use tokio::sync::Mutex; +use tokio::sync::{broadcast, Mutex}; + +use super::websockets::WsIn; // TODO: Define defaults via `Default` trait impl. @@ -97,7 +97,7 @@ pub struct UserConfigWithId { pub struct UserConfig { pub user_type: String, pub pretty_name: String, - pub api_keys: Vec, + pub api_keys: Vec, } // TODO: Add serialization/deserialization functions... @@ -148,9 +148,13 @@ impl NexusStore { match error { Some(e) => { Err(e) }, None => { - let mut key_ids: Vec = vec![]; + let mut key_ids = vec![]; + let mut key_configs = vec![]; for key_config in user_config.api_keys.iter() { - key_ids.push(key_config.key.clone()); + let key = gen_api_key_with_check(self).await; + + key_ids.push(key.clone()); + key_configs.push((key.clone(), key_config.clone())); } let id = gen_uid_with_check(self).await; @@ -166,9 +170,9 @@ impl NexusStore { }, ); - for key_config in user_config.api_keys.iter() { + for (key, key_config) in key_configs { self.api_keys.lock().await.insert( - key_config.key.clone(), + key.clone(), ApiKey { allowed_events_to: key_config.allowed_events_to.clone(), allowed_events_from: key_config.allowed_events_from.clone(), @@ -196,10 +200,9 @@ impl NexusStore { let ret = UserConfig { pretty_name: pretty_name.clone(), user_type: NexusStore::MASTER_USER_TYPE.to_string(), - api_keys: vec![ApiKeyWithKeyWithoutUID { + api_keys: vec![ApiKeyWithoutUID { allowed_events_to: vec![".*".to_string()], allowed_events_from: vec![".*".to_string()], - key: gen_api_key_with_check(self).await, echo: true, proxy: true }], @@ -207,4 +210,33 @@ impl NexusStore { 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> { + 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 user_to_server = to_server.clone(); + let user_from_server_tx = from_server_tx.clone(); + Ok(( + NexusUser::new( + false, + status, + cancellation_token, + api_key.to_api_key_with_key(&api_key_str.clone()), + user_to_server, + user_from_server_tx + ), + to_server, + from_server_tx + )) + }, + 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 3076ee0..809518c 100644 --- a/src/server/websockets.rs +++ b/src/server/websockets.rs @@ -29,7 +29,7 @@ use serde::{ use std::collections::HashMap; use std::sync::Arc; use std::time::SystemTime; -use tokio::sync::Mutex; +use tokio::sync::{broadcast, Mutex}; use tokio::sync::mpsc::{ self, UnboundedSender, @@ -52,7 +52,7 @@ pub async fn handle_ws_client( auth: (UserWithId, ApiKeyWithKey, ClientWithId, Session), ws: warp::ws::Ws, store: Arc>, - to_clients_tx: Arc>>>>, + to_clients_tx: Arc>>>>, from_clients_tx: Arc>, ) -> Result { let user = auth.0.clone(); diff --git a/src/user/mod.rs b/src/user/mod.rs index 49c8167..670491e 100644 --- a/src/user/mod.rs +++ b/src/user/mod.rs @@ -12,9 +12,9 @@ use crate::{arbiter::models::ApiKeyWithKey, client::ClientStatus, server::{model /// This struct basically just ensures that messages are formatted properly. #[derive(Debug, Clone)] pub struct NexusUser { - is_child: bool, + pub is_child: bool, connected: Arc>, - cancellation_token: CancellationToken, + pub cancellation_token: CancellationToken, api_key: ApiKeyWithKey, to_client: broadcast::Sender, from_server_tx: broadcast::Sender @@ -39,7 +39,7 @@ impl NexusUser { } } - pub fn send_msg(&self, kind: &String, message: &String) -> Result<(), anyhow::Error> { + pub fn send(&self, kind: &String, message: &String) -> Result<(), anyhow::Error> { match self.to_client.send(WsIn { kind: kind.clone(), message: message.clone(),