path: root/crates/daemon/src/aclient/history/mod.rs

use std::time::Duration;

use rmp::decode::DecodeStringError;
use rmp::decode::ValueReadError;
use rmp::{Marker, decode::Bytes};
use turtle_api::history::History;

use turtle_common::record::DecryptedData;

use eyre::{Result, bail, eyre};

use time::OffsetDateTime;

pub(crate) mod store;

const HISTORY_VERSION_V0: &str = "v0";
const HISTORY_VERSION_V1: &str = "v1";
const HISTORY_RECORD_VERSION_V0: u16 = 0;
const HISTORY_RECORD_VERSION_V1: u16 = 1;
const HISTORY_VERSION: &str = HISTORY_VERSION_V1;
const HISTORY_TAG: &str = "history";

trait HistoryExt: Sized {
    fn serialize(&self) -> Result<DecryptedData>;
    fn read_optional_string(bytes: &[u8]) -> Result<(Option<String>, &[u8])>;
    fn deserialize_v0(bytes: &[u8]) -> Result<Self>;
    fn deserialize_v1(bytes: &[u8]) -> Result<Self>;
    fn deserialize(bytes: &[u8], version: &str) -> Result<Self>;
}

impl HistoryExt for History {
    fn serialize(&self) -> Result<DecryptedData> {
        // This is pretty much the same as what we used for the old history, with one difference -
        // it uses integers for timestamps rather than a string format.

        use rmp::encode;

        let mut output = vec![];

        // write the version
        encode::write_u16(&mut output, HISTORY_RECORD_VERSION_V1)?;
        let include_intent = self.intent.is_some();
        encode::write_array_len(&mut output, 10 + u32::from(include_intent))?;

        encode::write_str(&mut output, &self.id.to_string())?;
        encode::write_u64(&mut output, self.timestamp.unix_timestamp_nanos() as u64)?;
        encode::write_sint(
            &mut output,
            i64::try_from(self.duration.as_nanos()).expect("should be small enough"),
        )?;
        encode::write_sint(&mut output, self.exit)?;
        encode::write_str(&mut output, &self.command)?;
        encode::write_str(&mut output, &self.cwd)?;
        encode::write_str(&mut output, &self.session)?;
        encode::write_str(&mut output, &self.hostname)?;

        match self.deleted_at {
            Some(d) => encode::write_u64(&mut output, d.unix_timestamp_nanos() as u64)?,
            None => encode::write_nil(&mut output)?,
        }

        encode::write_str(&mut output, self.author.as_str())?;
        if let Some(intent) = &self.intent {
            encode::write_str(&mut output, intent.as_str())?;
        }

        Ok(DecryptedData(output))
    }

    fn read_optional_string(bytes: &[u8]) -> Result<(Option<String>, &[u8])> {
        use rmp::decode;

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

        match decode::read_str_from_slice(bytes) {
            Ok((value, bytes)) => Ok((Some(value.to_owned()), bytes)),
            Err(DecodeStringError::TypeMismatch(Marker::Null)) => {
                let mut cursor = Bytes::new(bytes);
                decode::read_nil(&mut cursor).map_err(error_report)?;

                Ok((None, cursor.remaining_slice()))
            }
            Err(err) => Err(error_report(err)),
        }
    }

    fn deserialize_v0(bytes: &[u8]) -> Result<Self> {
        use rmp::decode;

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

        let mut bytes = Bytes::new(bytes);

        let version = decode::read_u16(&mut bytes).map_err(error_report)?;

        if version != HISTORY_RECORD_VERSION_V0 {
            bail!("expected decoding v0 record, found v{version}");
        }

        let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?;

        if nfields != 9 {
            bail!("cannot decrypt history from a different version of Atuin");
        }

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

        let mut bytes = Bytes::new(bytes);
        let timestamp = decode::read_u64(&mut bytes).map_err(error_report)?;
        let duration = decode::read_int(&mut bytes)
            .map(|int: i64| {
                Duration::from_nanos(u64::try_from(int).expect("should be small enough"))
            })
            .map_err(error_report)?;
        let exit = decode::read_int(&mut bytes).map_err(error_report)?;

        let bytes = bytes.remaining_slice();
        let (command, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?;
        let (cwd, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?;
        let (session, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?;
        let (hostname, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?;

        let mut bytes = Bytes::new(bytes);

        let (deleted_at, bytes) = match decode::read_u64(&mut bytes) {
            Ok(unix) => (Some(unix), bytes.remaining_slice()),
            // we accept null here
            Err(ValueReadError::TypeMismatch(Marker::Null)) => (None, bytes.remaining_slice()),
            Err(err) => return Err(error_report(err)),
        };
        if !bytes.is_empty() {
            bail!("trailing bytes in encoded history. malformed")
        }

        Ok(Self {
            id: id.to_owned().into(),
            timestamp: OffsetDateTime::from_unix_timestamp_nanos(i128::from(timestamp))?,
            duration,
            exit,
            command: command.to_owned(),
            cwd: cwd.to_owned(),
            session: session.to_owned(),
            hostname: hostname.to_owned(),
            author: Self::author_from_hostname(hostname),
            intent: None,
            deleted_at: deleted_at
                .map(|t| OffsetDateTime::from_unix_timestamp_nanos(i128::from(t)))
                .transpose()?,
        })
    }

    fn deserialize_v1(bytes: &[u8]) -> Result<Self> {
        use rmp::decode;

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

        let mut bytes = Bytes::new(bytes);

        let version = decode::read_u16(&mut bytes).map_err(error_report)?;

        if version != HISTORY_RECORD_VERSION_V1 {
            bail!("expected decoding v1 record, found v{version}");
        }

        let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?;

        if !(10..=11).contains(&nfields) {
            bail!("cannot decrypt history from a different version of Atuin");
        }

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

        let mut bytes = Bytes::new(bytes);
        let timestamp = decode::read_u64(&mut bytes).map_err(error_report)?;
        let duration = decode::read_int(&mut bytes)
            .map(|int: i64| Duration::from_nanos(u64::try_from(int).expect("to be small enough")))
            .map_err(error_report)?;
        let exit = decode::read_int(&mut bytes).map_err(error_report)?;

        let bytes = bytes.remaining_slice();
        let (command, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?;
        let (cwd, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?;
        let (session, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?;
        let (hostname, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?;

        let mut bytes = Bytes::new(bytes);

        let (deleted_at, bytes) = match decode::read_u64(&mut bytes) {
            Ok(unix) => (Some(unix), bytes.remaining_slice()),
            // we accept null here
            Err(ValueReadError::TypeMismatch(Marker::Null)) => (None, bytes.remaining_slice()),
            Err(err) => return Err(error_report(err)),
        };
        let (author, bytes) = Self::read_optional_string(bytes)?;
        let (intent, bytes) = if nfields > 10 {
            Self::read_optional_string(bytes)?
        } else {
            (None, bytes)
        };

        if !bytes.is_empty() {
            bail!("trailing bytes in encoded history. malformed")
        }

        Ok(Self {
            id: id.to_owned().into(),
            timestamp: OffsetDateTime::from_unix_timestamp_nanos(i128::from(timestamp))?,
            duration,
            exit,
            command: command.to_owned(),
            cwd: cwd.to_owned(),
            session: session.to_owned(),
            hostname: hostname.to_owned(),
            author: author.unwrap_or_else(|| Self::author_from_hostname(hostname)),
            intent,
            deleted_at: deleted_at
                .map(|t| OffsetDateTime::from_unix_timestamp_nanos(i128::from(t)))
                .transpose()?,
        })
    }

    fn deserialize(bytes: &[u8], version: &str) -> Result<Self> {
        match version {
            HISTORY_VERSION_V0 => Self::deserialize_v0(bytes),
            HISTORY_VERSION_V1 => Self::deserialize_v1(bytes),

            _ => bail!("unknown version {version:?}"),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use time::macros::datetime;

    use crate::aclient::history::{HISTORY_VERSION, HistoryExt};

    use super::History;

    #[test]
    fn test_serialize_deserialize() {
        let history = History {
            id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(),
            timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00),
            duration: Duration::from_nanos(49_206_000),
            exit: 0,
            command: "git status".to_owned(),
            cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(),
            session: "b97d9a306f274473a203d2eba41f9457".to_owned(),
            hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(),
            author: "conrad.ludgate".to_owned(),
            intent: None,
            deleted_at: None,
        };

        let serialized = history.serialize().expect("failed to serialize history");
        assert_eq!(
            &serialized.0[0..3],
            [205, 0, 1],
            "should encode as history v1"
        );

        let deserialized = History::deserialize(&serialized.0, HISTORY_VERSION)
            .expect("failed to deserialize history");
        assert_eq!(history, deserialized);
    }

    #[test]
    fn test_serialize_deserialize_deleted() {
        let history = History {
            id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(),
            timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00),
            duration: Duration::from_nanos(49_206_000),
            exit: 0,
            command: "git status".to_owned(),
            cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(),
            session: "b97d9a306f274473a203d2eba41f9457".to_owned(),
            hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(),
            author: "conrad.ludgate".to_owned(),
            intent: None,
            deleted_at: Some(datetime!(2023-11-19 20:18 +00:00)),
        };

        let serialized = history.serialize().expect("failed to serialize history");

        let deserialized = History::deserialize(&serialized.0, HISTORY_VERSION)
            .expect("failed to deserialize history");

        assert_eq!(history, deserialized);
    }

    #[test]
    fn test_serialize_deserialize_with_author_and_intent() {
        let history = History {
            id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(),
            timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00),
            duration: Duration::from_nanos(49_206_000),
            exit: 0,
            command: "git status".to_owned(),
            cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(),
            session: "b97d9a306f274473a203d2eba41f9457".to_owned(),
            hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(),
            author: "claude".to_owned(),
            intent: Some("check repository status".to_owned()),
            deleted_at: None,
        };

        let serialized = history.serialize().expect("failed to serialize history");
        let deserialized = History::deserialize(&serialized.0, HISTORY_VERSION)
            .expect("failed to deserialize history");

        assert_eq!(history, deserialized);
    }

    #[test]
    fn test_serialize_deserialize_version() {
        // v0
        let bytes_v0 = [
            205, 0, 0, 153, 217, 32, 54, 54, 100, 49, 54, 99, 98, 101, 101, 55, 99, 100, 52, 55,
            53, 51, 56, 101, 53, 99, 53, 98, 56, 98, 52, 52, 101, 57, 48, 48, 54, 101, 207, 23, 99,
            98, 117, 24, 210, 246, 128, 206, 2, 238, 210, 240, 0, 170, 103, 105, 116, 32, 115, 116,
            97, 116, 117, 115, 217, 42, 47, 85, 115, 101, 114, 115, 47, 99, 111, 110, 114, 97, 100,
            46, 108, 117, 100, 103, 97, 116, 101, 47, 68, 111, 99, 117, 109, 101, 110, 116, 115,
            47, 99, 111, 100, 101, 47, 97, 116, 117, 105, 110, 217, 32, 98, 57, 55, 100, 57, 97,
            51, 48, 54, 102, 50, 55, 52, 52, 55, 51, 97, 50, 48, 51, 100, 50, 101, 98, 97, 52, 49,
            102, 57, 52, 53, 55, 187, 102, 118, 102, 103, 57, 51, 54, 99, 48, 107, 112, 102, 58,
            99, 111, 110, 114, 97, 100, 46, 108, 117, 100, 103, 97, 116, 101, 192,
        ];

        let deserialized = History::deserialize(&bytes_v0, "v0");
        assert!(deserialized.is_ok());

        let deserialized = History::deserialize(&bytes_v0, HISTORY_VERSION);
        assert!(deserialized.is_err());

        let current = History {
            id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(),
            timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00),
            duration: Duration::from_nanos(49_206_000),
            exit: 0,
            command: "git status".to_owned(),
            cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(),
            session: "b97d9a306f274473a203d2eba41f9457".to_owned(),
            hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(),
            author: "conrad.ludgate".to_owned(),
            intent: None,
            deleted_at: None,
        };

        let bytes_v1 = current.serialize().expect("failed to serialize history");
        let deserialized = History::deserialize(&bytes_v1.0, HISTORY_VERSION);
        assert!(deserialized.is_ok());

        let deserialized = History::deserialize(&bytes_v1.0, "v0");
        assert!(deserialized.is_err());
    }
}