Skip to content

Commit f5fc40f

Browse files
jot2rejot2re
committed
refactor: backup_priv_data from macro to function
1 parent 22ecd85 commit f5fc40f

File tree

1 file changed

+84
-66
lines changed

1 file changed

+84
-66
lines changed

‎core/service/src/engine/context_manager.rs‎

Lines changed: 84 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use crate::cryptography::backup_pke::{self, BackupCiphertext};
55
use crate::cryptography::internal_crypto_types::PrivateSigKey;
66
use crate::engine::context::ContextInfo;
77
use crate::engine::threshold::service::ThresholdFheKeys;
8+
use crate::vault::Vault;
89
use crate::{
910
engine::{
1011
base::BaseKmsStruct, threshold::traits::ContextManager, validation::validate_request_id,
@@ -90,37 +91,54 @@ where
9091
}
9192
}
9293

93-
macro_rules! backup_priv_data {
94-
($rng:expr, $guarded_priv_storage:expr, $guarded_backup_vault:expr, $cur_type:expr, $data_type:ty, $pub_enc_key:expr) => {
95-
let data_ids = $guarded_priv_storage
96-
.all_data_ids(&$cur_type.to_string())
94+
async fn backup_priv_data<
95+
S1: Storage + Sync + Send + 'static,
96+
T: serde::de::DeserializeOwned
97+
+ tfhe::Unversionize
98+
+ tfhe::named::Named
99+
+ Send
100+
+ serde::ser::Serialize
101+
+ tfhe::Versionize
102+
+ Sync
103+
+ 'static,
104+
>(
105+
rng: &mut AesRng,
106+
priv_storage: &S1,
107+
backup_vault: &mut Vault,
108+
data_type_enum: PrivDataType,
109+
pub_enc_key: &backup_pke::BackupPublicKey,
110+
) -> anyhow::Result<()>
111+
where
112+
for<'a> <T as tfhe::Versionize>::Versioned<'a>: Send + Sync,
113+
{
114+
let data_ids = priv_storage
115+
.all_data_ids(&data_type_enum.to_string())
116+
.await?;
117+
for data_id in data_ids.iter() {
118+
let data: T = priv_storage
119+
.read_data(&data_id, &data_type_enum.to_string())
97120
.await?;
98-
for data_id in data_ids {
99-
let data: $data_type = $guarded_priv_storage
100-
.read_data(&data_id, &$cur_type.to_string())
101-
.await?;
102-
let mut serialized_data = Vec::new();
103-
safe_serialize(&data, &mut serialized_data, SAFE_SER_SIZE_LIMIT)?;
104-
let encrypted_data = $pub_enc_key.encrypt($rng, &serialized_data)?;
105-
let enc_ct = BackupCiphertext {
106-
ciphertext: encrypted_data,
107-
priv_data_type: $cur_type,
108-
};
109-
110-
// Delete the old backup data
111-
// Observe that no backups from previous contexts are deleted, only current context.
112-
$guarded_backup_vault
113-
.delete_data(&data_id, &$cur_type.to_string())
114-
.await?;
115-
$guarded_backup_vault
116-
.store_data(
117-
&enc_ct,
118-
&data_id,
119-
&BackupDataType::PrivData($cur_type).to_string(),
120-
)
121-
.await?;
122-
}
123-
};
121+
let mut serialized_data = Vec::new();
122+
safe_serialize(&data, &mut serialized_data, SAFE_SER_SIZE_LIMIT)?;
123+
let encrypted_data = pub_enc_key.encrypt(rng, &serialized_data)?;
124+
let enc_ct = BackupCiphertext {
125+
ciphertext: encrypted_data,
126+
priv_data_type: data_type_enum,
127+
};
128+
// Delete the old backup data
129+
// Observe that no backups from previous contexts are deleted, only backups for current custodian context in case they exist.
130+
backup_vault
131+
.delete_data(&data_id, &data_type_enum.to_string())
132+
.await?;
133+
backup_vault
134+
.store_data(
135+
&enc_ct,
136+
&data_id,
137+
&BackupDataType::PrivData(data_type_enum).to_string(),
138+
)
139+
.await?;
140+
}
141+
Ok(())
124142
}
125143

126144
impl<PubS, PrivS> RealContextManager<PubS, PrivS>
@@ -185,68 +203,68 @@ where
185203
// We need to match on each type to manually specify the data type and to ensure that we do not forget anything in case the enum is extended
186204
match cur_type {
187205
PrivDataType::SigningKey => {
188-
backup_priv_data!(
206+
backup_priv_data::<PrivS, PrivateSigKey>(
189207
&mut rng,
190-
guarded_priv_storage,
191-
guarded_backup_vault,
208+
&guarded_priv_storage,
209+
&mut guarded_backup_vault,
192210
cur_type,
193-
PrivateSigKey,
194-
backup_enc_key
195-
);
211+
&backup_enc_key,
212+
)
213+
.await?;
196214
}
197215
PrivDataType::FheKeyInfo => {
198-
backup_priv_data!(
216+
backup_priv_data::<PrivS, ThresholdFheKeys>(
199217
&mut rng,
200-
guarded_priv_storage,
201-
guarded_backup_vault,
218+
&guarded_priv_storage,
219+
&mut guarded_backup_vault,
202220
cur_type,
203-
ThresholdFheKeys,
204-
backup_enc_key
205-
);
221+
&backup_enc_key,
222+
)
223+
.await?;
206224
}
207225
PrivDataType::CrsInfo => {
208-
backup_priv_data!(
226+
backup_priv_data::<PrivS, SignedPubDataHandleInternal>(
209227
&mut rng,
210-
guarded_priv_storage,
211-
guarded_backup_vault,
228+
&guarded_priv_storage,
229+
&mut guarded_backup_vault,
212230
cur_type,
213-
SignedPubDataHandleInternal,
214-
backup_enc_key
215-
);
231+
&backup_enc_key,
232+
)
233+
.await?;
216234
}
217235
PrivDataType::FhePrivateKey => {
218-
backup_priv_data!(
236+
backup_priv_data::<PrivS, ClientKey>(
219237
&mut rng,
220-
guarded_priv_storage,
221-
guarded_backup_vault,
238+
&guarded_priv_storage,
239+
&mut guarded_backup_vault,
222240
cur_type,
223-
ClientKey,
224-
backup_enc_key
225-
);
241+
&backup_enc_key,
242+
)
243+
.await?;
226244
}
227245
PrivDataType::PrssSetup => {
228246
// We will not back up PRSS setup data
229247
continue;
230248
}
231249
PrivDataType::CustodianInfo => {
232-
backup_priv_data!(
250+
backup_priv_data::<PrivS, InternalCustodianContext>(
233251
&mut rng,
234-
guarded_priv_storage,
235-
guarded_backup_vault,
252+
&guarded_priv_storage,
253+
&mut guarded_backup_vault,
236254
cur_type,
237-
InternalCustodianContext,
238-
backup_enc_key
239-
);
255+
&backup_enc_key,
256+
)
257+
.await?;
240258
}
241259
PrivDataType::ContextInfo => {
242-
backup_priv_data!(
260+
backup_priv_data::<PrivS, ContextInfo>(
243261
&mut rng,
244-
guarded_priv_storage,
245-
guarded_backup_vault,
262+
&guarded_priv_storage,
263+
&mut guarded_backup_vault,
246264
cur_type,
247-
ContextInfo,
248-
backup_enc_key
249-
);
265+
&backup_enc_key,
266+
)
267+
.await?;
250268
}
251269
}
252270
}

0 commit comments

Comments
 (0)