cipherstash_dynamodb/crypto/
sealed.rs

1use crate::{
2    crypto::attrs::FlattenedEncryptedAttributes,
3    encrypted_table::TableEntry,
4    traits::{ReadConversionError, WriteConversionError},
5    Decryptable, Identifiable,
6};
7use aws_sdk_dynamodb::{primitives::Blob, types::AttributeValue};
8use cipherstash_client::{
9    credentials::{service_credentials::ServiceToken, Credentials},
10    encryption::Encryption,
11};
12use itertools::Itertools;
13use std::{borrow::Cow, collections::HashMap};
14
15use super::{attrs::NormalizedProtectedAttributes, SealError, Unsealed};
16
17// FIXME: Move this to a separate file
18/// Wrapped to indicate that the value is encrypted
19pub struct SealedTableEntry(pub(super) TableEntry);
20
21pub struct UnsealSpec<'a> {
22    pub(crate) protected_attributes: Cow<'a, [Cow<'a, str>]>,
23
24    /// The prefix used for sort keys.
25    /// If None, the type name will be used.
26    /// This *must* be the same as the value used when encrypting the data
27    /// so that descriptors can be correctly matched.
28    /// See [TableAttribute::as_encrypted_record]
29    pub(crate) sort_key_prefix: String,
30}
31
32impl UnsealSpec<'static> {
33    pub fn new_for_decryptable<D>() -> Self
34    where
35        D: Decryptable + Identifiable,
36    {
37        Self {
38            protected_attributes: D::protected_attributes(),
39            sort_key_prefix: D::sort_key_prefix()
40                .as_deref()
41                .map(ToOwned::to_owned)
42                .unwrap_or(D::type_name().to_string()),
43        }
44    }
45}
46
47impl SealedTableEntry {
48    pub fn vec_from<O: TryInto<Self>>(
49        items: impl IntoIterator<Item = O>,
50    ) -> Result<Vec<Self>, <O as TryInto<Self>>::Error> {
51        items.into_iter().map(Self::from_inner).collect()
52    }
53
54    pub(super) fn from_inner<O: TryInto<Self>>(
55        item: O,
56    ) -> Result<Self, <O as TryInto<Self>>::Error> {
57        item.try_into()
58    }
59
60    pub(crate) fn inner(&self) -> &TableEntry {
61        &self.0
62    }
63
64    pub(crate) fn into_inner(self) -> TableEntry {
65        self.0
66    }
67
68    /// Unseal a list of [`Sealed`] values in an efficient manner that optimizes for bulk
69    /// decryptions
70    ///
71    /// This should be used over [`Sealed::unseal`] when multiple values need to be unsealed.
72    pub(crate) async fn unseal_all(
73        items: Vec<Self>,
74        spec: UnsealSpec<'_>,
75        cipher: &Encryption<impl Credentials<Token = ServiceToken>>,
76    ) -> Result<Vec<Unsealed>, SealError> {
77        let UnsealSpec {
78            protected_attributes,
79            sort_key_prefix,
80        } = spec;
81
82        let items_len = items.len();
83        if items_len == 0 {
84            return Ok(Vec::new());
85        }
86
87        let mut protected_items = {
88            let capacity = items.len() * protected_attributes.len();
89            FlattenedEncryptedAttributes::with_capacity(capacity)
90        };
91        let mut unprotected_items = Vec::with_capacity(items.len());
92
93        for item in items.into_iter() {
94            let (protected, unprotected) = item
95                .into_inner()
96                .attributes
97                .partition(protected_attributes.as_ref());
98
99            protected_items.try_extend(protected, sort_key_prefix.clone())?;
100            unprotected_items.push(unprotected);
101        }
102
103        if protected_items.is_empty() {
104            unprotected_items
105                .into_iter()
106                .map(|unprotected| Ok(Unsealed::new_from_unprotected(unprotected)))
107                .collect()
108        } else {
109            let chunk_size =
110                protected_items
111                    .len()
112                    .checked_div(items_len)
113                    .ok_or(SealError::AssertionFailed(
114                        "Division by zero when calculating chunk size".to_string(),
115                    ))?;
116
117            protected_items
118                .decrypt_all(cipher)
119                .await?
120                .into_iter()
121                // TODO: Can we make decrypt_all return a Vec of FlattenedProtectedAttributes? (like the mirror of encrypt_all)
122                .chunks(chunk_size)
123                .into_iter()
124                .map(|fpa| fpa.into_iter().collect::<NormalizedProtectedAttributes>())
125                .zip_eq(unprotected_items.into_iter())
126                .map(|(fpa, unprotected)| Ok(Unsealed::new_from_parts(fpa, unprotected)))
127                .collect()
128        }
129    }
130
131    /// Unseal the current value and return it's plaintext representation
132    ///
133    /// If you need to unseal multiple values at once use [`Sealed::unseal_all`]
134    pub(crate) async fn unseal(
135        self,
136        spec: UnsealSpec<'_>,
137        cipher: &Encryption<impl Credentials<Token = ServiceToken>>,
138    ) -> Result<Unsealed, SealError> {
139        let mut vec = Self::unseal_all(vec![self], spec, cipher).await?;
140
141        if vec.len() != 1 {
142            let actual = vec.len();
143
144            return Err(SealError::AssertionFailed(format!(
145                "Expected unseal_all to return 1 result but got {actual}"
146            )));
147        }
148
149        Ok(vec.remove(0))
150    }
151}
152
153impl TryFrom<HashMap<String, AttributeValue>> for SealedTableEntry {
154    type Error = ReadConversionError;
155
156    fn try_from(item: HashMap<String, AttributeValue>) -> Result<Self, Self::Error> {
157        let pk = item
158            .get("pk")
159            .ok_or(ReadConversionError::NoSuchAttribute("pk".to_string()))?
160            .as_s()
161            .map_err(|_| ReadConversionError::InvalidFormat("pk".to_string()))?
162            .to_string();
163
164        let sk = item
165            .get("sk")
166            .ok_or(ReadConversionError::NoSuchAttribute("sk".to_string()))?
167            .as_s()
168            .map_err(|_| ReadConversionError::InvalidFormat("sk".to_string()))?
169            .to_string();
170
171        let mut table_entry = TableEntry::new(pk, sk);
172
173        // This prevents loading special columns when retrieving records
174        // pk/sk are handled specially or will be called __sk and __pk
175        // We never want to read term during queries
176        item.into_iter()
177            .filter(|(k, _)| k != "pk" && k != "sk" && k != "term")
178            .for_each(|(k, v)| {
179                table_entry.add_attribute(k, v.into());
180            });
181
182        Ok(SealedTableEntry(table_entry))
183    }
184}
185
186// TODO: Test this conversion
187impl TryFrom<SealedTableEntry> for HashMap<String, AttributeValue> {
188    type Error = WriteConversionError;
189
190    fn try_from(item: SealedTableEntry) -> Result<Self, Self::Error> {
191        let mut map = HashMap::new();
192
193        map.insert("pk".to_string(), AttributeValue::S(item.0.pk));
194        map.insert("sk".to_string(), AttributeValue::S(item.0.sk));
195
196        if let Some(term) = item.0.term {
197            map.insert("term".to_string(), AttributeValue::B(Blob::new(term)));
198        }
199
200        item.0.attributes.into_iter().for_each(|(k, v)| {
201            map.insert(k.into_stored_name(), v.into());
202        });
203
204        Ok(map)
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::SealedTableEntry;
211    use cipherstash_client::{
212        credentials::{auto_refresh::AutoRefresh, service_credentials::ServiceCredentials},
213        encryption::Encryption,
214        ConsoleConfig, ZeroKMS, ZeroKMSConfig,
215    };
216    use miette::IntoDiagnostic;
217    use std::borrow::Cow;
218
219    type Cipher = Encryption<AutoRefresh<ServiceCredentials>>;
220
221    // FIXME: Use the test cipher from CipherStash Client when that's ready
222    async fn get_cipher() -> Result<Cipher, Box<dyn std::error::Error>> {
223        let console_config = ConsoleConfig::builder().with_env().build()?;
224        let zero_kms_config = ZeroKMSConfig::builder()
225            .decryption_log(true)
226            .with_env()
227            .console_config(&console_config)
228            .build_with_client_key()?;
229
230        let zero_kms_client = ZeroKMS::new_with_client_key(
231            &zero_kms_config.base_url(),
232            AutoRefresh::new(zero_kms_config.credentials()),
233            zero_kms_config.decryption_log_path().as_deref(),
234            zero_kms_config.client_key(),
235        );
236
237        let config = zero_kms_client.load_dataset_config().await?;
238        Ok(Encryption::new(config.index_root_key, zero_kms_client))
239    }
240
241    #[tokio::test]
242    async fn test_unseal_all_empty() -> Result<(), Box<dyn std::error::Error>> {
243        let spec = super::UnsealSpec {
244            protected_attributes: Cow::Borrowed(&[]),
245            sort_key_prefix: "test".to_string(),
246        };
247        let cipher = get_cipher().await?;
248        let results = SealedTableEntry::unseal_all(vec![], spec, &cipher)
249            .await
250            .into_diagnostic()?;
251        assert!(results.is_empty());
252
253        Ok(())
254    }
255}
cipherstash_dynamodb/crypto/sealed.rs

cipherstash_dynamodb/crypto/
sealed.rs
