creddy/src-tauri/src/state.rs

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<AppConfig>,
    pub session: RwLock<Session>,
    pub request_count: RwLock<u64>,
    pub open_requests: RwLock<HashMap<u64, Sender<ipc::Approval>>>,
    pub bans: RwLock<std::collections::HashSet<Option<Client>>>,
    server: RwLock<Server>,
    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<ipc::Approval>) -> 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<Client>) {
        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::<AppState>();
            let mut bans = state.bans.write().await;
            bans.remove(&client);
        });
    }

    pub async fn is_banned(&self, client: &Option<Client>) -> 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<String, GetCredentialsError> {
        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<String, GetCredentialsError> {
        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(())
    }
}