use std::collections::{HashMap, HashSet}; use std::time::Duration; use tokio::{ sync::oneshot::Sender, sync::RwLock, time::sleep, }; use sqlx::SqlitePool; use tauri::async_runtime as runtime; use tauri::Manager; use crate::app::APP; use crate::credentials::{ Session, BaseCredentials, SessionCredentials, }; use crate::{config, config::AppConfig}; use crate::ipc::{self, Approval}; use crate::clientinfo::Client; use crate::errors::*; use crate::server::Server; #[derive(Debug)] pub struct AppState { pub config: RwLock, pub session: RwLock, pub request_count: RwLock, pub open_requests: RwLock>>, pub bans: RwLock>>, server: RwLock, pool: sqlx::SqlitePool, } impl AppState { pub fn new(config: AppConfig, session: Session, server: Server, pool: SqlitePool) -> AppState { AppState { config: RwLock::new(config), session: RwLock::new(session), request_count: RwLock::new(0), open_requests: RwLock::new(HashMap::new()), bans: RwLock::new(HashSet::new()), server: RwLock::new(server), pool, } } pub async fn new_creds(&self, base_creds: BaseCredentials, passphrase: &str) -> Result<(), UnlockError> { let locked = base_creds.encrypt(passphrase)?; // do this first so that if it fails we don't save bad credentials self.new_session(base_creds).await?; locked.save(&self.pool).await?; Ok(()) } pub async fn update_config(&self, new_config: AppConfig) -> Result<(), SetupError> { let mut live_config = self.config.write().await; if new_config.start_on_login != live_config.start_on_login { config::set_auto_launch(new_config.start_on_login)?; } if new_config.listen_addr != live_config.listen_addr || new_config.listen_port != live_config.listen_port { let mut sv = self.server.write().await; sv.rebind(new_config.listen_addr, new_config.listen_port).await?; } new_config.save(&self.pool).await?; *live_config = new_config; Ok(()) } pub async fn register_request(&self, chan: Sender) -> u64 { let count = { let mut c = self.request_count.write().await; *c += 1; c }; let mut open_requests = self.open_requests.write().await; open_requests.insert(*count, chan); // `count` is the request id *count } pub async fn unregister_request(&self, id: u64) { let mut open_requests = self.open_requests.write().await; open_requests.remove(&id); } pub async fn req_count(&self) -> usize { let open_requests = self.open_requests.read().await; open_requests.len() } pub async fn send_response(&self, response: ipc::RequestResponse) -> Result<(), SendResponseError> { if let Approval::Approved = response.approval { let mut session = self.session.write().await; session.renew_if_expired().await?; } let mut open_requests = self.open_requests.write().await; let chan = open_requests .remove(&response.id) .ok_or(SendResponseError::NotFound) ?; chan.send(response.approval) .map_err(|_e| SendResponseError::Abandoned) } pub async fn add_ban(&self, client: Option) { let mut bans = self.bans.write().await; bans.insert(client.clone()); runtime::spawn(async move { sleep(Duration::from_secs(5)).await; let app = APP.get().unwrap(); let state = app.state::(); let mut bans = state.bans.write().await; bans.remove(&client); }); } pub async fn is_banned(&self, client: &Option) -> bool { self.bans.read().await.contains(&client) } pub async fn unlock(&self, passphrase: &str) -> Result<(), UnlockError> { let base_creds = match *self.session.read().await { Session::Empty => {return Err(UnlockError::NoCredentials);}, Session::Unlocked{..} => {return Err(UnlockError::NotLocked);}, Session::Locked(ref locked) => locked.decrypt(passphrase)?, }; // Read lock is dropped here, so this doesn't deadlock self.new_session(base_creds).await?; Ok(()) } pub async fn serialize_base_creds(&self) -> Result { let session = self.session.read().await; match *session { Session::Unlocked{ref base, ..} => Ok(serde_json::to_string(base).unwrap()), Session::Locked(_) => Err(GetCredentialsError::Locked), Session::Empty => Err(GetCredentialsError::Empty), } } pub async fn serialize_session_creds(&self) -> Result { let session = self.session.read().await; match *session { Session::Unlocked{ref session, ..} => Ok(serde_json::to_string(session).unwrap()), Session::Locked(_) => Err(GetCredentialsError::Locked), Session::Empty => Err(GetCredentialsError::Empty), } } async fn new_session(&self, base: BaseCredentials) -> Result<(), GetSessionError> { let session = SessionCredentials::from_base(&base).await?; let mut app_session = self.session.write().await; *app_session = Session::Unlocked {base, session}; Ok(()) } }