refactor: commonize MessagePack (rmp) helpers

crates/atuin-client/Cargo.toml

@@ -47,7 +47,6 @@ serde_regex = "1.1.0"
 fs-err = { workspace = true }
 sql-builder = { workspace = true }
 memchr = "2.7"
-rmp = { version = "0.8.14" }
 typed-builder = { workspace = true }
 tokio = { workspace = true }
 semver = { workspace = true }
@@ -75,11 +74,12 @@ crossterm = { workspace = true, features = ["serde"] }
 palette = { version = "0.7.5", features = ["serializing"] }
 strum_macros = "0.27"
 strum = { version = "0.27", features = ["strum_macros"] }
+rmp = { version = "0.8.14" }
 
 [dependencies.atuin-common]
 path = "../atuin-common"
 version = "18.17.1"
-features = ["tracing"]
+features = ["tracing", "rmp", "eyre"]
 
 [dev-dependencies]
 tokio = { version = "1", features = ["full"] }

crates/atuin-client/src/encryption.rs

@@ -13,11 +13,12 @@ use std::io::prelude::*;
 use base64::prelude::{BASE64_STANDARD, Engine};
 pub use crypto_secretbox::Key;
 use crypto_secretbox::{KeyInit, XSalsa20Poly1305, aead::OsRng};
-use eyre::{Context, Result, bail, ensure, eyre};
+use eyre::{Context, Result, bail, ensure};
 use fs_err as fs;
-use rmp::Marker;
 
 use crate::settings::Settings;
+use atuin_common::rmp as atu_rmp;
+use atuin_common::rmp::decode::DecodeExt;
 
 pub fn generate_encoded_key() -> Result<(Key, String)> {
     let key = XSalsa20Poly1305::generate_key(&mut OsRng);
@@ -57,10 +58,10 @@ pub fn load_key(settings: &Settings) -> Result<Key> {
 
 pub fn encode_key(key: &Key) -> Result<String> {
     let mut buf = vec![];
-    rmp::encode::write_array_len(&mut buf, key.len() as u32)
+    atu_rmp::encode::write_array_len(&mut buf, key.len() as u32)
         .wrap_err("could not encode key to message pack")?;
     for b in key {
-        rmp::encode::write_uint(&mut buf, *b as u64)
+        atu_rmp::encode::write_uint(&mut buf, *b as u64)
             .wrap_err("could not encode key to message pack")?;
     }
     let buf = BASE64_STANDARD.encode(buf);
@@ -69,8 +70,6 @@ pub fn encode_key(key: &Key) -> Result<String> {
 }
 
 pub fn decode_key(key: String) -> Result<Key> {
-    use rmp::decode;
-
     let buf = BASE64_STANDARD
         .decode(key.trim_end())
         .wrap_err("encryption key is not a valid base64 encoding")?;
@@ -80,23 +79,22 @@ pub fn decode_key(key: String) -> Result<Key> {
     match <[u8; 32]>::try_from(&*buf) {
         Ok(key) => Ok(key.into()),
         Err(_) => {
-            let mut bytes = rmp::decode::Bytes::new(&buf);
+            let mut bytes = atu_rmp::decode::Bytes::new(&buf);
 
-            match Marker::from_u8(buf[0]) {
-                Marker::Bin8 => {
-                    let len = decode::read_bin_len(&mut bytes).map_err(|err| eyre!("{err:?}"))?;
+            match atu_rmp::decode::Marker::from_u8(buf[0]) {
+                atu_rmp::decode::Marker::Bin8 => {
+                    let len = rmp::decode::read_bin_len(&mut bytes).decode()?;
                     ensure!(len == 32, "encryption key is not the correct size");
                     let key = <[u8; 32]>::try_from(bytes.remaining_slice())
                         .context("could not decode encryption key")?;
                     Ok(key.into())
                 }
-                Marker::Array16 => {
-                    let len = decode::read_array_len(&mut bytes).map_err(|err| eyre!("{err:?}"))?;
-                    ensure!(len == 32, "encryption key is not the correct size");
+                atu_rmp::decode::Marker::Array16 => {
+                    atu_rmp::decode::expect_array_len(&mut bytes, 32)?;
 
                     let mut key = Key::default();
                     for i in &mut key {
-                        *i = rmp::decode::read_int(&mut bytes).map_err(|err| eyre!("{err:?}"))?;
+                        *i = rmp::decode::read_int::<u8, _>(&mut bytes).decode()?;
                     }
                     Ok(key)
                 }

crates/atuin-client/src/history.rs

@@ -1,5 +1,3 @@
-use rmp::decode::{self, Bytes};
-use rmp::encode;
 use std::env;
 
 use atuin_common::record::DecryptedData;
@@ -10,7 +8,8 @@ use eyre::{Result, bail};
 use crate::secrets::SECRET_PATTERNS_RE;
 use crate::settings::Settings;
 use crate::utils::get_host_user;
-use crate::utils::rmp::{DecodeError, EncodeError, read_optional, read_string, write_optional};
+use atuin_common::rmp as atu_rmp;
+use atuin_common::rmp::decode::DecodeExt;
 use time::OffsetDateTime;
 
 pub(crate) mod builder;
@@ -214,30 +213,38 @@ impl History {
     /// V2 is designed to allow new fields to be added without incrementing the version. V1 cannot
     /// accommodate this because its deserialization routine errors if more than 11 fields are
     /// provided.
-    pub fn serialize(&self) -> Result<DecryptedData, EncodeError> {
+    pub fn serialize(&self) -> Result<DecryptedData, atu_rmp::encode::EncodeError> {
         let mut output = vec![];
 
         // write the version
-        encode::write_u16(&mut output, Version::LATEST.as_int())?;
-        encode::write_array_len(&mut output, Version::LATEST.min_fields())?;
-
-        encode::write_str(&mut output, &self.id.0)?;
-        encode::write_u64(&mut output, self.timestamp.unix_timestamp_nanos() as u64)?;
-        encode::write_sint(&mut output, self.duration)?;
-        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)?;
-
-        write_optional(
+        atu_rmp::encode::write_u16(&mut output, Version::LATEST.as_int())?;
+        atu_rmp::encode::write_array_len(&mut output, Version::LATEST.min_fields())?;
+
+        atu_rmp::encode::write_str(&mut output, &self.id.0)?;
+        atu_rmp::encode::write_u64(&mut output, self.timestamp.unix_timestamp_nanos() as u64)?;
+        atu_rmp::encode::write_sint(&mut output, self.duration)?;
+        atu_rmp::encode::write_sint(&mut output, self.exit)?;
+        atu_rmp::encode::write_str(&mut output, &self.command)?;
+        atu_rmp::encode::write_str(&mut output, &self.cwd)?;
+        atu_rmp::encode::write_str(&mut output, &self.session)?;
+        atu_rmp::encode::write_str(&mut output, &self.hostname)?;
+
+        atu_rmp::encode::write_optional(
             &mut output,
             self.deleted_at.map(|d| d.unix_timestamp_nanos() as u64),
-            encode::write_u64,
+            atu_rmp::encode::write_u64,
+        )?;
+        atu_rmp::encode::write_str(&mut output, self.author.as_str())?;
+        atu_rmp::encode::write_optional(
+            &mut output,
+            self.intent.as_deref(),
+            atu_rmp::encode::write_str,
+        )?;
+        atu_rmp::encode::write_optional(
+            &mut output,
+            self.shell.as_deref(),
+            atu_rmp::encode::write_str,
         )?;
-        encode::write_str(&mut output, self.author.as_str())?;
-        write_optional(&mut output, self.intent.as_deref(), encode::write_str)?;
-        write_optional(&mut output, self.shell.as_deref(), encode::write_str)?;
         Ok(DecryptedData(output))
     }
 
@@ -246,32 +253,33 @@ impl History {
             bail!("unknown version {version:?}");
         };
 
-        let mut bytes = Bytes::new(bytes);
+        let mut bytes = atu_rmp::decode::Bytes::new(bytes);
 
-        let real_version = decode::read_u16(&mut bytes).map_err(DecodeError::from)?;
+        let real_version = rmp::decode::read_int::<u16, _>(&mut bytes).decode()?;
         if real_version != version.as_int() {
             bail!("expected to decode {version} record, found v{real_version}");
         }
 
-        let nfields = decode::read_array_len(&mut bytes).map_err(DecodeError::from)?;
+        let nfields = rmp::decode::read_array_len(&mut bytes).decode()?;
         let min_fields = version.min_fields();
         if nfields < min_fields || version.max_fields().is_some_and(|max| nfields > max) {
             bail!("unexpected number of fields ({nfields}) for history version {version}");
         }
 
-        let id = read_string(&mut bytes)?;
-        let timestamp = decode::read_u64(&mut bytes).map_err(DecodeError::from)?;
-        let duration = decode::read_int(&mut bytes).map_err(DecodeError::from)?;
-        let exit = decode::read_int(&mut bytes).map_err(DecodeError::from)?;
+        let id = atu_rmp::decode::read_string(&mut bytes).decode()?;
+        let timestamp = rmp::decode::read_int::<u64, _>(&mut bytes).decode()?;
+        let duration = rmp::decode::read_int::<i64, _>(&mut bytes).decode()?;
+        let exit = rmp::decode::read_int::<i64, _>(&mut bytes).decode()?;
 
-        let command = read_string(&mut bytes)?;
-        let cwd = read_string(&mut bytes)?;
-        let session = read_string(&mut bytes)?;
-        let hostname = read_string(&mut bytes)?;
-        let deleted_at = read_optional(&mut bytes, decode::read_u64)?;
+        let command = atu_rmp::decode::read_string(&mut bytes).decode()?;
+        let cwd = atu_rmp::decode::read_string(&mut bytes).decode()?;
+        let session = atu_rmp::decode::read_string(&mut bytes).decode()?;
+        let hostname = atu_rmp::decode::read_string(&mut bytes).decode()?;
+        let deleted_at =
+            atu_rmp::decode::read_optional(&mut bytes, rmp::decode::read_int::<u64, _>)?;
 
         let author = if version >= Version::One {
-            read_optional(&mut bytes, read_string)?
+            atu_rmp::decode::read_optional(&mut bytes, atu_rmp::decode::read_string)?
         } else {
             None
         };
@@ -281,13 +289,13 @@ impl History {
             Version::One => nfields > min_fields,
             Version::Two => true,
         } {
-            read_optional(&mut bytes, read_string)?
+            atu_rmp::decode::read_optional(&mut bytes, atu_rmp::decode::read_string)?
         } else {
             None
         };
 
         let shell = if version >= Version::Two {
-            read_optional(&mut bytes, read_string)?
+            atu_rmp::decode::read_optional(&mut bytes, atu_rmp::decode::read_string)?
         } else {
             None
         };

crates/atuin-client/src/history/store.rs

@@ -3,7 +3,6 @@ use std::{collections::HashSet, fmt::Write, time::Duration};
 use eyre::{Result, bail, eyre};
 use futures::{Stream, StreamExt, TryStreamExt, future, stream};
 use indicatif::{ProgressBar, ProgressState, ProgressStyle};
-use rmp::decode::Bytes;
 
 use crate::{
     database::{Database, current_context},
@@ -12,6 +11,8 @@ use crate::{
 use atuin_common::record::{DecryptedData, Host, HostId, Record, RecordId, RecordIdx};
 
 use super::{HISTORY_TAG, History, HistoryId, Version};
+use atuin_common::rmp as atu_rmp;
+use atuin_common::rmp::decode::DecodeExt;
 
 #[derive(Debug, Clone)]
 pub struct HistoryStore {
@@ -47,46 +48,39 @@ impl HistoryRecord {
     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)?;
+                atu_rmp::encode::write_u8(&mut output, 0)?;
 
                 let bytes = history.serialize()?;
 
-                encode::write_bin(&mut output, &bytes.0)?;
+                atu_rmp::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())?;
+                atu_rmp::encode::write_u8(&mut output, 1)?;
+                atu_rmp::encode::write_str(&mut output, id.0.as_str())?;
             }
         };
 
         Ok(DecryptedData(output))
     }
 
     pub fn deserialize(bytes: &DecryptedData, version: &str) -> Result<Self> {
-        use rmp::decode;
+        let mut bytes = atu_rmp::decode::Bytes::new(&bytes.0);
 
-        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)?;
+        let record_type = rmp::decode::read_int::<u8, _>(&mut bytes).decode()?;
 
         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 _ = rmp::decode::read_bin_len(&mut bytes).decode()?;
 
                 let record = History::deserialize(bytes.remaining_slice(), version)?;
 
@@ -95,16 +89,10 @@ impl HistoryRecord {
 
             // 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:?}"
-                    );
-                }
+                let id = atu_rmp::decode::read_string(&mut bytes).decode()?;
+                atu_rmp::decode::expect_eof(&bytes)?;
 
-                Ok(HistoryRecord::Delete(id.to_string().into()))
+                Ok(HistoryRecord::Delete(id.into()))
             }
 
             n => {

crates/atuin-client/src/utils.rs

@@ -1,5 +1,3 @@
-pub(crate) mod rmp;
-
 pub(crate) fn get_hostname() -> String {
     std::env::var("ATUIN_HOST_NAME")
         .unwrap_or_else(|_| whoami::hostname().unwrap_or_else(|_| "unknown-host".to_string()))