Skip to main content

mkit_core/store/
memory.rs

1use super::{MAX_RAW_OBJECT_SIZE, ObjectSource, StoreError, StoreResult, check_hash};
2use crate::hash::{Hash, to_hex};
3use crate::object::object_id_from_bytes;
4use std::collections::BTreeMap;
5/// Prefetched canonical bytes for synchronous builders such as
6/// [`crate::verify::build_disclosure_from`] or [`crate::ops::diff::diff_trees`].
7/// The host owns authorization and aggregate memory limits.
8#[derive(Debug, Default)]
9pub struct MemorySource {
10    objects: BTreeMap<Hash, Vec<u8>>,
11}
12impl MemorySource {
13    /// Insert verified canonical bytes; failed inserts preserve the previous value.
14    ///
15    /// # Errors
16    /// [`StoreError::ObjectTooLarge`] or [`StoreError::HashMismatch`].
17    pub fn insert(&mut self, id: Hash, bytes: Vec<u8>) -> StoreResult<()> {
18        if bytes.len() > MAX_RAW_OBJECT_SIZE {
19            return Err(StoreError::ObjectTooLarge);
20        }
21        check_hash(&id, &object_id_from_bytes(&bytes))?;
22        self.objects.insert(id, bytes);
23        Ok(())
24    }
25    fn bytes(&self, id: &Hash) -> StoreResult<&[u8]> {
26        self.objects
27            .get(id)
28            .map(Vec::as_slice)
29            .ok_or_else(|| StoreError::ObjectNotFound(to_hex(id)))
30    }
31}
32impl ObjectSource for MemorySource {
33    fn read(&self, id: &Hash) -> StoreResult<Vec<u8>> {
34        let bytes = self.bytes(id)?;
35        check_hash(id, &object_id_from_bytes(bytes))?;
36        Ok(bytes.to_vec())
37    }
38}
39#[cfg(test)]
40mod tests {
41    use super::*;
42    use crate::object::{Blob, ChunkedBlob, EntryMode, Tree, TreeEntry};
43    use crate::{object::Object, serialize};
44
45    fn put(source: &mut MemorySource, object: &Object) -> Hash {
46        let id = object.id().unwrap();
47        source
48            .insert(id, serialize::serialize(object).unwrap())
49            .unwrap();
50        id
51    }
52
53    #[test]
54    fn prefetched_blob_tree_and_manifest_use_their_canonical_ids() {
55        let mut source = MemorySource::default();
56        let blob = Object::Blob(Blob {
57            data: b"chunk".to_vec(),
58        });
59        let chunk = put(&mut source, &blob);
60        let manifest = Object::ChunkedBlob(ChunkedBlob {
61            total_size: 5,
62            chunk_size: 0,
63            chunks: vec![chunk],
64        });
65        let file = put(&mut source, &manifest);
66        let tree = Object::Tree(Tree {
67            entries: vec![TreeEntry {
68                name: b"file".to_vec(),
69                mode: EntryMode::Blob,
70                object_hash: file,
71            }],
72        });
73        let root = put(&mut source, &tree);
74        let source: &dyn ObjectSource = &source;
75        for (id, object) in [(chunk, blob), (file, manifest), (root, tree)] {
76            assert_eq!(
77                source.read(&id).unwrap(),
78                serialize::serialize(&object).unwrap()
79            );
80            assert_eq!(source.read_object(&id).unwrap(), object);
81        }
82        assert!(matches!(
83            source.read(&[0; 32]),
84            Err(StoreError::ObjectNotFound(_))
85        ));
86    }
87
88    #[test]
89    fn failed_insert_preserves_previous_bytes_and_reads_verify_again() {
90        let mut source = MemorySource::default();
91        let object = Object::Blob(Blob {
92            data: b"valid".to_vec(),
93        });
94        let id = put(&mut source, &object);
95        assert!(matches!(
96            source.insert(id, b"invalid".to_vec()),
97            Err(StoreError::HashMismatch { .. })
98        ));
99        assert_eq!(source.read_object(&id).unwrap(), object);
100        // Simulate memory corruption to prove insertion-time checking is not
101        // the only integrity boundary, including through read_unverified.
102        source.objects.get_mut(&id).unwrap()[0] ^= 0xff;
103        for result in [source.read(&id), source.read_unverified(&id)] {
104            assert!(matches!(result, Err(StoreError::HashMismatch { .. })));
105        }
106        assert!(matches!(
107            source.read_object(&id),
108            Err(StoreError::HashMismatch { .. })
109        ));
110    }
111
112    #[test]
113    fn prefetched_manifest_builds_the_same_disclosure_as_the_durable_source() {
114        use crate::object::{Commit, Identity};
115        use crate::sign::{KeyPair, sign_commit};
116        use crate::verify::{Selector, build_disclosure, build_disclosure_from, verify_disclosure};
117
118        let dir = tempfile::TempDir::new().unwrap();
119        let store = super::super::ObjectStore::init(&crate::layout::RepoLayout::single(dir.path()))
120            .unwrap();
121        let mut source = MemorySource::default();
122        let mut add = |object: Object| {
123            let bytes = serialize::serialize(&object).unwrap();
124            let id = store.write(&bytes).unwrap();
125            source.insert(id, bytes.clone()).unwrap();
126            id
127        };
128        let chunk = add(Object::Blob(Blob {
129            data: b"chunk".to_vec(),
130        }));
131        let manifest = add(Object::ChunkedBlob(ChunkedBlob {
132            total_size: 5,
133            chunk_size: 0,
134            chunks: vec![chunk],
135        }));
136        let tree = add(Object::Tree(Tree {
137            entries: vec![TreeEntry {
138                name: b"file".to_vec(),
139                mode: EntryMode::Blob,
140                object_hash: manifest,
141            }],
142        }));
143        let key = KeyPair::from_seed([7; 32]);
144        let mut commit = Commit {
145            tree_hash: tree,
146            parents: vec![],
147            author: Identity::ed25519(key.public.0),
148            signer: key.public.0,
149            message: b"prefetch fixture".to_vec(),
150            timestamp: 1_726_300_000,
151            message_hash: [0; 32],
152            content_digest: [0; 32],
153            signature: [0; 64],
154        };
155        commit.signature = sign_commit(&commit, &key).unwrap().0;
156        let id = add(Object::Commit(commit));
157        for selector in [Selector::Object, Selector::Chunk(0)] {
158            let bundle = build_disclosure_from(&source, &id, &[b"file"], selector).unwrap();
159            assert_eq!(
160                bundle,
161                build_disclosure(&store, &id, &[b"file"], selector).unwrap()
162            );
163            let verified = verify_disclosure(&id, &bundle).unwrap();
164            assert!(verified.signature_valid);
165            assert_eq!(verified.leaf_id, manifest);
166        }
167        // A Merkle object's canonical bytes must be checked on read too.
168        source
169            .objects
170            .get_mut(&tree)
171            .unwrap()
172            .last_mut()
173            .map(|byte| *byte ^= 1)
174            .unwrap();
175        assert!(matches!(
176            source.read(&tree),
177            Err(StoreError::HashMismatch { .. })
178        ));
179        assert!(matches!(
180            source.read_object(&tree),
181            Err(StoreError::HashMismatch { .. })
182        ));
183    }
184}