path: root/crates/turtle/src/atuin_common/record.rs

use std::collections::HashMap;

use eyre::Result;
use serde::{Deserialize, Serialize};
use typed_builder::TypedBuilder;
use uuid::Uuid;

#[derive(Clone, Debug, PartialEq)]
pub(crate) struct DecryptedData(pub(crate) Vec<u8>);

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub(crate) struct EncryptedData {
    pub(crate) data: String,
    pub(crate) content_encryption_key: String,
}

#[derive(Debug, PartialEq, PartialOrd, Ord, Eq)]
pub(crate) struct Diff {
    pub(crate) host: HostId,
    pub(crate) tag: String,
    pub(crate) local: Option<RecordIdx>,
    pub(crate) remote: Option<RecordIdx>,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
pub(crate) struct Host {
    pub(crate) id: HostId,
    pub(crate) name: String,
}

impl Host {
    pub(crate) fn new(id: HostId) -> Self {
        Self {
            id,
            name: String::new(),
        }
    }
}

new_uuid!(RecordId);
new_uuid!(HostId);

pub(crate) type RecordIdx = u64;

/// A single record stored inside of our local database
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TypedBuilder)]
pub(crate) struct Record<Data> {
    /// a unique ID
    #[builder(default = RecordId(crate::atuin_common::utils::uuid_v7()))]
    pub(crate) id: RecordId,

    /// The integer record ID. This is only unique per (host, tag).
    pub(crate) idx: RecordIdx,

    /// The unique ID of the host.
    // TODO(ellie): Optimize the storage here. We use a bunch of IDs, and currently store
    // as strings. I would rather avoid normalization, so store as UUID binary instead of
    // encoding to a string and wasting much more storage.
    pub(crate) host: Host,

    /// The creation time in nanoseconds since unix epoch
    #[builder(default = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as u64)]
    pub(crate) timestamp: u64,

    /// The version the data in the entry conforms to
    // However we want to track versions for this tag, eg v2
    pub(crate) version: String,

    /// The type of data we are storing here. Eg, "history"
    pub(crate) tag: String,

    /// Some data. This can be anything you wish to store. Use the tag field to know how to handle it.
    pub(crate) data: Data,
}

/// Extra data from the record that should be encoded in the data
#[derive(Debug, Copy, Clone)]
pub(crate) struct AdditionalData<'a> {
    pub(crate) id: &'a RecordId,
    pub(crate) idx: &'a u64,
    pub(crate) version: &'a str,
    pub(crate) tag: &'a str,
    pub(crate) host: &'a HostId,
}

/// An index representing the current state of the record stores
/// This can be both remote, or local, and compared in either direction
#[derive(Debug, Serialize, Deserialize)]
pub(crate) struct RecordStatus {
    // A map of host -> tag -> max(idx)
    pub(crate) hosts: HashMap<HostId, HashMap<String, RecordIdx>>,
}

impl Default for RecordStatus {
    fn default() -> Self {
        Self::new()
    }
}

impl Extend<(HostId, String, RecordIdx)> for RecordStatus {
    fn extend<T: IntoIterator<Item = (HostId, String, RecordIdx)>>(&mut self, iter: T) {
        for (host, tag, tail_idx) in iter {
            self.set_raw(host, tag, tail_idx);
        }
    }
}

impl RecordStatus {
    pub(crate) fn new() -> Self {
        Self {
            hosts: HashMap::new(),
        }
    }

    /// Insert a new tail record into the store
    pub(crate) fn set_raw(&mut self, host: HostId, tag: String, tail_id: RecordIdx) {
        self.hosts.entry(host).or_default().insert(tag, tail_id);
    }

    pub(crate) fn get(&self, host: HostId, tag: &str) -> Option<RecordIdx> {
        self.hosts.get(&host).and_then(|v| v.get(tag)).copied()
    }

    /// Diff this index with another, likely remote index.
    /// The two diffs can then be reconciled, and the optimal change set calculated
    /// Returns a tuple, with (host, tag, Option(OTHER))
    /// OTHER is set to the value of the idx on the other machine. If it is greater than our index,
    /// then we need to do some downloading. If it is smaller, then we need to do some uploading
    /// Note that we cannot upload if we are not the owner of the record store - hosts can only
    /// write to their own store.
    pub(crate) fn diff(&self, other: &Self) -> Vec<Diff> {
        let mut ret = Vec::new();

        // First, we check if other has everything that self has
        for (host, tag_map) in &self.hosts {
            for (tag, idx) in tag_map {
                match other.get(*host, tag) {
                    // The other store is all up to date! No diff.
                    Some(t) if t.eq(idx) => (),

                    // The other store does exist, and it is either ahead or behind us. A diff regardless
                    Some(t) => ret.push(Diff {
                        host: *host,
                        tag: tag.clone(),
                        local: Some(*idx),
                        remote: Some(t),
                    }),

                    // The other store does not exist :O
                    None => ret.push(Diff {
                        host: *host,
                        tag: tag.clone(),
                        local: Some(*idx),
                        remote: None,
                    }),
                }
            }
        }

        // At this point, there is a single case we have not yet considered.
        // If the other store knows of a tag that we are not yet aware of, then the diff will be missed

        // account for that!
        for (host, tag_map) in &other.hosts {
            for (tag, idx) in tag_map {
                match self.get(*host, tag) {
                    // If we have this host/tag combo, the comparison and diff will have already happened above
                    Some(_) => (),

                    None => ret.push(Diff {
                        host: *host,
                        tag: tag.clone(),
                        remote: Some(*idx),
                        local: None,
                    }),
                }
            }
        }

        // Stability is a nice property to have
        ret.sort();
        ret
    }
}

pub(crate) trait Encryption {
    fn re_encrypt(
        data: EncryptedData,
        ad: AdditionalData<'_>,
        old_key: &[u8; 32],
        new_key: &[u8; 32],
    ) -> Result<EncryptedData> {
        let data = Self::decrypt(data, ad, old_key)?;
        Ok(Self::encrypt(data, ad, new_key))
    }
    fn encrypt(data: DecryptedData, ad: AdditionalData<'_>, key: &[u8; 32]) -> EncryptedData;
    fn decrypt(
        data: EncryptedData,
        ad: AdditionalData<'_>,
        key: &[u8; 32],
    ) -> Result<DecryptedData>;
}

impl Record<DecryptedData> {
    pub(crate) fn encrypt<E: Encryption>(self, key: &[u8; 32]) -> Record<EncryptedData> {
        let ad = AdditionalData {
            id: &self.id,
            version: &self.version,
            tag: &self.tag,
            host: &self.host.id,
            idx: &self.idx,
        };
        Record {
            data: E::encrypt(self.data, ad, key),
            id: self.id,
            host: self.host,
            idx: self.idx,
            timestamp: self.timestamp,
            version: self.version,
            tag: self.tag,
        }
    }
}

impl Record<EncryptedData> {
    pub(crate) fn decrypt<E: Encryption>(self, key: &[u8; 32]) -> Result<Record<DecryptedData>> {
        let ad = AdditionalData {
            id: &self.id,
            version: &self.version,
            tag: &self.tag,
            host: &self.host.id,
            idx: &self.idx,
        };
        Ok(Record {
            data: E::decrypt(self.data, ad, key)?,
            id: self.id,
            host: self.host,
            idx: self.idx,
            timestamp: self.timestamp,
            version: self.version,
            tag: self.tag,
        })
    }

    pub(crate) fn re_encrypt<E: Encryption>(
        self,
        old_key: &[u8; 32],
        new_key: &[u8; 32],
    ) -> Result<Self> {
        let ad = AdditionalData {
            id: &self.id,
            version: &self.version,
            tag: &self.tag,
            host: &self.host.id,
            idx: &self.idx,
        };
        Ok(Self {
            data: E::re_encrypt(self.data, ad, old_key, new_key)?,
            id: self.id,
            host: self.host,
            idx: self.idx,
            timestamp: self.timestamp,
            version: self.version,
            tag: self.tag,
        })
    }
}

#[cfg(test)]
mod tests {
    use crate::atuin_common::record::{Host, HostId};

    use super::{DecryptedData, Record};

    fn test_record() -> Record<DecryptedData> {
        Record::builder()
            .host(Host::new(HostId(crate::atuin_common::utils::uuid_v7())))
            .version("v1".into())
            .tag(crate::atuin_common::utils::uuid_v7().simple().to_string())
            .data(DecryptedData(vec![0, 1, 2, 3]))
            .idx(0)
            .build()
    }
}