Skip to content

Commit 0435a6a

Browse files
Nicklas Warming JacobsenBennett Hardwick
Nicklas Warming Jacobsen
authored and
Bennett Hardwick
committed
Move fn type_name and fn sort_key_prefix to trait Identifiable.
Before the functions `fn type_name` and `fn sort_key_prefix` existed on both the trait `Encryptable` and `Decryptable`, which made it hard to make proper trait bounds. This PR moves both functions to the trait `Identifiable`.
1 parent a54e629 commit 0435a6a

File tree

12 files changed

+128
-99
lines changed

12 files changed

+128
-99
lines changed

‎cipherstash-dynamodb-derive/src/decryptable.rs‎

Lines changed: 6 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -12,25 +12,19 @@ pub(crate) fn derive_decryptable(input: DeriveInput) -> Result<TokenStream, syn:
1212
let protected_attributes = settings.protected_attributes();
1313
let plaintext_attributes = settings.plaintext_attributes();
1414

15-
let protected_attributes_cow = settings.protected_attributes()
15+
let protected_attributes_cow = settings
16+
.protected_attributes()
1617
.into_iter()
17-
.map(|x| quote!{ std::borrow::Cow::Borrowed(#x) });
18+
.map(|x| quote! { std::borrow::Cow::Borrowed(#x) });
1819

19-
let plaintext_attributes_cow = settings.plaintext_attributes()
20+
let plaintext_attributes_cow = settings
21+
.plaintext_attributes()
2022
.into_iter()
21-
.map(|x| quote!{ std::borrow::Cow::Borrowed(#x) });
23+
.map(|x| quote! { std::borrow::Cow::Borrowed(#x) });
2224

2325
let skipped_attributes = settings.skipped_attributes();
2426
let ident = settings.ident();
2527

26-
let type_name = &settings.type_name;
27-
28-
let sort_key_prefix_impl = if let Some(prefix) = &settings.sort_key_prefix {
29-
quote! { Some(std::borrow::Cow::Borrowed(#prefix)) }
30-
} else {
31-
quote! { None }
32-
};
33-
3428
let from_unsealed_impl = protected_attributes
3529
.iter()
3630
.map(|attr| {
@@ -58,16 +52,6 @@ pub(crate) fn derive_decryptable(input: DeriveInput) -> Result<TokenStream, syn:
5852
let expanded = quote! {
5953
#[automatically_derived]
6054
impl cipherstash_dynamodb::traits::Decryptable for #ident {
61-
#[inline]
62-
fn type_name() -> std::borrow::Cow<'static, str> {
63-
std::borrow::Cow::Borrowed(#type_name)
64-
}
65-
66-
#[inline]
67-
fn sort_key_prefix() -> Option<std::borrow::Cow<'static, str>> {
68-
#sort_key_prefix_impl
69-
}
70-
7155
fn protected_attributes() -> std::borrow::Cow<'static, [std::borrow::Cow<'static, str>]> {
7256
std::borrow::Cow::Borrowed(&[#(#protected_attributes_cow,)*])
7357
}

‎cipherstash-dynamodb-derive/src/encryptable.rs‎

Lines changed: 7 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -9,25 +9,18 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
99
.field_attributes(&input)?
1010
.build()?;
1111

12-
let type_name = settings.type_name.clone();
13-
14-
let sort_key_prefix_impl = if let Some(prefix) = &settings.sort_key_prefix {
15-
quote! { Some(std::borrow::Cow::Borrowed(#prefix)) }
16-
} else {
17-
quote! { None }
18-
};
19-
2012
let protected_attributes = settings.protected_attributes();
2113
let plaintext_attributes = settings.plaintext_attributes();
2214

23-
let protected_attributes_cow = settings.protected_attributes()
15+
let protected_attributes_cow = settings
16+
.protected_attributes()
2417
.into_iter()
25-
.map(|x| quote!{ std::borrow::Cow::Borrowed(#x) });
18+
.map(|x| quote! { std::borrow::Cow::Borrowed(#x) });
2619

27-
let plaintext_attributes_cow = settings.plaintext_attributes()
20+
let plaintext_attributes_cow = settings
21+
.plaintext_attributes()
2822
.into_iter()
29-
.map(|x| quote!{ std::borrow::Cow::Borrowed(#x) });
30-
23+
.map(|x| quote! { std::borrow::Cow::Borrowed(#x) });
3124

3225
let ident = settings.ident();
3326

@@ -51,16 +44,6 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
5144
let expanded = quote! {
5245
#[automatically_derived]
5346
impl cipherstash_dynamodb::traits::Encryptable for #ident {
54-
#[inline]
55-
fn type_name() -> std::borrow::Cow<'static, str> {
56-
std::borrow::Cow::Borrowed(#type_name)
57-
}
58-
59-
#[inline]
60-
fn sort_key_prefix() -> Option<std::borrow::Cow<'static, str>> {
61-
#sort_key_prefix_impl
62-
}
63-
6447
fn protected_attributes() -> std::borrow::Cow<'static, [std::borrow::Cow<'static, str>]> {
6548
std::borrow::Cow::Borrowed(&[#(#protected_attributes_cow,)*])
6649
}
@@ -71,7 +54,7 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
7154

7255
#[allow(clippy::needless_question_mark)]
7356
fn into_unsealed(self) -> cipherstash_dynamodb::crypto::Unsealed {
74-
let mut unsealed = cipherstash_dynamodb::crypto::Unsealed::new_with_descriptor(<Self as cipherstash_dynamodb::traits::Encryptable>::type_name());
57+
let mut unsealed = cipherstash_dynamodb::crypto::Unsealed::new_with_descriptor(<Self as cipherstash_dynamodb::traits::Identifiable>::type_name());
7558

7659
#(#into_unsealed_impl)*
7760

‎cipherstash-dynamodb-derive/src/identifiable.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,11 +53,28 @@ pub(crate) fn derive_identifiable(input: DeriveInput) -> Result<TokenStream, syn
5353
}
5454
}
5555
};
56+
let type_name = &settings.type_name;
57+
58+
let sort_key_prefix_impl = if let Some(prefix) = &settings.sort_key_prefix {
59+
quote! { Some(std::borrow::Cow::Borrowed(#prefix)) }
60+
} else {
61+
quote! { None }
62+
};
5663

5764
let expanded = quote! {
5865
impl cipherstash_dynamodb::traits::Identifiable for #ident {
5966
#primary_key_impl
6067

68+
#[inline]
69+
fn type_name() -> std::borrow::Cow<'static, str> {
70+
std::borrow::Cow::Borrowed(#type_name)
71+
}
72+
73+
#[inline]
74+
fn sort_key_prefix() -> Option<std::borrow::Cow<'static, str>> {
75+
#sort_key_prefix_impl
76+
}
77+
6178
fn is_pk_encrypted() -> bool {
6279
#is_partition_key_encrypted
6380
}

‎src/crypto/mod.rs‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,30 @@ where
106106
))
107107
}
108108

109+
// Contains all the necessary information to encrypt the primary key pair
110+
pub struct PreparedPrimaryKey {
111+
pub primary_key_parts: PrimaryKeyParts,
112+
pub is_pk_encrypted: bool,
113+
pub is_sk_encrypted: bool,
114+
}
115+
116+
impl PreparedPrimaryKey {
117+
pub fn new<R>(k: impl Into<R::PrimaryKey>) -> Self
118+
where
119+
R: Identifiable,
120+
{
121+
let primary_key_parts = k
122+
.into()
123+
.into_parts(&R::type_name(), R::sort_key_prefix().as_deref());
124+
125+
Self {
126+
primary_key_parts,
127+
is_pk_encrypted: R::is_pk_encrypted(),
128+
is_sk_encrypted: R::is_sk_encrypted(),
129+
}
130+
}
131+
}
132+
109133
pub fn encrypt_primary_key<I: Identifiable>(
110134
k: impl Into<I::PrimaryKey>,
111135
type_name: &str,

‎src/encrypted_table/mod.rs‎

Lines changed: 29 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ pub use self::{
88
use crate::{
99
crypto::*,
1010
errors::*,
11-
traits::{Decryptable, PrimaryKey, PrimaryKeyParts, Searchable},
11+
traits::{Decryptable, PrimaryKey, PrimaryKeyError, PrimaryKeyParts, Searchable},
1212
Identifiable, IndexType,
1313
};
1414
use aws_sdk_dynamodb::types::{AttributeValue, Delete, Put, TransactWriteItem};
@@ -264,12 +264,15 @@ impl<D> EncryptedTable<D> {
264264
&self,
265265
k: impl Into<E::PrimaryKey>,
266266
) -> Result<DynamoRecordPatch, DeleteError> {
267-
let PrimaryKeyParts { pk, sk } = encrypt_primary_key::<E>(
268-
k,
269-
&E::type_name(),
270-
E::sort_key_prefix().as_deref(),
271-
&self.cipher,
272-
)?;
267+
// let PrimaryKeyParts { pk, sk } = encrypt_primary_key::<E>(
268+
// k,
269+
// &E::type_name(),
270+
// E::sort_key_prefix().as_deref(),
271+
// &self.cipher,
272+
// )?;
273+
274+
let PrimaryKeyParts { pk, sk } =
275+
self.encrypt_primary_key_parts(PreparedPrimaryKey::new::<E>(k))?;
273276

274277
let delete_records = all_index_keys(&sk, E::protected_indexes())
275278
.into_iter()
@@ -343,6 +346,23 @@ impl<D> EncryptedTable<D> {
343346
delete_records,
344347
})
345348
}
349+
350+
pub fn encrypt_primary_key_parts(
351+
&self,
352+
prepared_primary_key: PreparedPrimaryKey,
353+
) -> Result<PrimaryKeyParts, PrimaryKeyError> {
354+
let PrimaryKeyParts { mut pk, mut sk } = prepared_primary_key.primary_key_parts;
355+
356+
if prepared_primary_key.is_pk_encrypted {
357+
pk = b64_encode(hmac(&pk, None, &self.cipher)?);
358+
}
359+
360+
if prepared_primary_key.is_sk_encrypted {
361+
sk = b64_encode(hmac(&sk, Some(pk.as_str()), &self.cipher)?);
362+
}
363+
364+
Ok(PrimaryKeyParts { pk, sk })
365+
}
346366
}
347367

348368
impl EncryptedTable<Dynamo> {
@@ -365,12 +385,8 @@ impl EncryptedTable<Dynamo> {
365385
where
366386
T: Decryptable + Identifiable,
367387
{
368-
let PrimaryKeyParts { pk, sk } = encrypt_primary_key::<T>(
369-
k,
370-
&T::type_name(),
371-
T::sort_key_prefix().as_deref(),
372-
&self.cipher,
373-
)?;
388+
let PrimaryKeyParts { pk, sk } =
389+
self.encrypt_primary_key_parts(PreparedPrimaryKey::new::<T>(k))?;
374390

375391
let result = self
376392
.db

‎src/encrypted_table/query.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use std::marker::PhantomData;
88

99
use crate::{
1010
traits::{Decryptable, Searchable},
11-
IndexType, SingleIndex,
11+
Identifiable, IndexType, SingleIndex,
1212
};
1313
use cipherstash_client::encryption::{compound_indexer::CompoundIndex, IndexTerm};
1414

@@ -115,7 +115,7 @@ where
115115

116116
impl<'t, S> QueryBuilder<'t, S, Dynamo>
117117
where
118-
S: Searchable,
118+
S: Searchable + Identifiable,
119119
{
120120
pub async fn load<T: Decryptable>(self) -> Result<Vec<T>, QueryError> {
121121
let (index_name, index, plaintext, builder) = self.build()?;
@@ -162,7 +162,7 @@ where
162162

163163
impl<'t, S> QueryBuilder<'t, S, Dynamo>
164164
where
165-
S: Searchable + Decryptable,
165+
S: Searchable + Decryptable + Identifiable,
166166
{
167167
pub async fn send(self) -> Result<Vec<S>, QueryError> {
168168
self.load::<S>().await

0 commit comments

Comments
 (0)