path: root/atuin-dotfiles/src/store.rs

use std::collections::BTreeMap;

use atuin_client::record::sqlite_store::SqliteStore;
// Sync aliases
// This will be noticeable similar to the kv store, though I expect the two shall diverge
// While we will support a range of shell config, I'd rather have a larger number of small records
// + stores, rather than one mega config store.
use atuin_common::record::{DecryptedData, Host, HostId};
use eyre::{bail, ensure, eyre, Result};

use atuin_client::record::encryption::PASETO_V4;
use atuin_client::record::store::Store;

use crate::shell::Alias;

const CONFIG_SHELL_ALIAS_VERSION: &str = "v0";
const CONFIG_SHELL_ALIAS_TAG: &str = "config-shell-alias";
const CONFIG_SHELL_ALIAS_FIELD_MAX_LEN: usize = 20000; // 20kb max total len, way more than should be needed.

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AliasRecord {
    Create(Alias),  // create a full record
    Delete(String), // delete by name
}

impl AliasRecord {
    pub fn serialize(&self) -> Result<DecryptedData> {
        use rmp::encode;

        let mut output = vec![];

        match self {
            AliasRecord::Create(alias) => {
                encode::write_u8(&mut output, 0)?; // create
                encode::write_array_len(&mut output, 2)?; // 2 fields

                encode::write_str(&mut output, alias.name.as_str())?;
                encode::write_str(&mut output, alias.value.as_str())?;
            }
            AliasRecord::Delete(name) => {
                encode::write_u8(&mut output, 1)?; // delete
                encode::write_array_len(&mut output, 1)?; // 1 field

                encode::write_str(&mut output, name.as_str())?;
            }
        }

        Ok(DecryptedData(output))
    }

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

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

        match version {
            CONFIG_SHELL_ALIAS_VERSION => {
                let mut bytes = decode::Bytes::new(&data.0);

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

                match record_type {
                    // create
                    0 => {
                        let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?;
                        ensure!(
                            nfields == 2,
                            "too many entries in v0 shell alias create record"
                        );

                        let bytes = bytes.remaining_slice();

                        let (key, bytes) =
                            decode::read_str_from_slice(bytes).map_err(error_report)?;
                        let (value, bytes) =
                            decode::read_str_from_slice(bytes).map_err(error_report)?;

                        if !bytes.is_empty() {
                            bail!("trailing bytes in encoded shell alias record. malformed")
                        }

                        Ok(AliasRecord::Create(Alias {
                            name: key.to_owned(),
                            value: value.to_owned(),
                        }))
                    }

                    // delete
                    1 => {
                        let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?;
                        ensure!(
                            nfields == 1,
                            "too many entries in v0 shell alias delete record"
                        );

                        let bytes = bytes.remaining_slice();

                        let (key, bytes) =
                            decode::read_str_from_slice(bytes).map_err(error_report)?;

                        if !bytes.is_empty() {
                            bail!("trailing bytes in encoded shell alias record. malformed")
                        }

                        Ok(AliasRecord::Delete(key.to_owned()))
                    }

                    n => {
                        bail!("unknown AliasRecord type {n}")
                    }
                }
            }
            _ => {
                bail!("unknown version {version:?}")
            }
        }
    }
}

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

impl AliasStore {
    // will want to init the actual kv store when that is done
    pub fn new(store: SqliteStore, host_id: HostId, encryption_key: [u8; 32]) -> AliasStore {
        AliasStore {
            store,
            host_id,
            encryption_key,
        }
    }

    pub async fn set(&self, name: &str, value: &str) -> Result<()> {
        if name.len() + value.len() > CONFIG_SHELL_ALIAS_FIELD_MAX_LEN {
            return Err(eyre!(
                "alias record too large: max len {} bytes",
                CONFIG_SHELL_ALIAS_FIELD_MAX_LEN
            ));
        }

        let record = AliasRecord::Create(Alias {
            name: name.to_string(),
            value: value.to_string(),
        });

        let bytes = record.serialize()?;

        let idx = self
            .store
            .last(self.host_id, CONFIG_SHELL_ALIAS_TAG)
            .await?
            .map_or(0, |entry| entry.idx + 1);

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

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

        Ok(())
    }

    pub async fn delete(&self, name: &str) -> Result<()> {
        if name.len() > CONFIG_SHELL_ALIAS_FIELD_MAX_LEN {
            return Err(eyre!(
                "alias record too large: max len {} bytes",
                CONFIG_SHELL_ALIAS_FIELD_MAX_LEN
            ));
        }

        let record = AliasRecord::Delete(name.to_string());

        let bytes = record.serialize()?;

        let idx = self
            .store
            .last(self.host_id, CONFIG_SHELL_ALIAS_TAG)
            .await?
            .map_or(0, |entry| entry.idx + 1);

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

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

        Ok(())
    }

    pub async fn aliases(&self) -> Result<Vec<Alias>> {
        let mut build = BTreeMap::new();

        // this is sorted, oldest to newest
        let tagged = self.store.all_tagged(CONFIG_SHELL_ALIAS_TAG).await?;

        for record in tagged {
            let version = record.version.clone();

            let decrypted = match version.as_str() {
                CONFIG_SHELL_ALIAS_VERSION => record.decrypt::<PASETO_V4>(&self.encryption_key)?,
                version => bail!("unknown version {version:?}"),
            };

            let ar = AliasRecord::deserialize(&decrypted.data, version.as_str())?;

            match ar {
                AliasRecord::Create(a) => {
                    build.insert(a.name.clone(), a);
                }
                AliasRecord::Delete(d) => {
                    build.remove(&d);
                }
            }
        }

        Ok(build.into_values().collect())
    }
}

#[cfg(test)]
pub(crate) fn test_sqlite_store_timeout() -> f64 {
    std::env::var("ATUIN_TEST_SQLITE_STORE_TIMEOUT")
        .ok()
        .and_then(|x| x.parse().ok())
        .