Skip to content

Commit cc4e796

Browse files
author
Bennett Hardwick
authored
Merge pull request #19 from cipherstash/feat/add-new-pk-v3
Add new partition and sort key schemes
2 parents 2dc7557 + e8c079e commit cc4e796

34 files changed

+877
-119
lines changed

cryptonamo-derive/src/encryptable.rs

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

12-
let partition_key = format_ident!("{}", settings.get_partition_key()?);
13-
let type_name = settings.sort_key_prefix.to_string();
12+
let partition_key_field = settings.get_partition_key()?;
13+
let partition_key = format_ident!("{partition_key_field}");
14+
let type_name = settings.type_name.clone();
15+
16+
let sort_key_prefix = settings
17+
.sort_key_prefix
18+
.as_ref()
19+
.map(|x| quote! { Some(#x) })
20+
.unwrap_or_else(|| quote! { None });
1421

1522
let protected_attributes = settings.protected_attributes();
1623
let plaintext_attributes = settings.plaintext_attributes();
1724
let ident = settings.ident();
1825

26+
let is_partition_key_encrypted = protected_attributes.contains(&partition_key_field.as_str());
27+
let is_sort_key_encrypted = settings
28+
.sort_key_field
29+
.as_ref()
30+
.map(|x| protected_attributes.contains(&x.as_str()))
31+
.unwrap_or(true);
32+
1933
let into_unsealed_impl = protected_attributes
2034
.iter()
2135
.map(|attr| {
@@ -35,7 +49,14 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
3549

3650
let sort_key_impl = if let Some(sort_key_field) = &settings.sort_key_field {
3751
let sort_key_attr = format_ident!("{sort_key_field}");
38-
quote! { format!("{}#{}", Self::type_name(), self.#sort_key_attr) }
52+
53+
quote! {
54+
if let Some(prefix) = Self::sort_key_prefix() {
55+
format!("{}#{}", prefix, self.#sort_key_attr)
56+
} else {
57+
self.#sort_key_attr.clone()
58+
}
59+
}
3960
} else {
4061
quote! { Self::type_name().into() }
4162
};
@@ -51,6 +72,7 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
5172
impl cryptonamo::traits::Encryptable for #ident {
5273
#primary_key_impl
5374

75+
#[inline]
5476
fn type_name() -> &'static str {
5577
#type_name
5678
}
@@ -59,6 +81,21 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
5981
#sort_key_impl
6082
}
6183

84+
#[inline]
85+
fn sort_key_prefix() -> Option<&'static str> {
86+
#sort_key_prefix
87+
}
88+
89+
#[inline]
90+
fn is_partition_key_encrypted() -> bool {
91+
#is_partition_key_encrypted
92+
}
93+
94+
#[inline]
95+
fn is_sort_key_encrypted() -> bool {
96+
#is_sort_key_encrypted
97+
}
98+
6299
fn partition_key(&self) -> String {
63100
self.#partition_key.to_string()
64101
}
@@ -71,6 +108,7 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
71108
vec![#(#plaintext_attributes,)*]
72109
}
73110

111+
#[allow(clippy::needless_question_mark)]
74112
fn into_sealer(self) -> Result<cryptonamo::crypto::Sealer<Self>, cryptonamo::crypto::SealError> {
75113
Ok(cryptonamo::crypto::Sealer::new_with_descriptor(self, Self::type_name())
76114
#(#into_unsealed_impl?)*)

cryptonamo-derive/src/settings/builder.rs

Lines changed: 120 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,28 @@ use proc_macro2::{Ident, Span};
33
use std::collections::HashMap;
44
use syn::{Data, DeriveInput, Fields, LitStr};
55

6+
enum SortKeyPrefix {
7+
Default,
8+
Value(String),
9+
None,
10+
}
11+
12+
impl SortKeyPrefix {
13+
fn into_prefix(self, default: &str) -> Option<String> {
14+
match self {
15+
Self::Default => Some(default.to_string()),
16+
Self::Value(v) => Some(v),
17+
Self::None => None,
18+
}
19+
}
20+
}
21+
22+
const RESERVED_FIELD_NAMES: &'static [&'static str] = &["term"];
23+
624
pub(crate) struct SettingsBuilder {
725
ident: Ident,
826
type_name: String,
9-
sort_key_prefix: Option<String>,
27+
sort_key_prefix: SortKeyPrefix,
1028
sort_key_field: Option<String>,
1129
partition_key_field: Option<String>,
1230
protected_attributes: Vec<String>,
@@ -33,7 +51,7 @@ impl SettingsBuilder {
3351
Self {
3452
ident: input.ident.clone(),
3553
type_name,
36-
sort_key_prefix: None,
54+
sort_key_prefix: SortKeyPrefix::Default,
3755
sort_key_field: None,
3856
partition_key_field: None,
3957
protected_attributes: Vec::new(),
@@ -54,9 +72,22 @@ impl SettingsBuilder {
5472
match ident.as_deref() {
5573
Some("sort_key_prefix") => {
5674
let value = meta.value()?;
57-
let t: LitStr = value.parse()?;
58-
let v = t.value().to_string();
59-
self.set_sort_key_prefix(v)
75+
76+
if let Ok(t) = value.parse::<LitStr>() {
77+
let v = t.value().to_string();
78+
self.set_sort_key_prefix(v)?;
79+
return Ok(());
80+
}
81+
82+
if let Ok(t) = value.parse::<Ident>() {
83+
let v = t.to_string();
84+
85+
if v == "None" {
86+
self.sort_key_prefix = SortKeyPrefix::None;
87+
}
88+
}
89+
90+
Ok(())
6091
}
6192
Some("partition_key") => {
6293
let value = meta.value()?;
@@ -85,25 +116,89 @@ impl SettingsBuilder {
85116
.flat_map(|x| x.ident.as_ref().map(|x| x.to_string()))
86117
.collect();
87118

119+
let explicit_pk = all_field_names.contains(&String::from("pk"));
120+
let explicit_sk = all_field_names.contains(&String::from("sk"));
121+
88122
let mut compound_indexes: HashMap<String, Vec<(String, String, Span)>> =
89123
Default::default();
90124

91125
for field in &fields_named.named {
92126
let ident = &field.ident;
93127
let mut attr_mode = AttributeMode::Protected;
94128

129+
let field_name = ident
130+
.as_ref()
131+
.ok_or_else(|| {
132+
syn::Error::new_spanned(
133+
field,
134+
"internal error: identifier was not Some",
135+
)
136+
})?
137+
.to_string();
138+
139+
if field_name.starts_with("__") {
140+
return Err(syn::Error::new_spanned(
141+
field,
142+
format!(
143+
"Invalid field '{field_name}': fields must not be prefixed with __"
144+
),
145+
));
146+
}
147+
148+
if RESERVED_FIELD_NAMES.contains(&field_name.as_str()) {
149+
return Err(syn::Error::new_spanned(
150+
field,
151+
format!(
152+
"Invalid field '{field_name}': name is reserved for internal use"
153+
),
154+
));
155+
}
156+
157+
if field_name == "pk" {
158+
let has_partition_key_attr = field
159+
.attrs
160+
.iter()
161+
.find(|x| x.path().is_ident("partition_key"))
162+
.is_some();
163+
164+
if !has_partition_key_attr {
165+
return Err(syn::Error::new_spanned(
166+
field,
167+
format!("field named 'pk' must be annotated with #[partition_key]"),
168+
));
169+
}
170+
}
171+
172+
if field_name == "sk" {
173+
let has_partition_key_attr = field
174+
.attrs
175+
.iter()
176+
.find(|x| x.path().is_ident("sort_key"))
177+
.is_some();
178+
179+
if !has_partition_key_attr {
180+
return Err(syn::Error::new_spanned(
181+
field,
182+
format!("field named 'sk' must be annotated with #[sort_key]"),
183+
));
184+
}
185+
}
186+
95187
// Parse the meta for the field
96188
for attr in &field.attrs {
97189
if attr.path().is_ident("sort_key") {
98-
let field_name = ident
99-
.as_ref()
100-
.ok_or_else(|| {
101-
syn::Error::new_spanned(
102-
field,
103-
"internal error: identifier was not Some",
104-
)
105-
})?
106-
.to_string();
190+
if explicit_sk && field_name != "sk" {
191+
return Err(syn::Error::new_spanned(
192+
field,
193+
format!("field '{field_name}' cannot be used as sort key as struct contains field named 'sk' which must be used")
194+
));
195+
}
196+
197+
if explicit_sk {
198+
// if the 'sk' field is set then there should be no prefix
199+
// otherwise when deserialising the sk value would be incorrect
200+
self.sort_key_prefix = SortKeyPrefix::None;
201+
}
107202

108203
if let Some(f) = &self.sort_key_field {
109204
return Err(syn::Error::new_spanned(
@@ -112,19 +207,16 @@ impl SettingsBuilder {
112207
));
113208
}
114209

115-
self.sort_key_field = Some(field_name);
210+
self.sort_key_field = Some(field_name.clone());
116211
}
117212

118213
if attr.path().is_ident("partition_key") {
119-
let field_name = ident
120-
.as_ref()
121-
.ok_or_else(|| {
122-
syn::Error::new_spanned(
123-
field,
124-
"internal error: identifier was not Some",
125-
)
126-
})?
127-
.to_string();
214+
if explicit_pk && field_name != "pk" {
215+
return Err(syn::Error::new_spanned(
216+
field,
217+
format!("field '{field_name}' cannot be used as partition key as struct contains field named 'pk' which must be used")
218+
));
219+
}
128220

129221
if let Some(f) = &self.partition_key_field {
130222
return Err(syn::Error::new_spanned(
@@ -133,7 +225,7 @@ impl SettingsBuilder {
133225
));
134226
}
135227

136-
self.partition_key_field = Some(field_name);
228+
self.partition_key_field = Some(field_name.clone());
137229
}
138230

139231
if attr.path().is_ident("cryptonamo") {
@@ -255,7 +347,7 @@ impl SettingsBuilder {
255347
indexes,
256348
} = self;
257349

258-
let sort_key_prefix = sort_key_prefix.unwrap_or(type_name);
350+
let sort_key_prefix = sort_key_prefix.into_prefix(&type_name);
259351

260352
let partition_key = partition_key.ok_or_else(|| {
261353
syn::Error::new(
@@ -267,6 +359,7 @@ impl SettingsBuilder {
267359
Ok(Settings {
268360
ident,
269361
sort_key_prefix,
362+
type_name,
270363
sort_key_field,
271364
partition_key_field: Some(partition_key), // TODO: Remove the Some
272365
protected_attributes,
@@ -277,7 +370,7 @@ impl SettingsBuilder {
277370
}
278371

279372
pub(crate) fn set_sort_key_prefix(&mut self, value: String) -> Result<(), syn::Error> {
280-
self.sort_key_prefix = Some(value);
373+
self.sort_key_prefix = SortKeyPrefix::Value(value);
281374
Ok(())
282375
}
283376

cryptonamo-derive/src/settings/mod.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,8 @@ pub(crate) enum AttributeMode {
1515

1616
pub(crate) struct Settings {
1717
ident: Ident,
18-
pub(crate) sort_key_prefix: String,
18+
pub(crate) sort_key_prefix: Option<String>,
19+
pub(crate) type_name: String,
1920
pub(crate) sort_key_field: Option<String>,
2021
pub(crate) partition_key_field: Option<String>,
2122
protected_attributes: Vec<String>,

src/crypto/mod.rs

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,10 @@ mod unsealed;
55
use crate::traits::{Encryptable, ReadConversionError, Searchable, WriteConversionError};
66
use cipherstash_client::{
77
credentials::{vitur_credentials::ViturToken, Credentials},
8-
encryption::{Encryption, EncryptionError, Plaintext, TypeParseError},
9-
schema::column::Index,
8+
encryption::{
9+
compound_indexer::{CompoundIndex, ExactIndex},
10+
Encryption, EncryptionError, Plaintext, TypeParseError,
11+
},
1012
};
1113
use thiserror::Error;
1214

@@ -56,18 +58,26 @@ pub(crate) fn all_index_keys<E: Searchable + Encryptable>(sort_key: &str) -> Vec
5658
.collect()
5759
}
5860

59-
pub(crate) fn encrypt_partition_key<C>(
61+
pub(crate) fn hmac<C>(
62+
field: &str,
6063
value: &str,
64+
salt: Option<&str>,
6165
cipher: &Encryption<C>,
6266
) -> Result<String, EncryptionError>
6367
where
6468
C: Credentials<Token = ViturToken>,
6569
{
6670
let plaintext = Plaintext::Utf8Str(Some(value.to_string()));
67-
let index_type = Index::new_unique().index_type;
71+
let index = CompoundIndex::new(ExactIndex::new(field, vec![]));
6872

6973
cipher
70-
.index(&plaintext, &index_type)?
74+
.compound_index(
75+
&index,
76+
plaintext,
77+
// passing None here results in no terms so pass an empty string
78+
Some(salt.unwrap_or("")),
79+
32,
80+
)?
7181
.as_binary()
7282
.map(hex::encode)
7383
.ok_or(EncryptionError::IndexingError(

0 commit comments

Comments
 (0)