Skip to main content

corium_store/
encrypted_store.rs

1use std::collections::BTreeMap;
2use std::time::SystemTime;
3
4use async_trait::async_trait;
5use corium_crypt::{SecretKey, decrypt_blob, encrypt_blob, parse_blob_header};
6
7use crate::{BlobId, BlobIdStream, BlobStore, RootStore, StoreError, digest};
8
9/// Blob-store decorator that exposes plaintext while storing deterministic
10/// authenticated ciphertext.
11///
12/// Blob identifiers are digests of the encrypted object. Listing, deletion,
13/// timestamps, and cleartext root records pass through to the wrapped store.
14/// The key set is an immutable snapshot; a manifest reload or rotation replaces
15/// the decorator rather than resolving a KMS key on each blob read.
16#[derive(Clone)]
17pub struct EncryptedBlobStore<S> {
18    inner: S,
19    current_epoch: u32,
20    keys: BTreeMap<u32, SecretKey>,
21}
22
23impl<S> EncryptedBlobStore<S> {
24    /// Creates a decorator with all readable epochs and the epoch used by new
25    /// writes.
26    ///
27    /// # Errors
28    ///
29    /// Returns [`StoreError::MissingEncryptionKey`] when `current_epoch` is not
30    /// present in `keys`.
31    pub fn new(
32        inner: S,
33        current_epoch: u32,
34        keys: impl IntoIterator<Item = (u32, SecretKey)>,
35    ) -> Result<Self, StoreError> {
36        let keys = keys.into_iter().collect::<BTreeMap<_, _>>();
37        if !keys.contains_key(&current_epoch) {
38            return Err(StoreError::MissingEncryptionKey(current_epoch));
39        }
40        Ok(Self {
41            inner,
42            current_epoch,
43            keys,
44        })
45    }
46
47    /// Creates a single-epoch decorator.
48    #[must_use]
49    pub fn with_key(inner: S, epoch: u32, key: SecretKey) -> Self {
50        Self {
51            inner,
52            current_epoch: epoch,
53            keys: BTreeMap::from([(epoch, key)]),
54        }
55    }
56
57    /// Returns the wrapped ciphertext store.
58    #[must_use]
59    pub fn inner(&self) -> &S {
60        &self.inner
61    }
62
63    /// Consumes the decorator and returns the wrapped store.
64    #[must_use]
65    pub fn into_inner(self) -> S {
66        self.inner
67    }
68
69    /// Returns the epoch used for new writes.
70    #[must_use]
71    pub fn current_epoch(&self) -> u32 {
72        self.current_epoch
73    }
74
75    fn current_key(&self) -> &SecretKey {
76        // Constructors establish this invariant, and there is no key mutation.
77        self.keys
78            .get(&self.current_epoch)
79            .expect("current encryption epoch must have key material")
80    }
81
82    fn key(&self, epoch: u32) -> Result<&SecretKey, StoreError> {
83        self.keys
84            .get(&epoch)
85            .ok_or(StoreError::MissingEncryptionKey(epoch))
86    }
87}
88
89#[async_trait]
90impl<S: BlobStore> BlobStore for EncryptedBlobStore<S> {
91    async fn put(&self, bytes: &[u8]) -> Result<BlobId, StoreError> {
92        let encrypted = encrypt_blob(self.current_key(), self.current_epoch, bytes)?;
93        let expected = digest(&encrypted);
94        let actual = self.inner.put(&encrypted).await?;
95        if actual != expected {
96            return Err(StoreError::BlobIdMismatch { expected, actual });
97        }
98        Ok(actual)
99    }
100
101    async fn get(&self, id: &BlobId) -> Result<Option<Vec<u8>>, StoreError> {
102        let Some(encrypted) = self.inner.get(id).await? else {
103            return Ok(None);
104        };
105        if digest(&encrypted) != *id {
106            return Err(StoreError::CorruptBlob(id.clone()));
107        }
108        let header = parse_blob_header(&encrypted)?;
109        let plaintext = decrypt_blob(self.key(header.epoch)?, &encrypted)?;
110        Ok(Some(plaintext))
111    }
112
113    async fn contains(&self, id: &BlobId) -> Result<bool, StoreError> {
114        self.inner.contains(id).await
115    }
116
117    async fn put_if_absent(&self, bytes: &[u8]) -> Result<BlobId, StoreError> {
118        let encrypted = encrypt_blob(self.current_key(), self.current_epoch, bytes)?;
119        let expected = digest(&encrypted);
120        if self.inner.contains(&expected).await? {
121            return Ok(expected);
122        }
123        // Another writer may store the same object after this check. That race
124        // is benign because deterministic encryption gives both writers the
125        // same bytes and content id, and BlobStore::put is idempotent.
126        let actual = self.inner.put(&encrypted).await?;
127        if actual != expected {
128            return Err(StoreError::BlobIdMismatch { expected, actual });
129        }
130        Ok(actual)
131    }
132
133    async fn delete(&self, id: &BlobId) -> Result<(), StoreError> {
134        self.inner.delete(id).await
135    }
136
137    async fn list(&self) -> Result<BlobIdStream, StoreError> {
138        self.inner.list().await
139    }
140
141    async fn modified_at(&self, id: &BlobId) -> Result<Option<SystemTime>, StoreError> {
142        self.inner.modified_at(id).await
143    }
144}
145
146#[async_trait]
147impl<S: RootStore> RootStore for EncryptedBlobStore<S> {
148    async fn get_root(&self, name: &str) -> Result<Option<Vec<u8>>, StoreError> {
149        self.inner.get_root(name).await
150    }
151
152    async fn cas_root(
153        &self,
154        name: &str,
155        expected: Option<&[u8]>,
156        new: &[u8],
157    ) -> Result<(), StoreError> {
158        self.inner.cas_root(name, expected, new).await
159    }
160
161    async fn delete_root(&self, name: &str) -> Result<(), StoreError> {
162        self.inner.delete_root(name).await
163    }
164
165    async fn list_roots(&self, prefix: &str) -> Result<Vec<String>, StoreError> {
166        self.inner.list_roots(prefix).await
167    }
168}