atuin_common/
record.rs

1use std::collections::HashMap;
2
3use eyre::Result;
4use serde::{Deserialize, Serialize};
5use typed_builder::TypedBuilder;
6use uuid::Uuid;
7
8#[derive(Clone, Debug, PartialEq)]
9pub struct DecryptedData(pub Vec<u8>);
10
11#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
12pub struct EncryptedData {
13    pub data: String,
14    pub content_encryption_key: String,
15}
16
17#[derive(Debug, PartialEq, PartialOrd, Ord, Eq)]
18pub struct Diff {
19    pub host: HostId,
20    pub tag: String,
21    pub local: Option<RecordIdx>,
22    pub remote: Option<RecordIdx>,
23}
24
25#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
26pub struct Host {
27    pub id: HostId,
28    pub name: String,
29}
30
31impl Host {
32    pub fn new(id: HostId) -> Self {
33        Host {
34            id,
35            name: String::new(),
36        }
37    }
38}
39
40new_uuid!(RecordId);
41new_uuid!(HostId);
42
43pub type RecordIdx = u64;
44
45/// A single record stored inside of our local database
46#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TypedBuilder)]
47pub struct Record<Data> {
48    /// a unique ID
49    #[builder(default = RecordId(crate::utils::uuid_v7()))]
50    pub id: RecordId,
51
52    /// The integer record ID. This is only unique per (host, tag).
53    pub idx: RecordIdx,
54
55    /// The unique ID of the host.
56    // TODO(ellie): Optimize the storage here. We use a bunch of IDs, and currently store
57    // as strings. I would rather avoid normalization, so store as UUID binary instead of
58    // encoding to a string and wasting much more storage.
59    pub host: Host,
60
61    /// The creation time in nanoseconds since unix epoch
62    #[builder(default = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as u64)]
63    pub timestamp: u64,
64
65    /// The version the data in the entry conforms to
66    // However we want to track versions for this tag, eg v2
67    pub version: String,
68
69    /// The type of data we are storing here. Eg, "history"
70    pub tag: String,
71
72    /// Some data. This can be anything you wish to store. Use the tag field to know how to handle it.
73    pub data: Data,
74}
75
76/// Extra data from the record that should be encoded in the data
77#[derive(Debug, Copy, Clone)]
78pub struct AdditionalData<'a> {
79    pub id: &'a RecordId,
80    pub idx: &'a u64,
81    pub version: &'a str,
82    pub tag: &'a str,
83    pub host: &'a HostId,
84}
85
86impl<Data> Record<Data> {
87    pub fn append(&self, data: Vec<u8>) -> Record<DecryptedData> {
88        Record::builder()
89            .host(self.host.clone())
90            .version(self.version.clone())
91            .idx(self.idx + 1)
92            .tag(self.tag.clone())
93            .data(DecryptedData(data))
94            .build()
95    }
96}
97
98/// An index representing the current state of the record stores
99/// This can be both remote, or local, and compared in either direction
100#[derive(Debug, Serialize, Deserialize)]
101pub struct RecordStatus {
102    // A map of host -> tag -> max(idx)
103    pub hosts: HashMap<HostId, HashMap<String, RecordIdx>>,
104}
105
106impl Default for RecordStatus {
107    fn default() -> Self {
108        Self::new()
109    }
110}
111
112impl Extend<(HostId, String, RecordIdx)> for RecordStatus {
113    fn extend<T: IntoIterator<Item = (HostId, String, RecordIdx)>>(&mut self, iter: T) {
114        for (host, tag, tail_idx) in iter {
115            self.set_raw(host, tag, tail_idx);
116        }
117    }
118}
119
120impl RecordStatus {
121    pub fn new() -> RecordStatus {
122        RecordStatus {
123            hosts: HashMap::new(),
124        }
125    }
126
127    /// Insert a new tail record into the store
128    pub fn set(&mut self, tail: Record<DecryptedData>) {
129        self.set_raw(tail.host.id, tail.tag, tail.idx)
130    }
131
132    pub fn set_raw(&mut self, host: HostId, tag: String, tail_id: RecordIdx) {
133        self.hosts.entry(host).or_default().insert(tag, tail_id);
134    }
135
136    pub fn get(&self, host: HostId, tag: String) -> Option<RecordIdx> {
137        self.hosts.get(&host).and_then(|v| v.get(&tag)).cloned()
138    }
139
140    /// Diff this index with another, likely remote index.
141    /// The two diffs can then be reconciled, and the optimal change set calculated
142    /// Returns a tuple, with (host, tag, Option(OTHER))
143    /// OTHER is set to the value of the idx on the other machine. If it is greater than our index,
144    /// then we need to do some downloading. If it is smaller, then we need to do some uploading
145    /// Note that we cannot upload if we are not the owner of the record store - hosts can only
146    /// write to their own store.
147    pub fn diff(&self, other: &Self) -> Vec<Diff> {
148        let mut ret = Vec::new();
149
150        // First, we check if other has everything that self has
151        for (host, tag_map) in self.hosts.iter() {
152            for (tag, idx) in tag_map.iter() {
153                match other.get(*host, tag.clone()) {
154                    // The other store is all up to date! No diff.
155                    Some(t) if t.eq(idx) => continue,
156
157                    // The other store does exist, and it is either ahead or behind us. A diff regardless
158                    Some(t) => ret.push(Diff {
159                        host: *host,
160                        tag: tag.clone(),
161                        local: Some(*idx),
162                        remote: Some(t),
163                    }),
164
165                    // The other store does not exist :O
166                    None => ret.push(Diff {
167                        host: *host,
168                        tag: tag.clone(),
169                        local: Some(*idx),
170                        remote: None,
171                    }),
172                };
173            }
174        }
175
176        // At this point, there is a single case we have not yet considered.
177        // If the other store knows of a tag that we are not yet aware of, then the diff will be missed
178
179        // account for that!
180        for (host, tag_map) in other.hosts.iter() {
181            for (tag, idx) in tag_map.iter() {
182                match self.get(*host, tag.clone()) {
183                    // If we have this host/tag combo, the comparison and diff will have already happened above
184                    Some(_) => continue,
185
186                    None => ret.push(Diff {
187                        host: *host,
188                        tag: tag.clone(),
189                        remote: Some(*idx),
190                        local: None,
191                    }),
192                };
193            }
194        }
195
196        // Stability is a nice property to have
197        ret.sort();
198        ret
199    }
200}
201
202pub trait Encryption {
203    fn re_encrypt(
204        data: EncryptedData,
205        ad: AdditionalData,
206        old_key: &[u8; 32],
207        new_key: &[u8; 32],
208    ) -> Result<EncryptedData> {
209        let data = Self::decrypt(data, ad, old_key)?;
210        Ok(Self::encrypt(data, ad, new_key))
211    }
212    fn encrypt(data: DecryptedData, ad: AdditionalData, key: &[u8; 32]) -> EncryptedData;
213    fn decrypt(data: EncryptedData, ad: AdditionalData, key: &[u8; 32]) -> Result<DecryptedData>;
214}
215
216impl Record<DecryptedData> {
217    pub fn encrypt<E: Encryption>(self, key: &[u8; 32]) -> Record<EncryptedData> {
218        let ad = AdditionalData {
219            id: &self.id,
220            version: &self.version,
221            tag: &self.tag,
222            host: &self.host.id,
223            idx: &self.idx,
224        };
225        Record {
226            data: E::encrypt(self.data, ad, key),
227            id: self.id,
228            host: self.host,
229            idx: self.idx,
230            timestamp: self.timestamp,
231            version: self.version,
232            tag: self.tag,
233        }
234    }
235}
236
237impl Record<EncryptedData> {
238    pub fn decrypt<E: Encryption>(self, key: &[u8; 32]) -> Result<Record<DecryptedData>> {
239        let ad = AdditionalData {
240            id: &self.id,
241            version: &self.version,
242            tag: &self.tag,
243            host: &self.host.id,
244            idx: &self.idx,
245        };
246        Ok(Record {
247            data: E::decrypt(self.data, ad, key)?,
248            id: self.id,
249            host: self.host,
250            idx: self.idx,
251            timestamp: self.timestamp,
252            version: self.version,
253            tag: self.tag,
254        })
255    }
256
257    pub fn re_encrypt<E: Encryption>(
258        self,
259        old_key: &[u8; 32],
260        new_key: &[u8; 32],
261    ) -> Result<Record<EncryptedData>> {
262        let ad = AdditionalData {
263            id: &self.id,
264            version: &self.version,
265            tag: &self.tag,
266            host: &self.host.id,
267            idx: &self.idx,
268        };
269        Ok(Record {
270            data: E::re_encrypt(self.data, ad, old_key, new_key)?,
271            id: self.id,
272            host: self.host,
273            idx: self.idx,
274            timestamp: self.timestamp,
275            version: self.version,
276            tag: self.tag,
277        })
278    }
279}
280
281#[cfg(test)]
282mod tests {
283    use crate::record::{Host, HostId};
284
285    use super::{DecryptedData, Diff, Record, RecordStatus};
286    use pretty_assertions::assert_eq;
287
288    fn test_record() -> Record<DecryptedData> {
289        Record::builder()
290            .host(Host::new(HostId(crate::utils::uuid_v7())))
291            .version("v1".into())
292            .tag(crate::utils::uuid_v7().simple().to_string())
293            .data(DecryptedData(vec![0, 1, 2, 3]))
294            .idx(0)
295            .build()
296    }
297
298    #[test]
299    fn record_index() {
300        let mut index = RecordStatus::new();
301        let record = test_record();
302
303        index.set(record.clone());
304
305        let tail = index.get(record.host.id, record.tag);
306
307        assert_eq!(
308            record.idx,
309            tail.expect("tail not in store"),
310            "tail in store did not match"
311        );
312    }
313
314    #[test]
315    fn record_index_overwrite() {
316        let mut index = RecordStatus::new();
317        let record = test_record();
318        let child = record.append(vec![1, 2, 3]);
319
320        index.set(record.clone());
321        index.set(child.clone());
322
323        let tail = index.get(record.host.id, record.tag);
324
325        assert_eq!(
326            child.idx,
327            tail.expect("tail not in store"),
328            "tail in store did not match"
329        );
330    }
331
332    #[test]
333    fn record_index_no_diff() {
334        // Here, they both have the same version and should have no diff
335
336        let mut index1 = RecordStatus::new();
337        let mut index2 = RecordStatus::new();
338
339        let record1 = test_record();
340
341        index1.set(record1.clone());
342        index2.set(record1);
343
344        let diff = index1.diff(&index2);
345
346        assert_eq!(0, diff.len(), "expected empty diff");
347    }
348
349    #[test]
350    fn record_index_single_diff() {
351        // Here, they both have the same stores, but one is ahead by a single record
352
353        let mut index1 = RecordStatus::new();
354        let mut index2 = RecordStatus::new();
355
356        let record1 = test_record();
357        let record2 = record1.append(vec![1, 2, 3]);
358
359        index1.set(record1);
360        index2.set(record2.clone());
361
362        let diff = index1.diff(&index2);
363
364        assert_eq!(1, diff.len(), "expected single diff");
365        assert_eq!(
366            diff[0],
367            Diff {
368                host: record2.host.id,
369                tag: record2.tag,
370                remote: Some(1),
371                local: Some(0)
372            }
373        );
374    }
375
376    #[test]
377    fn record_index_multi_diff() {
378        // A much more complex case, with a bunch more checks
379        let mut index1 = RecordStatus::new();
380        let mut index2 = RecordStatus::new();
381
382        let store1record1 = test_record();
383        let store1record2 = store1record1.append(vec![1, 2, 3]);
384
385        let store2record1 = test_record();
386        let store2record2 = store2record1.append(vec![1, 2, 3]);
387
388        let store3record1 = test_record();
389
390        let store4record1 = test_record();
391
392        // index1 only knows about the first two entries of the first two stores
393        index1.set(store1record1);
394        index1.set(store2record1);
395
396        // index2 is fully up to date with the first two stores, and knows of a third
397        index2.set(store1record2);
398        index2.set(store2record2);
399        index2.set(store3record1);
400
401        // index1 knows of a 4th store
402        index1.set(store4record1);
403
404        let diff1 = index1.diff(&index2);
405        let diff2 = index2.diff(&index1);
406
407        // both diffs the same length
408        assert_eq!(4, diff1.len());
409        assert_eq!(4, diff2.len());
410
411        dbg!(&diff1, &diff2);
412
413        // both diffs should be ALMOST the same. They will agree on which hosts and tags
414        // require updating, but the "other" value will not be the same.
415        let smol_diff_1: Vec<(HostId, String)> =
416            diff1.iter().map(|v| (v.host, v.tag.clone())).collect();
417        let smol_diff_2: Vec<(HostId, String)> =
418            diff1.iter().map(|v| (v.host, v.tag.clone())).collect();
419
420        assert_eq!(smol_diff_1, smol_diff_2);
421
422        // diffing with yourself = no diff
423        assert_eq!(index1.diff(&index1).len(), 0);
424        assert_eq!(index2.diff(&index2).len(), 0);
425    }
426}
atuin_common/record.rs

atuin_common/
record.rs
