identity_storage/key_storage/
memstore.rs

1// Copyright 2020-2025 IOTA Stiftung, Fondazione Links
2// SPDX-License-Identifier: Apache-2.0
3
4use core::fmt::Debug;
5use std::collections::HashMap;
6use std::fmt::Display;
7use std::str::FromStr;
8
9use async_trait::async_trait;
10use fastcrypto::ed25519::Ed25519KeyPair;
11use fastcrypto::ed25519::Ed25519Signature;
12use fastcrypto::traits::KeyPair as _;
13use fastcrypto::traits::Signer;
14use identity_verification::jose::jwk::EdCurve;
15use identity_verification::jose::jwk::Jwk;
16use identity_verification::jose::jwk::JwkType;
17use identity_verification::jose::jws::JwsAlgorithm;
18use identity_verification::jwk::BlsCurve;
19use identity_verification::jwk::FromJwk as _;
20use identity_verification::jwk::ToJwk as _;
21use rand::distributions::DistString;
22use shared::Shared;
23use tokio::sync::RwLockReadGuard;
24use tokio::sync::RwLockWriteGuard;
25
26use super::jwk_gen_output::JwkGenOutput;
27use super::KeyId;
28use super::KeyStorageError;
29use super::KeyStorageErrorKind;
30use super::KeyStorageResult;
31use super::KeyType;
32use crate::key_storage::JwkStorage;
33
34/// The map from key ids to JWKs.
35type JwkKeyStore = HashMap<KeyId, Jwk>;
36
37/// An insecure, in-memory [`JwkStorage`] implementation that serves as an example and may be used in tests.
38#[derive(Debug)]
39pub struct JwkMemStore {
40  jwk_store: Shared<JwkKeyStore>,
41}
42
43impl JwkMemStore {
44  /// Creates a new, empty `JwkMemStore` instance.
45  pub fn new() -> Self {
46    Self {
47      jwk_store: Shared::new(HashMap::new()),
48    }
49  }
50
51  /// Returns the number of items contained in the [`JwkMemStore`].
52  pub async fn count(&self) -> usize {
53    self.jwk_store.read().await.keys().count()
54  }
55}
56
57// Refer to the `JwkStorage` interface docs for high-level documentation of the individual methods.
58#[cfg_attr(not(feature = "send-sync-storage"), async_trait(?Send))]
59#[cfg_attr(feature = "send-sync-storage", async_trait)]
60impl JwkStorage for JwkMemStore {
61  async fn generate(&self, key_type: KeyType, alg: JwsAlgorithm) -> KeyStorageResult<JwkGenOutput> {
62    let key_type: MemStoreKeyType = MemStoreKeyType::try_from(&key_type)?;
63
64    check_key_alg_compatibility(key_type, &alg)?;
65
66    let keypair = match key_type {
67      MemStoreKeyType::Ed25519 => Ed25519KeyPair::generate(&mut rand::thread_rng()),
68      other => {
69        return Err(
70          KeyStorageError::new(KeyStorageErrorKind::UnsupportedKeyType)
71            .with_custom_message(format!("{other} is not supported")),
72        );
73      }
74    };
75
76    let kid: KeyId = random_key_id();
77
78    let mut jwk: Jwk = keypair.to_jwk().map_err(|err| {
79      KeyStorageError::new(KeyStorageErrorKind::Unspecified)
80        .with_custom_message("could not convert `Ed25519KeyPair` to `Jwk`")
81        .with_source(err)
82    })?;
83    jwk.set_alg(alg.name());
84    jwk.set_kid(jwk.thumbprint_sha256_b64());
85    let public_jwk: Jwk = jwk.to_public().expect("should only panic if kty == oct");
86
87    let mut jwk_store: RwLockWriteGuard<'_, JwkKeyStore> = self.jwk_store.write().await;
88    jwk_store.insert(kid.clone(), jwk);
89
90    Ok(JwkGenOutput::new(kid, public_jwk))
91  }
92
93  async fn insert(&self, jwk: Jwk) -> KeyStorageResult<KeyId> {
94    let key_type = MemStoreKeyType::try_from(&jwk)?;
95
96    if !jwk.is_private() {
97      return Err(
98        KeyStorageError::new(KeyStorageErrorKind::Unspecified)
99          .with_custom_message("expected a Jwk with all private key components set"),
100      );
101    }
102
103    match jwk.alg() {
104      Some(alg) => {
105        let alg: JwsAlgorithm = JwsAlgorithm::from_str(alg)
106          .map_err(|err| KeyStorageError::new(KeyStorageErrorKind::UnsupportedSignatureAlgorithm).with_source(err))?;
107        check_key_alg_compatibility(key_type, &alg)?;
108      }
109      None => {
110        return Err(
111          KeyStorageError::new(KeyStorageErrorKind::UnsupportedSignatureAlgorithm)
112            .with_custom_message("expected a Jwk with an `alg` parameter"),
113        );
114      }
115    }
116
117    if jwk.alg().is_none() {
118      return Err(
119        KeyStorageError::new(KeyStorageErrorKind::UnsupportedSignatureAlgorithm)
120          .with_custom_message("expected a Jwk with an `alg` parameter"),
121      );
122    }
123
124    let key_id: KeyId = random_key_id();
125
126    let mut jwk_store: RwLockWriteGuard<'_, JwkKeyStore> = self.jwk_store.write().await;
127
128    jwk_store.insert(key_id.clone(), jwk);
129
130    Ok(key_id)
131  }
132
133  async fn sign(&self, key_id: &KeyId, data: &[u8], public_key: &Jwk) -> KeyStorageResult<Vec<u8>> {
134    let jwk_store: RwLockReadGuard<'_, JwkKeyStore> = self.jwk_store.read().await;
135
136    // Extract the required alg from the given public key
137    let alg = public_key
138      .alg()
139      .ok_or(KeyStorageErrorKind::UnsupportedSignatureAlgorithm)
140      .and_then(|alg_str| {
141        JwsAlgorithm::from_str(alg_str).map_err(|_| KeyStorageErrorKind::UnsupportedSignatureAlgorithm)
142      })?;
143
144    // Check that `kty` is `Okp` and `crv = Ed25519`.
145    match alg {
146      JwsAlgorithm::EdDSA => {
147        let okp_params = public_key.try_okp_params().map_err(|err| {
148          KeyStorageError::new(KeyStorageErrorKind::Unspecified)
149            .with_custom_message(format!("expected a Jwk with Okp params in order to sign with {alg}"))
150            .with_source(err)
151        })?;
152        if okp_params.crv != EdCurve::Ed25519.name() {
153          return Err(
154            KeyStorageError::new(KeyStorageErrorKind::Unspecified).with_custom_message(format!(
155              "expected Jwk with Okp {} crv in order to sign with {alg}",
156              EdCurve::Ed25519
157            )),
158          );
159        }
160      }
161      other => {
162        return Err(
163          KeyStorageError::new(KeyStorageErrorKind::UnsupportedSignatureAlgorithm)
164            .with_custom_message(format!("{other} is not supported")),
165        );
166      }
167    };
168
169    // Obtain the corresponding private key and sign `data`.
170    let jwk: &Jwk = jwk_store
171      .get(key_id)
172      .ok_or_else(|| KeyStorageError::new(KeyStorageErrorKind::KeyNotFound))?;
173    let secret_key = Ed25519KeyPair::from_jwk(jwk).map_err(|err| {
174      KeyStorageError::new(KeyStorageErrorKind::Unspecified)
175        .with_custom_message("could not convert `Jwk` to `Ed25519KeyPair`")
176        .with_source(err)
177    })?;
178    Ok(Signer::<Ed25519Signature>::sign(&secret_key, data).as_ref().to_vec())
179  }
180
181  async fn delete(&self, key_id: &KeyId) -> KeyStorageResult<()> {
182    let mut jwk_store: RwLockWriteGuard<'_, JwkKeyStore> = self.jwk_store.write().await;
183
184    jwk_store
185      .remove(key_id)
186      .map(|_| ())
187      .ok_or_else(|| KeyStorageError::new(KeyStorageErrorKind::KeyNotFound))
188  }
189
190  async fn exists(&self, key_id: &KeyId) -> KeyStorageResult<bool> {
191    let jwk_store: RwLockReadGuard<'_, JwkKeyStore> = self.jwk_store.read().await;
192    Ok(jwk_store.contains_key(key_id))
193  }
194}
195
196#[derive(Debug, Copy, Clone)]
197enum MemStoreKeyType {
198  Ed25519,
199  BLS12381G2,
200}
201
202impl JwkMemStore {
203  const ED25519_KEY_TYPE_STR: &'static str = "Ed25519";
204  /// The Ed25519 key type.
205  pub const ED25519_KEY_TYPE: KeyType = KeyType::from_static_str(Self::ED25519_KEY_TYPE_STR);
206
207  const BLS12381G2_KEY_TYPE_STR: &'static str = "BLS12381G2";
208  /// The BLS12381G2 key type
209  pub const BLS12381G2_KEY_TYPE: KeyType = KeyType::from_static_str(Self::BLS12381G2_KEY_TYPE_STR);
210
211  const PQ_KEY_TYPE_STR: &'static str = "AKP";
212  /// ML-DSA algorithms key types;
213  pub const PQ_KEY_TYPE: KeyType = KeyType::from_static_str(Self::PQ_KEY_TYPE_STR);
214}
215
216impl MemStoreKeyType {
217  const fn name(&self) -> &'static str {
218    match self {
219      MemStoreKeyType::Ed25519 => JwkMemStore::ED25519_KEY_TYPE_STR,
220      MemStoreKeyType::BLS12381G2 => JwkMemStore::BLS12381G2_KEY_TYPE_STR,
221    }
222  }
223}
224
225impl Display for MemStoreKeyType {
226  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
227    f.write_str(self.name())
228  }
229}
230
231impl TryFrom<&KeyType> for MemStoreKeyType {
232  type Error = KeyStorageError;
233
234  fn try_from(value: &KeyType) -> Result<Self, Self::Error> {
235    match value.as_str() {
236      JwkMemStore::ED25519_KEY_TYPE_STR => Ok(MemStoreKeyType::Ed25519),
237      JwkMemStore::BLS12381G2_KEY_TYPE_STR => Ok(MemStoreKeyType::BLS12381G2),
238      _ => Err(KeyStorageError::new(KeyStorageErrorKind::UnsupportedKeyType)),
239    }
240  }
241}
242
243impl TryFrom<&Jwk> for MemStoreKeyType {
244  type Error = KeyStorageError;
245
246  fn try_from(jwk: &Jwk) -> Result<Self, Self::Error> {
247    match jwk.kty() {
248      JwkType::Okp => {
249        let okp_params = jwk.try_okp_params().map_err(|err| {
250          KeyStorageError::new(KeyStorageErrorKind::UnsupportedKeyType)
251            .with_custom_message("expected Okp parameters for a JWK with `kty` Okp")
252            .with_source(err)
253        })?;
254        match okp_params.try_ed_curve().map_err(|err| {
255          KeyStorageError::new(KeyStorageErrorKind::UnsupportedKeyType)
256            .with_custom_message("only Ed curves are supported for signing")
257            .with_source(err)
258        })? {
259          EdCurve::Ed25519 => Ok(MemStoreKeyType::Ed25519),
260          curve => Err(
261            KeyStorageError::new(KeyStorageErrorKind::UnsupportedKeyType)
262              .with_custom_message(format!("{curve} not supported")),
263          ),
264        }
265      }
266      JwkType::Ec => {
267        let ec_params = jwk.try_ec_params().map_err(|err| {
268          KeyStorageError::new(KeyStorageErrorKind::UnsupportedKeyType)
269            .with_custom_message("expected EC parameters for a JWK with `kty` Ec")
270            .with_source(err)
271        })?;
272        match ec_params.try_bls_curve().map_err(|err| {
273          KeyStorageError::new(KeyStorageErrorKind::UnsupportedKeyType)
274            .with_custom_message("only Ed curves are supported for signing")
275            .with_source(err)
276        })? {
277          BlsCurve::BLS12381G2 => Ok(MemStoreKeyType::BLS12381G2),
278          curve => Err(
279            KeyStorageError::new(KeyStorageErrorKind::UnsupportedKeyType)
280              .with_custom_message(format!("{curve} not supported")),
281          ),
282        }
283      }
284      other => Err(
285        KeyStorageError::new(KeyStorageErrorKind::UnsupportedKeyType)
286          .with_custom_message(format!("Jwk `kty` {other} not supported")),
287      ),
288    }
289  }
290}
291
292impl Default for JwkMemStore {
293  fn default() -> Self {
294    Self::new()
295  }
296}
297
298/// Generate a random alphanumeric string of len 32.
299fn random_key_id() -> KeyId {
300  KeyId::new(rand::distributions::Alphanumeric.sample_string(&mut rand::thread_rng(), 32))
301}
302
303/// Check that the key type can be used with the algorithm.
304fn check_key_alg_compatibility(key_type: MemStoreKeyType, alg: &JwsAlgorithm) -> KeyStorageResult<()> {
305  match (key_type, alg) {
306    (MemStoreKeyType::Ed25519, JwsAlgorithm::EdDSA) => Ok(()),
307    (key_type, alg) => Err(
308      KeyStorageError::new(crate::key_storage::KeyStorageErrorKind::KeyAlgorithmMismatch)
309        .with_custom_message(format!("`cannot use key type `{key_type}` with algorithm `{alg}`")),
310    ),
311  }
312}
313
314#[cfg(feature = "pqc-liboqs")]
315mod pqc_liboqs {
316  use async_trait::async_trait;
317  use identity_verification::jose::jwk::Jwk;
318  use identity_verification::jose::jwk::JwkType;
319  use identity_verification::jose::jws::JwsAlgorithm;
320  use identity_verification::jwk::JwkParams;
321  use identity_verification::jwk::PostQuantumJwk;
322  use identity_verification::jwu;
323  use oqs::sig::Algorithm;
324  use oqs::sig::Sig;
325  use std::str::FromStr;
326  use tokio::sync::RwLockReadGuard;
327  use tokio::sync::RwLockWriteGuard;
328
329  use super::random_key_id;
330  use super::JwkKeyStore;
331  use super::JwkMemStore;
332  use super::KeyId;
333  use super::KeyStorageError;
334  use super::KeyStorageErrorKind;
335  use super::KeyStorageResult;
336  use super::KeyType;
337  use crate::key_storage::jwk_storage_pqc::JwkStoragePQ;
338  use crate::JwkGenOutput;
339
340  fn check_pq_alg_compatibility(alg: &JwsAlgorithm) -> KeyStorageResult<Algorithm> {
341    match alg {
342      JwsAlgorithm::ML_DSA_44 => Ok(Algorithm::MlDsa44),
343      JwsAlgorithm::ML_DSA_65 => Ok(Algorithm::MlDsa65),
344      JwsAlgorithm::ML_DSA_87 => Ok(Algorithm::MlDsa87),
345      JwsAlgorithm::SLH_DSA_SHA2_128s => Ok(Algorithm::SphincsSha2128sSimple),
346      JwsAlgorithm::SLH_DSA_SHAKE_128s => Ok(Algorithm::SphincsShake128sSimple),
347      JwsAlgorithm::SLH_DSA_SHA2_128f => Ok(Algorithm::SphincsSha2128fSimple),
348
349      JwsAlgorithm::SLH_DSA_SHAKE_128f => Ok(Algorithm::SphincsShake128fSimple),
350      JwsAlgorithm::SLH_DSA_SHA2_192s => Ok(Algorithm::SphincsSha2192sSimple),
351      JwsAlgorithm::SLH_DSA_SHAKE_192s => Ok(Algorithm::SphincsShake192sSimple),
352      JwsAlgorithm::SLH_DSA_SHA2_192f => Ok(Algorithm::SphincsSha2192fSimple),
353      JwsAlgorithm::SLH_DSA_SHAKE_192f => Ok(Algorithm::SphincsShake192fSimple),
354      JwsAlgorithm::SLH_DSA_SHA2_256s => Ok(Algorithm::SphincsSha2256sSimple),
355      JwsAlgorithm::SLH_DSA_SHAKE_256s => Ok(Algorithm::SphincsShake256sSimple),
356      JwsAlgorithm::SLH_DSA_SHA2_256f => Ok(Algorithm::SphincsSha2256fSimple),
357      JwsAlgorithm::SLH_DSA_SHAKE_256f => Ok(Algorithm::SphincsShake256fSimple),
358
359      JwsAlgorithm::FALCON512 => Ok(Algorithm::Falcon512),
360      JwsAlgorithm::FALCON1024 => Ok(Algorithm::Falcon1024),
361      other => Err(
362        KeyStorageError::new(KeyStorageErrorKind::UnsupportedSignatureAlgorithm)
363          .with_custom_message(format!("{other} is not supported")),
364      ),
365    }
366  }
367
368  /// JwkStoragePQ implementation for JwkMemStore based on liboqs.
369  #[cfg_attr(not(feature = "send-sync-storage"), async_trait(?Send))]
370  #[cfg_attr(feature = "send-sync-storage", async_trait)]
371  impl JwkStoragePQ for JwkMemStore {
372    async fn generate_pq_key(&self, key_type: KeyType, alg: JwsAlgorithm) -> KeyStorageResult<JwkGenOutput> {
373      if key_type != JwkMemStore::PQ_KEY_TYPE {
374        return Err(
375          KeyStorageError::new(KeyStorageErrorKind::UnsupportedKeyType)
376            .with_custom_message(format!("unsupported key type {key_type}")),
377        );
378      }
379
380      let oqs_alg = check_pq_alg_compatibility(&alg)?;
381      oqs::init();
382
383      let scheme = Sig::new(oqs_alg).map_err(|err| {
384        KeyStorageError::new(KeyStorageErrorKind::Unspecified)
385          .with_custom_message("signature scheme init failed".to_string())
386          .with_source(err)
387      })?;
388      let (pk, sk) = scheme.keypair().map_err(|err| {
389        KeyStorageError::new(KeyStorageErrorKind::Unspecified)
390          .with_custom_message("keypair generation failed!".to_string())
391          .with_source(err)
392      })?;
393
394      let kid: KeyId = random_key_id();
395
396      let public = jwu::encode_b64(pk.into_vec());
397      let private = jwu::encode_b64(sk.into_vec());
398
399      let mut jwk_params = match alg {
400        JwsAlgorithm::ML_DSA_44 => JwkParams::new(JwkType::Akp),
401        JwsAlgorithm::ML_DSA_65 => JwkParams::new(JwkType::Akp),
402        JwsAlgorithm::ML_DSA_87 => JwkParams::new(JwkType::Akp),
403        JwsAlgorithm::SLH_DSA_SHA2_128s => JwkParams::new(JwkType::Akp),
404        JwsAlgorithm::SLH_DSA_SHAKE_128s => JwkParams::new(JwkType::Akp),
405        JwsAlgorithm::SLH_DSA_SHA2_128f => JwkParams::new(JwkType::Akp),
406        JwsAlgorithm::SLH_DSA_SHAKE_128f => JwkParams::new(JwkType::Akp),
407        JwsAlgorithm::SLH_DSA_SHA2_192s => JwkParams::new(JwkType::Akp),
408        JwsAlgorithm::SLH_DSA_SHAKE_192s => JwkParams::new(JwkType::Akp),
409        JwsAlgorithm::SLH_DSA_SHA2_192f => JwkParams::new(JwkType::Akp),
410        JwsAlgorithm::SLH_DSA_SHAKE_192f => JwkParams::new(JwkType::Akp),
411        JwsAlgorithm::SLH_DSA_SHA2_256s => JwkParams::new(JwkType::Akp),
412        JwsAlgorithm::SLH_DSA_SHAKE_256s => JwkParams::new(JwkType::Akp),
413        JwsAlgorithm::SLH_DSA_SHA2_256f => JwkParams::new(JwkType::Akp),
414        JwsAlgorithm::SLH_DSA_SHAKE_256f => JwkParams::new(JwkType::Akp),
415        JwsAlgorithm::FALCON512 => JwkParams::new(JwkType::Akp),
416        JwsAlgorithm::FALCON1024 => JwkParams::new(JwkType::Akp),
417        other => {
418          return Err(
419            KeyStorageError::new(KeyStorageErrorKind::UnsupportedSignatureAlgorithm)
420              .with_custom_message(format!("{other} is not supported")),
421          );
422        }
423      };
424
425      match jwk_params {
426        JwkParams::Akp(ref mut params) => {
427          params.public = public;
428          params.private = Some(private);
429        }
430        _ => {
431          return Err(
432            KeyStorageError::new(KeyStorageErrorKind::UnsupportedKeyType).with_custom_message("Should NOT happen!"),
433          )
434        }
435      }
436
437      let mut jwk = Jwk::from_params(jwk_params);
438
439      jwk.set_alg(alg.name());
440      jwk.set_kid(jwk.thumbprint_sha256_b64());
441      let public_jwk: Jwk = jwk.to_public().expect("should only panic if kty == oct");
442
443      let mut jwk_store: RwLockWriteGuard<'_, JwkKeyStore> = self.jwk_store.write().await;
444      jwk_store.insert(kid.clone(), jwk);
445
446      Ok(JwkGenOutput::new(kid, public_jwk))
447    }
448
449    async fn pq_sign(
450      &self,
451      key_id: &KeyId,
452      data: &[u8],
453      public_key: &PostQuantumJwk,
454      ctx: Option<&[u8]>,
455    ) -> KeyStorageResult<Vec<u8>> {
456      let jwk_store: RwLockReadGuard<'_, JwkKeyStore> = self.jwk_store.read().await;
457
458      // Extract the required alg from the given public key
459      let alg = public_key
460        .alg()
461        .ok_or(KeyStorageErrorKind::UnsupportedSignatureAlgorithm)
462        .and_then(|alg_str| {
463          JwsAlgorithm::from_str(alg_str).map_err(|_| KeyStorageErrorKind::UnsupportedSignatureAlgorithm)
464        })?;
465
466      let oqs_alg = check_pq_alg_compatibility(&alg)?;
467
468      // Check that `kty` is `AKP`.
469      match alg {
470        JwsAlgorithm::ML_DSA_44
471        | JwsAlgorithm::ML_DSA_65
472        | JwsAlgorithm::ML_DSA_87
473        | JwsAlgorithm::SLH_DSA_SHA2_128s
474        | JwsAlgorithm::SLH_DSA_SHAKE_128s
475        | JwsAlgorithm::SLH_DSA_SHA2_128f
476        | JwsAlgorithm::SLH_DSA_SHAKE_128f
477        | JwsAlgorithm::SLH_DSA_SHA2_192s
478        | JwsAlgorithm::SLH_DSA_SHAKE_192s
479        | JwsAlgorithm::SLH_DSA_SHA2_192f
480        | JwsAlgorithm::SLH_DSA_SHAKE_192f
481        | JwsAlgorithm::SLH_DSA_SHA2_256s
482        | JwsAlgorithm::SLH_DSA_SHAKE_256s
483        | JwsAlgorithm::SLH_DSA_SHA2_256f
484        | JwsAlgorithm::SLH_DSA_SHAKE_256f
485        | JwsAlgorithm::FALCON512
486        | JwsAlgorithm::FALCON1024 => public_key.try_akp_params().map_err(|err| {
487          KeyStorageError::new(KeyStorageErrorKind::Unspecified)
488            .with_custom_message(format!("expected a Jwk with AKP params in order to sign with {alg}"))
489            .with_source(err)
490        })?,
491        other => {
492          return Err(
493            KeyStorageError::new(KeyStorageErrorKind::UnsupportedSignatureAlgorithm)
494              .with_custom_message(format!("{other} is not supported")),
495          );
496        }
497      };
498
499      // Obtain the corresponding private key and sign `data`.
500      let jwk: &Jwk = jwk_store
501        .get(key_id)
502        .ok_or_else(|| KeyStorageError::new(KeyStorageErrorKind::KeyNotFound))?;
503
504      let params = jwk.try_akp_params().unwrap();
505
506      let sk_bytes = params
507        .private
508        .as_deref()
509        .map(jwu::decode_b64)
510        .ok_or_else(|| {
511          KeyStorageError::new(KeyStorageErrorKind::Unspecified)
512            .with_custom_message("expected Jwk `private` param to be present")
513        })?
514        .map_err(|err| {
515          KeyStorageError::new(KeyStorageErrorKind::Unspecified)
516            .with_custom_message("unable to decode `private` param")
517            .with_source(err)
518        })?;
519      oqs::init();
520
521      let scheme = Sig::new(oqs_alg).map_err(|err| {
522        KeyStorageError::new(KeyStorageErrorKind::Unspecified)
523          .with_custom_message("signature scheme init failed".to_string())
524          .with_source(err)
525      })?;
526
527      let secret_key = scheme.secret_key_from_bytes(&sk_bytes).ok_or(
528        KeyStorageError::new(KeyStorageErrorKind::Unspecified).with_custom_message("invalid private key".to_string()),
529      )?;
530
531      let signature: oqs::sig::Signature;
532
533      if let Some(ctx) = ctx {
534        if !scheme.has_ctx_str_support() {
535          return Err(
536            KeyStorageError::new(KeyStorageErrorKind::Unspecified)
537              .with_custom_message("signature with ctx is not supported with this algorithm".to_string()),
538          );
539        }
540        signature = scheme.sign_with_ctx_str(data, ctx, secret_key).map_err(|err| {
541          KeyStorageError::new(KeyStorageErrorKind::Unspecified)
542            .with_custom_message("signature computation failed".to_string())
543            .with_source(err)
544        })?;
545      } else {
546        signature = scheme.sign(data, secret_key).map_err(|err| {
547          KeyStorageError::new(KeyStorageErrorKind::Unspecified)
548            .with_custom_message("signature computation failed".to_string())
549            .with_source(err)
550        })?;
551      }
552
553      Ok(signature.into_vec())
554    }
555  }
556}
557
558#[cfg(feature = "jpt-bbs-plus")]
559mod bbs_plus_impl {
560  use std::str::FromStr as _;
561
562  use crate::key_storage::bls::encode_bls_jwk;
563  use crate::key_storage::bls::expand_bls_jwk;
564  use crate::key_storage::bls::generate_bbs_keypair;
565  use crate::key_storage::bls::sign_bbs;
566  use crate::key_storage::bls::update_bbs_signature;
567  use crate::JwkGenOutput;
568  use crate::JwkMemStore;
569  use crate::JwkStorageBbsPlusExt;
570  use crate::KeyId;
571  use crate::KeyStorageError;
572  use crate::KeyStorageErrorKind;
573  use crate::KeyStorageResult;
574  use crate::KeyType;
575  use crate::ProofUpdateCtx;
576  use async_trait::async_trait;
577  use identity_verification::jwk::BlsCurve;
578  use identity_verification::jwk::Jwk;
579  use jsonprooftoken::jpa::algs::ProofAlgorithm;
580
581  use super::random_key_id;
582
583  /// JwkStorageBbsPlusExt implementation for JwkMemStore
584  #[cfg_attr(not(feature = "send-sync-storage"), async_trait(?Send))]
585  #[cfg_attr(feature = "send-sync-storage", async_trait)]
586  impl JwkStorageBbsPlusExt for JwkMemStore {
587    async fn generate_bbs(&self, key_type: KeyType, alg: ProofAlgorithm) -> KeyStorageResult<JwkGenOutput> {
588      if key_type != JwkMemStore::BLS12381G2_KEY_TYPE {
589        return Err(
590          KeyStorageError::new(KeyStorageErrorKind::UnsupportedKeyType)
591            .with_custom_message(format!("unsupported key type {key_type}")),
592        );
593      }
594
595      let (private_key, public_key) = generate_bbs_keypair(alg)?;
596      let (jwk, public_jwk) = encode_bls_jwk(&private_key, &public_key, alg);
597
598      let kid: KeyId = random_key_id();
599      let mut jwk_store = self.jwk_store.write().await;
600      jwk_store.insert(kid.clone(), jwk);
601
602      Ok(JwkGenOutput::new(kid, public_jwk))
603    }
604
605    async fn sign_bbs(
606      &self,
607      key_id: &KeyId,
608      data: &[Vec<u8>],
609      header: &[u8],
610      public_key: &Jwk,
611    ) -> KeyStorageResult<Vec<u8>> {
612      let jwk_store = self.jwk_store.read().await;
613
614      // Extract the required alg from the given public key
615      let alg = public_key
616        .alg()
617        .and_then(|alg_str| ProofAlgorithm::from_str(alg_str).ok())
618        .ok_or(KeyStorageErrorKind::UnsupportedProofAlgorithm)?;
619
620      // Check the provided JWK represents a BLS12381G2 key.
621      if !public_key
622        .try_ec_params()
623        .map(|ec| ec.crv == BlsCurve::BLS12381G2.to_string())
624        .unwrap_or(false)
625      {
626        return Err(
627          KeyStorageError::new(KeyStorageErrorKind::UnsupportedKeyType)
628            .with_custom_message(format!("expected a key from the {} curve", BlsCurve::BLS12381G2)),
629        );
630      }
631
632      // Obtain the corresponding private key.
633      let jwk: &Jwk = jwk_store.get(key_id).ok_or(KeyStorageErrorKind::KeyNotFound)?;
634      let (sk, pk) = expand_bls_jwk(jwk)?;
635
636      sign_bbs(alg, data, &sk.expect("jwk is private"), &pk, header)
637    }
638
639    async fn update_signature(
640      &self,
641      key_id: &KeyId,
642      public_key: &Jwk,
643      signature: &[u8],
644      ctx: ProofUpdateCtx,
645    ) -> KeyStorageResult<Vec<u8>> {
646      let jwk_store = self.jwk_store.read().await;
647
648      // Extract the required alg from the given public key
649      let alg = public_key
650        .alg()
651        .ok_or(KeyStorageErrorKind::UnsupportedProofAlgorithm)
652        .and_then(|alg_str| {
653          ProofAlgorithm::from_str(alg_str).map_err(|_| KeyStorageErrorKind::UnsupportedProofAlgorithm)
654        })?;
655
656      // Check the provided JWK represents a BLS12381G2 key.
657      if !public_key
658        .try_ec_params()
659        .map(|ec| ec.crv == BlsCurve::BLS12381G2.to_string())
660        .unwrap_or(false)
661      {
662        return Err(
663          KeyStorageError::new(KeyStorageErrorKind::UnsupportedKeyType)
664            .with_custom_message(format!("expected a key from the {} curve", BlsCurve::BLS12381G2)),
665        );
666      }
667
668      // Obtain the corresponding private key.
669      let jwk = jwk_store.get(key_id).ok_or(KeyStorageErrorKind::KeyNotFound)?;
670      let sk = expand_bls_jwk(jwk)?.0.expect("jwk is private");
671
672      // Update the signature.
673      update_bbs_signature(alg, signature, &sk, &ctx)
674    }
675  }
676}
677pub(crate) mod shared {
678  use core::fmt::Debug;
679  use core::fmt::Formatter;
680  use tokio::sync::RwLock;
681  use tokio::sync::RwLockReadGuard;
682  use tokio::sync::RwLockWriteGuard;
683
684  #[derive(Default)]
685  pub(crate) struct Shared<T>(RwLock<T>);
686
687  impl<T> Shared<T> {
688    pub(crate) fn new(data: T) -> Self {
689      Self(RwLock::new(data))
690    }
691
692    pub(crate) async fn read(&self) -> RwLockReadGuard<'_, T> {
693      self.0.read().await
694    }
695
696    pub(crate) async fn write(&self) -> RwLockWriteGuard<'_, T> {
697      self.0.write().await
698    }
699  }
700
701  impl<T: Debug> Debug for Shared<T> {
702    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
703      Debug::fmt(&self.0, f)
704    }
705  }
706}
707
708#[cfg(test)]
709mod tests {
710  use crate::key_storage::tests::utils::generate_ed25519;
711  use fastcrypto::ed25519::Ed25519PublicKey;
712  use fastcrypto::traits::ToFromBytes as _;
713  use fastcrypto::traits::VerifyingKey as _;
714  use identity_verification::jose::jwk::EcCurve;
715  use identity_verification::jose::jwk::JwkParamsEc;
716  use identity_verification::jwk::FromJwk as _;
717
718  use super::*;
719
720  #[tokio::test]
721  async fn generate_and_sign() {
722    let test_msg: &[u8] = b"test";
723    let store: JwkMemStore = JwkMemStore::new();
724
725    let JwkGenOutput { key_id, jwk } = store
726      .generate(JwkMemStore::ED25519_KEY_TYPE, JwsAlgorithm::EdDSA)
727      .await
728      .unwrap();
729
730    let signature = store.sign(&key_id, test_msg, &jwk.to_public().unwrap()).await.unwrap();
731
732    let public_key = Ed25519PublicKey::from_jwk(&jwk).unwrap();
733    let signature = Ed25519Signature::from_bytes(&signature).unwrap();
734
735    assert!(public_key.verify(test_msg, &signature).is_ok());
736    assert!(store.exists(&key_id).await.unwrap());
737    store.delete(&key_id).await.unwrap();
738  }
739
740  #[tokio::test]
741  async fn insert() {
742    let store: JwkMemStore = JwkMemStore::new();
743
744    let key_pair = generate_ed25519();
745    let mut jwk: Jwk = key_pair.to_jwk().unwrap();
746
747    // INVALID: Inserting a Jwk without an `alg` parameter should fail.
748    let err = store.insert(jwk.clone()).await.unwrap_err();
749    assert!(matches!(err.kind(), KeyStorageErrorKind::UnsupportedSignatureAlgorithm));
750
751    // VALID: Inserting a Jwk with all private key components set should succeed.
752    jwk.set_alg(JwsAlgorithm::EdDSA.name());
753    store.insert(jwk.clone()).await.unwrap();
754
755    // INVALID: Inserting a Jwk with all private key components unset should fail.
756    let err = store.insert(jwk.to_public().unwrap()).await.unwrap_err();
757    assert!(matches!(err.kind(), KeyStorageErrorKind::Unspecified))
758  }
759
760  #[tokio::test]
761  async fn exists() {
762    let store: JwkMemStore = JwkMemStore::new();
763    assert!(!store.exists(&KeyId::new("non-existent-id")).await.unwrap());
764  }
765
766  #[tokio::test]
767  async fn incompatible_key_type() {
768    let store: JwkMemStore = JwkMemStore::new();
769
770    let mut ec_params = JwkParamsEc::new();
771    ec_params.crv = EcCurve::P256.name().to_string();
772    ec_params.x = String::new();
773    ec_params.y = String::new();
774    ec_params.d = Some(String::new());
775    let jwk_ec = Jwk::from_params(ec_params);
776
777    let err = store.insert(jwk_ec).await.unwrap_err();
778    assert!(matches!(err.kind(), KeyStorageErrorKind::UnsupportedKeyType));
779  }
780
781  #[tokio::test]
782  async fn incompatible_key_alg() {
783    let store: JwkMemStore = JwkMemStore::new();
784
785    let key_pair = generate_ed25519();
786    let mut jwk: Jwk = key_pair.to_jwk().unwrap();
787    jwk.set_alg(JwsAlgorithm::ES256.name());
788
789    // INVALID: Inserting an Ed25519 key with the ES256 alg is not compatible.
790    let err = store.insert(jwk.clone()).await.unwrap_err();
791    assert!(matches!(err.kind(), KeyStorageErrorKind::KeyAlgorithmMismatch));
792  }
793}