creddy/src-tauri/src/state.rs

use std::collections::HashMap;
use std::sync::RwLock;

use serde::{Serialize, Deserialize};
use tokio::sync::oneshot::Sender;
use sqlx::{SqlitePool, sqlite::SqlitePoolOptions, sqlite::SqliteConnectOptions};
use sodiumoxide::crypto::{
        pwhash,
        pwhash::Salt, 
        secretbox, 
        secretbox::{Nonce, Key}
};
use tauri::async_runtime as runtime;

use crate::ipc;
use crate::errors::*;


#[derive(Serialize, Deserialize)]
#[serde(rename_all = "PascalCase")]
#[serde(untagged)]
pub enum Credentials {
    LongLived {
        #[serde(rename = "AccessKeyId")]
        access_key_id: String,
        #[serde(rename = "SecretAccessKey")]
        secret_access_key: String,
    },
    ShortLived {
        #[serde(rename = "AccessKeyId")]
        access_key_id: String,
        #[serde(rename = "SecretAccessKey")]
        secret_access_key: String,
        #[serde(rename = "Token")]
        token: String,
        #[serde(rename = "Expiration")]
        expiration: String,
    },
}


pub struct LockedCredentials {
    access_key_id: String,
    secret_key_enc: Vec<u8>,
    salt: Salt,
    nonce: Nonce,
}


pub enum Session {
    Unlocked(Credentials),
    Locked(LockedCredentials),
    Empty,
}


// #[derive(Serialize, Deserialize)]
// pub enum SessionStatus {
//   Unlocked,
//   Locked,
//   Empty,
// }


pub struct AppState {
    pub session: RwLock<Session>,
    pub request_count: RwLock<u64>,
    pub open_requests: RwLock<HashMap<u64, Sender<ipc::Approval>>>,
    pool: SqlitePool,
}

impl AppState {
    pub fn new() -> Result<Self, SetupError> {
        let conn_opts = SqliteConnectOptions::new()
            .filename("creddy.db")
            .create_if_missing(true);
        let pool_opts = SqlitePoolOptions::new();

        let pool: SqlitePool = runtime::block_on(pool_opts.connect_with(conn_opts))?;
        runtime::block_on(sqlx::migrate!().run(&pool))?;
        let creds = runtime::block_on(Self::load_creds(&pool))?;

        let state = AppState {
            session: RwLock::new(creds),
            request_count: RwLock::new(0),
            open_requests: RwLock::new(HashMap::new()),
            pool,
        };

        Ok(state)
    }

    async fn load_creds(pool: &SqlitePool) -> Result<Session, SetupError> {
        let res = sqlx::query!("SELECT * FROM credentials")
            .fetch_optional(pool)
            .await?;
        let row = match res {
            Some(r) => r,
            None => {return Ok(Session::Empty);}
        };

        let salt_buf: [u8; 32] = row.salt
            .try_into()
            .map_err(|_e| SetupError::InvalidRecord)?;
        let nonce_buf: [u8; 24] = row.nonce
            .try_into()
            .map_err(|_e| SetupError::InvalidRecord)?;

        let creds = LockedCredentials {
            access_key_id: row.access_key_id,
            secret_key_enc: row.secret_key_enc,
            salt: Salt(salt_buf),
            nonce: Nonce(nonce_buf),
        };
        Ok(Session::Locked(creds))
    }

    pub async fn save_creds(&self, creds: Credentials, passphrase: &str) -> Result<(), sqlx::error::Error> {
        let (key_id, secret_key) = match creds {
            Credentials::LongLived {access_key_id, secret_access_key} => {
                (access_key_id, secret_access_key)
            },
            _ => unreachable!(),
        };
        let salt = pwhash::gen_salt();
        let mut key_buf = [0; secretbox::KEYBYTES];
        pwhash::derive_key_interactive(&mut key_buf, passphrase.as_bytes(), &salt).unwrap();
        let key = Key(key_buf);
        // not sure we need both salt AND nonce given that we generate a
        // fresh salt every time we encrypt, but better safe than sorry
        let nonce = secretbox::gen_nonce();
        let key_enc = secretbox::seal(secret_key.as_bytes(), &nonce, &key);
        // insert into database
        Ok(())
    }

    pub fn register_request(&self, chan: Sender<ipc::Approval>) -> u64 {
        let count = {
            let mut c = self.request_count.write().unwrap();
            *c += 1;
            c
        };

        let mut open_requests = self.open_requests.write().unwrap();
        open_requests.insert(*count, chan); // `count` is the request id
        *count
    }

    pub fn send_response(&self, response: ipc::RequestResponse) -> Result<(), SendResponseError> {
        let mut open_requests = self.open_requests.write().unwrap();
        let chan = open_requests
            .remove(&response.id)
            .ok_or(SendResponseError::NotFound)
            ?;

        chan.send(response.approval)
            .map_err(|_e| SendResponseError::Abandoned)
    }

    pub async fn decrypt(&self, passphrase: &str) -> Result<(), UnlockError> {
        let session = self.session.read().unwrap();
        let locked = match *session {
            Session::Empty => {return Err(UnlockError::NoCredentials);},
            Session::Unlocked(_) => {return Err(UnlockError::NotLocked);},
            Session::Locked(ref c) => c,
        };

        let mut key_buf = [0; secretbox::KEYBYTES];
        // pretty sure this only fails if we're out of memory
        pwhash::derive_key_interactive(&mut key_buf, passphrase.as_bytes(), &locked.salt).unwrap();
        let decrypted = secretbox::open(&locked.secret_key_enc, &locked.nonce, &Key(key_buf))
            .map_err(|_e| UnlockError::BadPassphrase)?;

        let secret_str = String::from_utf8(decrypted).map_err(|_e| UnlockError::InvalidUtf8)?;
        let mut session = self.session.write().unwrap();
        let creds = Credentials::LongLived {
            access_key_id: locked.access_key_id.clone(),
            secret_access_key: secret_str,
        };
        *session = Session::Unlocked(creds);
        Ok(())
    }

    pub fn get_creds_serialized(&self) -> Result<String, GetCredentialsError> {
        let session = self.session.read().unwrap();
        match *session {
            Session::Unlocked(ref creds) => Ok(serde_json::to_string(creds).unwrap()),
            Session::Locked(_) => Err(GetCredentialsError::Locked),
            Session::Empty => Err(GetCredentialsError::Empty),
        }
    }
}