path: root/atuin-client/src/history/store.rs

use std::collections::HashSet;

use eyre::{bail, eyre, Result};
use rmp::decode::Bytes;

use crate::{
    database::{self, Database},
    record::{encryption::PASETO_V4, sqlite_store::SqliteStore, store::Store},
};
use atuin_common::record::{DecryptedData, Host, HostId, Record, RecordId, RecordIdx};

use super::{History, HistoryId, HISTORY_TAG, HISTORY_VERSION};

#[derive(Debug)]
pub struct HistoryStore {
    pub store: SqliteStore,
    pub host_id: HostId,
    pub encryption_key: [u8; 32],
}

#[derive(Debug, Eq, PartialEq, Clone)]
pub enum HistoryRecord {
    Create(History),   // Create a history record
    Delete(HistoryId), // Delete a history record, identified by ID
}

impl HistoryRecord {
    /// Serialize a history record, returning DecryptedData
    /// The record will be of a certain type
    /// We map those like so:
    ///
    /// HistoryRecord::Create -> 0
    /// HistoryRecord::Delete-> 1
    ///
    /// This numeric identifier is then written as the first byte to the buffer. For history, we
    /// append the serialized history right afterwards, to avoid having to handle serialization
    /// twice.
    ///
    /// Deletion simply refers to the history by ID
    pub fn serialize(&self) -> Result<DecryptedData> {
        // probably don't actually need to use rmp here, but if we ever need to extend it, it's a
        // nice wrapper around raw byte stuff
        use rmp::encode;

        let mut output = vec![];

        match self {
            HistoryRecord::Create(history) => {
                // 0 -> a history create
                encode::write_u8(&mut output, 0)?;

                let bytes = history.serialize()?;

                encode::write_bin(&mut output, &bytes.0)?;
            }
            HistoryRecord::Delete(id) => {
                // 1 -> a history delete
                encode::write_u8(&mut output, 1)?;
                encode::write_str(&mut output, id.0.as_str())?;
            }
        };

        Ok(DecryptedData(output))
    }

    pub fn deserialize(bytes: &DecryptedData, version: &str) -> Result<Self> {
        use rmp::decode;

        fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report {
            eyre!("{err:?}")
        }

        let mut bytes = Bytes::new(&bytes.0);

        let record_type = decode::read_u8(&mut bytes).map_err(error_report)?;

        match record_type {
            // 0 -> HistoryRecord::Create
            0 => {
                // not super useful to us atm, but perhaps in the future
                // written by write_bin above
                let _ = decode::read_bin_len(&mut bytes).map_err(error_report)?;

                let record = History::deserialize(bytes.remaining_slice(), version)?;

                Ok(HistoryRecord::Create(record))
            }

            // 1 -> HistoryRecord::Delete
            1 => {
                let bytes = bytes.remaining_slice();
                let (id, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?;

                if !bytes.is_empty() {
                    bail!(
                        "trailing bytes decoding HistoryRecord::Delete - malformed? got {bytes:?}"
                    );
                }

                Ok(HistoryRecord::Delete(id.to_string().into()))
            }

            n => {
                bail!("unknown HistoryRecord type {n}")
            }
        }
    }
}

impl HistoryStore {
    pub fn new(store: SqliteStore, host_id: HostId, encryption_key: [u8; 32]) -> Self {
        HistoryStore {
            store,
            host_id,
            encryption_key,
        }
    }

    async fn push_record(&self, record: HistoryRecord) -> Result<(RecordId, RecordIdx)> {
        let bytes = record.serialize()?;
        let idx = self
            .store
            .last(self.host_id, HISTORY_TAG)
            .await?
            .map_or(0, |p| p.idx + 1);

        let record = Record::builder()
            .host(Host::new(self.host_id))
            .version(HISTORY_VERSION.to_string())
            .tag(HISTORY_TAG.to_string())
            .idx(idx)
            .data(bytes)
            .build();

        let id = record.id;

        self.store
            .push(&record.encrypt::<PASETO_V4>(&self.encryption_key))
            .await?;

        Ok((id, idx))
    }

    pub async fn delete(&self, id: HistoryId) -> Result<(RecordId, RecordIdx)> {
        let record = HistoryRecord::Delete(id);

        self.push_record(record).await
    }

    pub async fn push(&self, history: History) -> Result<(RecordId, RecordIdx)> {
        // TODO(ellie): move the history store to its own file
        // it's tiny rn so fine as is
        let record = HistoryRecord::Create(history);

        self.push_record(record).await
    }

    pub async fn history(&self) -> Result<Vec<HistoryRecord>> {
        // Atm this loads all history into memory
        // Not ideal as that is potentially quite a lot, although history will be small.
        let records = self.store.all_tagged(HISTORY_TAG).await?;
        let mut ret = Vec::with_capacity(records.len());

        for record in records.into_iter() {
            let hist = match record.version.as_str() {
                HISTORY_VERSION => {
                    let decrypted = record.decrypt::<PASETO_V4>(&self.encryption_key);

                    let decrypted = match decrypted {
                        Ok(d) => d,
                        Err(e) => {
                            println!("failed to decrypt history: {e}");
                            continue;
                        }
                    };

                    HistoryRecord::deserialize(&decrypted.data, HISTORY_VERSION)
                }
                version => bail!("unknown history version {version:?}"),
            }?;

            ret.push(hist);
        }

        Ok(ret)
    }

    pub async fn build(&self, database: &dyn Database) -> Result<()> {
        // I'd like to change how we rebuild and not couple this with the database, but need to
        // consider the structure more deeply. This will be easy to change.

        // TODO(ellie): page or iterate this
        let history = self.history().await?;

        // In theory we could flatten this here
        // The current issue is that the database may have history in it already, from the old sync
        // This didn't actually delete old history
        // If we're sure we have a DB only maintained by the new store, we can flatten
        // create/delete before we even get to sqlite
        let mut creates = Vec::new();
        let mut deletes = Vec::new();

        for i in history {
            match i {
                HistoryRecord::Create(h) => {
                    creates.push(h);
                }
                HistoryRecord::Delete(id) => {
                    deletes.push(id);
                }
            }
        }

        database.save_bulk(&creates).await?;
        database.delete_rows(&deletes).await?;

        Ok(())
    }

    pub async fn incremental_build(&self, database: &dyn Database, ids: &[RecordId]) -> Result<()> {
        for id in ids {
            let record = self.store.get(*id).await?;

            if record.tag != HISTORY_TAG {
                continue;
            }

            let decrypted = record.decrypt::<PASETO_V4>(&self.encr