featherChat/warzone/crates/warzone-client/src/storage.rs

//! Local sled database: sessions, pre-keys, message history.

use anyhow::{Context, Result};
use warzone_protocol::ratchet::RatchetState;
use warzone_protocol::types::Fingerprint;
use x25519_dalek::StaticSecret;

pub struct LocalDb {
    sessions: sled::Tree,
    pre_keys: sled::Tree,
    _db: sled::Db,
}

impl LocalDb {
    pub fn open() -> Result<Self> {
        let path = crate::keystore::data_dir().join("db");
        let db = sled::open(&path).context("failed to open local database")?;
        let sessions = db.open_tree("sessions")?;
        let pre_keys = db.open_tree("pre_keys")?;
        Ok(LocalDb {
            sessions,
            pre_keys,
            _db: db,
        })
    }

    /// Save a ratchet session for a peer.
    pub fn save_session(&self, peer: &Fingerprint, state: &RatchetState) -> Result<()> {
        let key = peer.to_hex();
        let data = bincode::serialize(state).context("failed to serialize session")?;
        self.sessions.insert(key.as_bytes(), data)?;
        self.sessions.flush()?;
        Ok(())
    }

    /// Load a ratchet session for a peer.
    pub fn load_session(&self, peer: &Fingerprint) -> Result<Option<RatchetState>> {
        let key = peer.to_hex();
        match self.sessions.get(key.as_bytes())? {
            Some(data) => {
                let state = bincode::deserialize(&data)
                    .context("failed to deserialize session")?;
                Ok(Some(state))
            }
            None => Ok(None),
        }
    }

    /// Store the signed pre-key secret (for X3DH respond).
    pub fn save_signed_pre_key(&self, id: u32, secret: &StaticSecret) -> Result<()> {
        let key = format!("spk:{}", id);
        self.pre_keys
            .insert(key.as_bytes(), secret.to_bytes().as_slice())?;
        self.pre_keys.flush()?;
        Ok(())
    }

    /// Load the signed pre-key secret.
    pub fn load_signed_pre_key(&self, id: u32) -> Result<Option<StaticSecret>> {
        let key = format!("spk:{}", id);
        match self.pre_keys.get(key.as_bytes())? {
            Some(data) => {
                let mut bytes = [0u8; 32];
                bytes.copy_from_slice(&data);
                Ok(Some(StaticSecret::from(bytes)))
            }
            None => Ok(None),
        }
    }

    /// Store a one-time pre-key secret.
    pub fn save_one_time_pre_key(&self, id: u32, secret: &StaticSecret) -> Result<()> {
        let key = format!("otpk:{}", id);
        self.pre_keys
            .insert(key.as_bytes(), secret.to_bytes().as_slice())?;
        self.pre_keys.flush()?;
        Ok(())
    }

    /// Load and remove a one-time pre-key secret.
    pub fn take_one_time_pre_key(&self, id: u32) -> Result<Option<StaticSecret>> {
        let key = format!("otpk:{}", id);
        match self.pre_keys.remove(key.as_bytes())? {
            Some(data) => {
                let mut bytes = [0u8; 32];
                bytes.copy_from_slice(&data);
                self.pre_keys.flush()?;
                Ok(Some(StaticSecret::from(bytes)))
            }
            None => Ok(None),
        }
    }
}