review comments

iotaledger · Apr 25, 2024 · 23a2f46 · 23a2f46
1 parent 6a32429
commit 23a2f46
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 16 deletions.
diff --git a/identity_storage/src/key_storage/bls.rs b/identity_storage/src/key_storage/bls.rs
@@ -173,15 +173,21 @@ where
 /// Updates BBS+ signature's timeframe data.
 pub fn update_bbs_signature(
   alg: ProofAlgorithm,
-  sig: &[u8; 80],
+  sig: &[u8],
   sk: &BBSplusSecretKey,
   update_ctx: &ProofUpdateCtx,
-) -> Result<[u8; 80], KeyStorageError> {
+) -> KeyStorageResult<Vec<u8>> {
+  let exact_size_signature = sig.try_into().map_err(|_| {
+    KeyStorageError::new(KeyStorageErrorKind::Unspecified).with_custom_message("invalid signature size".to_owned())
+  })?;
   match alg {
-    ProofAlgorithm::BLS12381_SHA256 => _update_bbs_signature::<Bls12381Sha256>(sig, sk, update_ctx),
-    ProofAlgorithm::BLS12381_SHAKE256 => _update_bbs_signature::<Bls12381Shake256>(sig, sk, update_ctx),
+    ProofAlgorithm::BLS12381_SHA256 => _update_bbs_signature::<Bls12381Sha256>(exact_size_signature, sk, update_ctx),
+    ProofAlgorithm::BLS12381_SHAKE256 => {
+      _update_bbs_signature::<Bls12381Shake256>(exact_size_signature, sk, update_ctx)
+    }
     _ => return Err(KeyStorageErrorKind::UnsupportedProofAlgorithm.into()),
   }
+  .map(Vec::from)
   .map_err(|e| {
     KeyStorageError::new(KeyStorageErrorKind::Unspecified)
       .with_custom_message("signature failed")

diff --git a/identity_storage/src/key_storage/jwk_storage_bbs_plus_ext.rs b/identity_storage/src/key_storage/jwk_storage_bbs_plus_ext.rs
@@ -1,7 +1,6 @@
 use async_trait::async_trait;
 use identity_verification::jwk::Jwk;
 use jsonprooftoken::jpa::algs::ProofAlgorithm;
-use zkryptium::bbsplus::signature::BBSplusSignature;
 
 use crate::JwkGenOutput;
 use crate::JwkStorage;
@@ -32,7 +31,7 @@ pub trait JwkStorageBbsPlusExt: JwkStorage {
     &self,
     key_id: &KeyId,
     public_key: &Jwk,
-    signature: &[u8; BBSplusSignature::BYTES],
+    signature: &[u8],
     ctx: ProofUpdateCtx,
-  ) -> KeyStorageResult<[u8; BBSplusSignature::BYTES]>;
+  ) -> KeyStorageResult<Vec<u8>>;
 }
diff --git a/identity_storage/src/key_storage/memstore.rs b/identity_storage/src/key_storage/memstore.rs
@@ -323,7 +323,6 @@ mod bbs_plus_impl {
   use identity_verification::jwk::BlsCurve;
   use identity_verification::jwk::Jwk;
   use jsonprooftoken::jpa::algs::ProofAlgorithm;
-  use zkryptium::bbsplus::signature::BBSplusSignature;
 
   use super::random_key_id;
 
@@ -386,9 +385,9 @@ mod bbs_plus_impl {
       &self,
       key_id: &KeyId,
       public_key: &Jwk,
-      signature: &[u8; BBSplusSignature::BYTES],
+      signature: &[u8],
       ctx: ProofUpdateCtx,
-    ) -> KeyStorageResult<[u8; BBSplusSignature::BYTES]> {
+    ) -> KeyStorageResult<Vec<u8>> {
       let jwk_store = self.jwk_store.read().await;
 
       // Extract the required alg from the given public key

diff --git a/identity_stronghold/src/stronghold_jwk_storage_ext.rs b/identity_stronghold/src/stronghold_jwk_storage_ext.rs
@@ -23,7 +23,6 @@ use jsonprooftoken::jpa::algs::ProofAlgorithm;
 use std::str::FromStr;
 use zeroize::Zeroizing;
 use zkryptium::bbsplus::keys::BBSplusSecretKey;
-use zkryptium::bbsplus::signature::BBSplusSignature;
 
 use crate::stronghold_key_type::*;
 use crate::utils::*;
@@ -46,10 +45,14 @@ impl JwkStorageBbsPlusExt for StrongholdStorage {
       return Err(KeyStorageErrorKind::UnsupportedProofAlgorithm.into());
     }
 
+    // Get a key id that's not already used.
+    let mut kid = random_key_id();
+    while self.exists(&kid).await? {
+      kid = random_key_id();
+    }
+
     let stronghold = self.get_stronghold().await;
     let client = get_client(&stronghold)?;
-
-    let kid: KeyId = random_key_id();
     let target_key_location = Location::generic(
       IDENTITY_VAULT_PATH.as_bytes().to_vec(),
       kid.to_string().as_bytes().to_vec(),
@@ -109,7 +112,11 @@ impl JwkStorageBbsPlusExt for StrongholdStorage {
     client
       .get_guards([sk_location], |[sk]| {
         let sk = BBSplusSecretKey::from_bytes(&sk.borrow()).map_err(|e| FatalProcedureError::from(e.to_string()))?;
-        sign_bbs(alg, data, &sk, &pk, header).map_err(|e| FatalProcedureError::from(e.to_string()))
+        let signature_result =
+          sign_bbs(alg, data, &sk, &pk, header).map_err(|e| FatalProcedureError::from(e.to_string()));
+        // clean up `sk` to avoid leaking.
+        drop(Zeroizing::new(sk.to_bytes()));
+        signature_result
       })
       .map(|sig| sig.to_vec())
       .map_err(|e| {
@@ -123,9 +130,9 @@ impl JwkStorageBbsPlusExt for StrongholdStorage {
     &self,
     key_id: &KeyId,
     public_key: &Jwk,
-    signature: &[u8; BBSplusSignature::BYTES],
+    signature: &[u8],
     ctx: ProofUpdateCtx,
-  ) -> KeyStorageResult<[u8; BBSplusSignature::BYTES]> {
+  ) -> KeyStorageResult<Vec<u8>> {
     // Extract the required alg from the given public key
     let alg = public_key
       .alg()