jwt_simple/algorithms/
mldsa.rs

1use std::convert::TryInto;
2
3use ct_codecs::{Base64UrlSafeNoPadding, Encoder};
4use serde::{de::DeserializeOwned, Serialize};
5use superboring::mldsa::{Algorithm, MlDsaPrivateKey, MlDsaPublicKey};
6
7use crate::algorithms::jwk::{
8    decode_fixed, decode_secret, encode, encode_secret, is_encoding_of, thumbprint, Jwk,
9    JwkDescriptor, JwkMembers, KeyRole, KeyType,
10};
11use crate::claims::*;
12use crate::common::*;
13#[cfg(feature = "cwt")]
14use crate::cwt_token::*;
15use crate::error::*;
16use crate::jwt_header::*;
17use crate::token::*;
18
19fn mldsa_jwk_descriptor(alg: &'static str, role: KeyRole) -> JwkDescriptor {
20    JwkDescriptor::new(KeyType::Akp, None, alg, role)
21}
22
23#[doc(hidden)]
24#[derive(Debug, Clone)]
25pub struct MLDSAPublicKey(MlDsaPublicKey);
26
27impl AsRef<MlDsaPublicKey> for MLDSAPublicKey {
28    fn as_ref(&self) -> &MlDsaPublicKey {
29        &self.0
30    }
31}
32
33impl MLDSAPublicKey {
34    pub fn from_bytes(algorithm: Algorithm, raw: &[u8]) -> Result<Self, Error> {
35        let mldsa_pk = MlDsaPublicKey::from_slice(algorithm, raw);
36        Ok(MLDSAPublicKey(
37            mldsa_pk.map_err(|_| JWTError::InvalidPublicKey)?,
38        ))
39    }
40
41    pub fn to_bytes(&self) -> Vec<u8> {
42        self.0.to_bytes().expect("failed to serialize public key")
43    }
44
45    pub(crate) fn from_jwk(
46        jwk: &str,
47        alg: &'static str,
48        algorithm: Algorithm,
49    ) -> Result<(Self, Option<String>), Error> {
50        let jwk = Jwk::parse(jwk, &mldsa_jwk_descriptor(alg, KeyRole::SignaturePublic))?;
51        let mut raw = vec![0u8; algorithm.public_key_bytes()];
52        ensure!(
53            decode_fixed(jwk.public(), &mut raw),
54            JWTError::InvalidPublicKey
55        );
56        Ok((
57            Self::from_bytes(algorithm, &raw)?,
58            jwk.kid().map(Into::into),
59        ))
60    }
61
62    pub(crate) fn to_jwk(&self, alg: &'static str, key_id: Option<&str>) -> String {
63        let public = encode(&self.to_bytes());
64        let members = JwkMembers {
65            public: Some(&public),
66            ..Default::default()
67        };
68        mldsa_jwk_descriptor(alg, KeyRole::SignaturePublic).export(members, key_id)
69    }
70
71    pub(crate) fn jwk_thumbprint(&self, alg: &str) -> String {
72        thumbprint(&[
73            ("alg", alg),
74            ("kty", "AKP"),
75            ("pub", &encode(&self.to_bytes())),
76        ])
77    }
78}
79
80#[doc(hidden)]
81#[derive(Clone)]
82pub struct MLDSAKeyPair {
83    mldsa_sk: MlDsaPrivateKey,
84    metadata: Option<KeyMetadata>,
85}
86
87impl std::fmt::Debug for MLDSAKeyPair {
88    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89        f.debug_struct("PKey")
90            .field("algorithm", &"ML-DSA")
91            .finish()
92    }
93}
94
95impl AsRef<MlDsaPrivateKey> for MLDSAKeyPair {
96    fn as_ref(&self) -> &MlDsaPrivateKey {
97        &self.mldsa_sk
98    }
99}
100
101impl MLDSAKeyPair {
102    /// The raw representation of an ML-DSA key pair is its 32-byte seed, as
103    /// mandated for JOSE and COSE.
104    pub fn from_bytes(algorithm: Algorithm, raw: &[u8]) -> Result<Self, Error> {
105        let seed = raw.try_into().map_err(|_| JWTError::InvalidKeyPair)?;
106        let mldsa_sk =
107            MlDsaPrivateKey::from_seed(algorithm, &seed).map_err(|_| JWTError::InvalidKeyPair)?;
108        Ok(MLDSAKeyPair {
109            mldsa_sk,
110            metadata: None,
111        })
112    }
113
114    pub fn to_bytes(&self) -> Vec<u8> {
115        self.mldsa_sk.seed_bytes().to_vec()
116    }
117
118    pub fn public_key(&self) -> MLDSAPublicKey {
119        let mldsa_pk = self
120            .mldsa_sk
121            .public_key()
122            .expect("failed to create public key");
123        MLDSAPublicKey(mldsa_pk)
124    }
125
126    pub fn generate(algorithm: Algorithm) -> Self {
127        let (_, mldsa_sk) =
128            MlDsaPrivateKey::generate(algorithm).expect("failed to generate key pair");
129        MLDSAKeyPair {
130            mldsa_sk,
131            metadata: None,
132        }
133    }
134
135    pub(crate) fn from_jwk(
136        jwk: &str,
137        alg: &'static str,
138        algorithm: Algorithm,
139    ) -> Result<(Self, Option<String>), Error> {
140        let jwk = Jwk::parse(jwk, &mldsa_jwk_descriptor(alg, KeyRole::SignaturePrivate))?;
141        let seed = decode_secret(jwk.private(), 32).ok_or(JWTError::InvalidKeyPair)?;
142        let key_pair = Self::from_bytes(algorithm, &seed)?;
143        ensure!(
144            is_encoding_of(jwk.public(), &key_pair.public_key().to_bytes()),
145            JWTError::InvalidKeyPair
146        );
147        Ok((key_pair, jwk.kid().map(Into::into)))
148    }
149
150    pub(crate) fn to_jwk(&self, alg: &'static str, key_id: Option<&str>) -> String {
151        let public = encode(&self.public_key().to_bytes());
152        let private = encode_secret(self.mldsa_sk.seed_bytes());
153        let members = JwkMembers {
154            public: Some(&public),
155            private: Some(&private),
156            ..Default::default()
157        };
158        mldsa_jwk_descriptor(alg, KeyRole::SignaturePrivate).export(members, key_id)
159    }
160}
161
162pub trait MLDSAKeyPairLike {
163    fn jwt_alg_name() -> &'static str;
164    fn key_pair(&self) -> &MLDSAKeyPair;
165    fn key_id(&self) -> &Option<String>;
166    fn metadata(&self) -> &Option<KeyMetadata>;
167    fn attach_metadata(&mut self, metadata: KeyMetadata) -> Result<(), Error>;
168
169    fn sign<CustomClaims: Serialize>(
170        &self,
171        claims: JWTClaims<CustomClaims>,
172    ) -> Result<String, Error> {
173        self.sign_with_options(claims, &Default::default())
174    }
175
176    fn sign_with_options<CustomClaims: Serialize>(
177        &self,
178        claims: JWTClaims<CustomClaims>,
179        opts: &HeaderOptions,
180    ) -> Result<String, Error> {
181        let jwt_header = JWTHeader::new(Self::jwt_alg_name().to_string(), self.key_id().clone())
182            .with_key_metadata(self.metadata())
183            .with_options(opts);
184        Token::build(&jwt_header, claims, |authenticated| {
185            let signature = self.key_pair().as_ref().sign(authenticated.as_bytes())?;
186            Ok(signature)
187        })
188    }
189}
190
191pub trait MLDSAPublicKeyLike {
192    fn jwt_alg_name() -> &'static str;
193    fn public_key(&self) -> &MLDSAPublicKey;
194    fn key_id(&self) -> &Option<String>;
195    fn set_key_id(&mut self, key_id: String);
196
197    fn verify_token<CustomClaims: DeserializeOwned>(
198        &self,
199        token: &str,
200        options: Option<VerificationOptions>,
201    ) -> Result<JWTClaims<CustomClaims>, Error> {
202        Token::verify(
203            Self::jwt_alg_name(),
204            token,
205            options,
206            |authenticated, signature| {
207                self.public_key()
208                    .as_ref()
209                    .verify(authenticated.as_bytes(), signature)
210                    .map_err(|_| JWTError::InvalidSignature)?;
211                Ok(())
212            },
213            |_salt: Option<&[u8]>| Ok(()),
214        )
215    }
216
217    #[cfg(feature = "cwt")]
218    fn verify_cwt_token<CustomClaims: DeserializeOwned>(
219        &self,
220        token: &[u8],
221        options: Option<VerificationOptions>,
222    ) -> Result<JWTClaims<NoCustomClaims>, Error> {
223        CWTToken::verify(
224            Self::jwt_alg_name(),
225            token,
226            options,
227            |authenticated, signature| {
228                self.public_key()
229                    .as_ref()
230                    .verify(authenticated.as_bytes(), signature)
231                    .map_err(|_| JWTError::InvalidSignature)?;
232                Ok(())
233            },
234        )
235    }
236
237    /// Decode CWT token metadata that can be useful prior to signature/tag verification
238    #[cfg(feature = "cwt")]
239    fn decode_cwt_metadata(&self, token: impl AsRef<[u8]>) -> Result<TokenMetadata, Error> {
240        CWTToken::decode_metadata(token)
241    }
242
243    fn create_key_id(&mut self) -> &str {
244        self.set_key_id(
245            Base64UrlSafeNoPadding::encode_to_string(hmac_sha256::Hash::hash(
246                &self.public_key().to_bytes(),
247            ))
248            .unwrap(),
249        );
250        self.key_id().as_ref().map(|x| x.as_str()).unwrap()
251    }
252}
253
254#[derive(Clone)]
255pub struct MLDSA44KeyPair {
256    key_pair: MLDSAKeyPair,
257    key_id: Option<String>,
258}
259
260impl std::fmt::Debug for MLDSA44KeyPair {
261    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
262        f.debug_struct("PKey")
263            .field("algorithm", &"ML-DSA-44")
264            .finish()
265    }
266}
267
268#[derive(Debug, Clone)]
269pub struct MLDSA44PublicKey {
270    pk: MLDSAPublicKey,
271    key_id: Option<String>,
272}
273
274impl MLDSAKeyPairLike for MLDSA44KeyPair {
275    fn jwt_alg_name() -> &'static str {
276        "ML-DSA-44"
277    }
278
279    fn key_pair(&self) -> &MLDSAKeyPair {
280        &self.key_pair
281    }
282
283    fn key_id(&self) -> &Option<String> {
284        &self.key_id
285    }
286
287    fn metadata(&self) -> &Option<KeyMetadata> {
288        &self.key_pair.metadata
289    }
290
291    fn attach_metadata(&mut self, metadata: KeyMetadata) -> Result<(), Error> {
292        self.key_pair.metadata = Some(metadata);
293        Ok(())
294    }
295}
296
297impl MLDSA44KeyPair {
298    pub fn from_bytes(raw: &[u8]) -> Result<Self, Error> {
299        Ok(MLDSA44KeyPair {
300            key_pair: MLDSAKeyPair::from_bytes(Algorithm::MlDsa44, raw)?,
301            key_id: None,
302        })
303    }
304
305    pub fn to_bytes(&self) -> Vec<u8> {
306        self.key_pair.to_bytes()
307    }
308
309    pub fn public_key(&self) -> MLDSA44PublicKey {
310        MLDSA44PublicKey {
311            pk: self.key_pair.public_key(),
312            key_id: self.key_id.clone(),
313        }
314    }
315
316    pub fn generate() -> Self {
317        MLDSA44KeyPair {
318            key_pair: MLDSAKeyPair::generate(Algorithm::MlDsa44),
319            key_id: None,
320        }
321    }
322
323    pub fn with_key_id(mut self, key_id: &str) -> Self {
324        self.key_id = Some(key_id.to_string());
325        self
326    }
327
328    /// Import a key pair from a private AKP JWK, with `alg: ML-DSA-44`.
329    ///
330    /// `priv` must contain the 32-byte seed used by `from_bytes()`.
331    /// See [strict JWK validation](crate#strict-jwk-validation) for the other rules.
332    pub fn from_jwk(jwk: &str) -> Result<Self, Error> {
333        let (key_pair, key_id) =
334            MLDSAKeyPair::from_jwk(jwk, Self::jwt_alg_name(), Algorithm::MlDsa44)?;
335        Ok(MLDSA44KeyPair { key_pair, key_id })
336    }
337
338    /// Export the key pair as a JWK, private key included.
339    ///
340    /// Use `public_key().to_jwk()` to share the public key.
341    pub fn to_jwk(&self) -> String {
342        self.key_pair
343            .to_jwk(Self::jwt_alg_name(), self.key_id.as_deref())
344    }
345
346    /// The JWK thumbprint of the public key, which includes its algorithm.
347    pub fn jwk_thumbprint(&self) -> String {
348        self.key_pair
349            .public_key()
350            .jwk_thumbprint(Self::jwt_alg_name())
351    }
352}
353
354impl MLDSAPublicKeyLike for MLDSA44PublicKey {
355    fn jwt_alg_name() -> &'static str {
356        "ML-DSA-44"
357    }
358
359    fn public_key(&self) -> &MLDSAPublicKey {
360        &self.pk
361    }
362
363    fn key_id(&self) -> &Option<String> {
364        &self.key_id
365    }
366
367    fn set_key_id(&mut self, key_id: String) {
368        self.key_id = Some(key_id);
369    }
370}
371
372impl MLDSA44PublicKey {
373    pub fn from_bytes(raw: &[u8]) -> Result<Self, Error> {
374        Ok(MLDSA44PublicKey {
375            pk: MLDSAPublicKey::from_bytes(Algorithm::MlDsa44, raw)?,
376            key_id: None,
377        })
378    }
379
380    pub fn to_bytes(&self) -> Vec<u8> {
381        self.pk.to_bytes()
382    }
383
384    pub fn with_key_id(mut self, key_id: &str) -> Self {
385        self.key_id = Some(key_id.to_string());
386        self
387    }
388
389    /// Import a public key from an AKP JWK, with `alg: ML-DSA-44`.
390    ///
391    /// See [strict JWK validation](crate#strict-jwk-validation) for the rules.
392    pub fn from_jwk(jwk: &str) -> Result<Self, Error> {
393        let (pk, key_id) = MLDSAPublicKey::from_jwk(jwk, Self::jwt_alg_name(), Algorithm::MlDsa44)?;
394        Ok(MLDSA44PublicKey { pk, key_id })
395    }
396
397    /// Export the public key as an AKP JWK, with `alg: ML-DSA-44` and `use: sig`.
398    pub fn to_jwk(&self) -> String {
399        self.pk.to_jwk(Self::jwt_alg_name(), self.key_id.as_deref())
400    }
401
402    /// The JWK thumbprint of the key, which includes its algorithm.
403    pub fn jwk_thumbprint(&self) -> String {
404        self.pk.jwk_thumbprint(Self::jwt_alg_name())
405    }
406}
407
408//
409
410#[derive(Clone)]
411pub struct MLDSA65KeyPair {
412    key_pair: MLDSAKeyPair,
413    key_id: Option<String>,
414}
415
416impl std::fmt::Debug for MLDSA65KeyPair {
417    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
418        f.debug_struct("PKey")
419            .field("algorithm", &"ML-DSA-65")
420            .finish()
421    }
422}
423
424#[derive(Debug, Clone)]
425pub struct MLDSA65PublicKey {
426    pk: MLDSAPublicKey,
427    key_id: Option<String>,
428}
429
430impl MLDSAKeyPairLike for MLDSA65KeyPair {
431    fn jwt_alg_name() -> &'static str {
432        "ML-DSA-65"
433    }
434
435    fn key_pair(&self) -> &MLDSAKeyPair {
436        &self.key_pair
437    }
438
439    fn key_id(&self) -> &Option<String> {
440        &self.key_id
441    }
442
443    fn metadata(&self) -> &Option<KeyMetadata> {
444        &self.key_pair.metadata
445    }
446
447    fn attach_metadata(&mut self, metadata: KeyMetadata) -> Result<(), Error> {
448        self.key_pair.metadata = Some(metadata);
449        Ok(())
450    }
451}
452
453impl MLDSA65KeyPair {
454    pub fn from_bytes(raw: &[u8]) -> Result<Self, Error> {
455        Ok(MLDSA65KeyPair {
456            key_pair: MLDSAKeyPair::from_bytes(Algorithm::MlDsa65, raw)?,
457            key_id: None,
458        })
459    }
460
461    pub fn to_bytes(&self) -> Vec<u8> {
462        self.key_pair.to_bytes()
463    }
464
465    pub fn public_key(&self) -> MLDSA65PublicKey {
466        MLDSA65PublicKey {
467            pk: self.key_pair.public_key(),
468            key_id: self.key_id.clone(),
469        }
470    }
471
472    pub fn generate() -> Self {
473        MLDSA65KeyPair {
474            key_pair: MLDSAKeyPair::generate(Algorithm::MlDsa65),
475            key_id: None,
476        }
477    }
478
479    pub fn with_key_id(mut self, key_id: &str) -> Self {
480        self.key_id = Some(key_id.to_string());
481        self
482    }
483
484    /// Import a key pair from a private AKP JWK, with `alg: ML-DSA-65`.
485    ///
486    /// `priv` must contain the 32-byte seed used by `from_bytes()`.
487    /// See [strict JWK validation](crate#strict-jwk-validation) for the other rules.
488    pub fn from_jwk(jwk: &str) -> Result<Self, Error> {
489        let (key_pair, key_id) =
490            MLDSAKeyPair::from_jwk(jwk, Self::jwt_alg_name(), Algorithm::MlDsa65)?;
491        Ok(MLDSA65KeyPair { key_pair, key_id })
492    }
493
494    /// Export the key pair as a JWK, private key included.
495    ///
496    /// Use `public_key().to_jwk()` to share the public key.
497    pub fn to_jwk(&self) -> String {
498        self.key_pair
499            .to_jwk(Self::jwt_alg_name(), self.key_id.as_deref())
500    }
501
502    /// The JWK thumbprint of the public key, which includes its algorithm.
503    pub fn jwk_thumbprint(&self) -> String {
504        self.key_pair
505            .public_key()
506            .jwk_thumbprint(Self::jwt_alg_name())
507    }
508}
509
510impl MLDSAPublicKeyLike for MLDSA65PublicKey {
511    fn jwt_alg_name() -> &'static str {
512        "ML-DSA-65"
513    }
514
515    fn public_key(&self) -> &MLDSAPublicKey {
516        &self.pk
517    }
518
519    fn key_id(&self) -> &Option<String> {
520        &self.key_id
521    }
522
523    fn set_key_id(&mut self, key_id: String) {
524        self.key_id = Some(key_id);
525    }
526}
527
528impl MLDSA65PublicKey {
529    pub fn from_bytes(raw: &[u8]) -> Result<Self, Error> {
530        Ok(MLDSA65PublicKey {
531            pk: MLDSAPublicKey::from_bytes(Algorithm::MlDsa65, raw)?,
532            key_id: None,
533        })
534    }
535
536    pub fn to_bytes(&self) -> Vec<u8> {
537        self.pk.to_bytes()
538    }
539
540    pub fn with_key_id(mut self, key_id: &str) -> Self {
541        self.key_id = Some(key_id.to_string());
542        self
543    }
544
545    /// Import a public key from an AKP JWK, with `alg: ML-DSA-65`.
546    ///
547    /// See [strict JWK validation](crate#strict-jwk-validation) for the rules.
548    pub fn from_jwk(jwk: &str) -> Result<Self, Error> {
549        let (pk, key_id) = MLDSAPublicKey::from_jwk(jwk, Self::jwt_alg_name(), Algorithm::MlDsa65)?;
550        Ok(MLDSA65PublicKey { pk, key_id })
551    }
552
553    /// Export the public key as an AKP JWK, with `alg: ML-DSA-65` and `use: sig`.
554    pub fn to_jwk(&self) -> String {
555        self.pk.to_jwk(Self::jwt_alg_name(), self.key_id.as_deref())
556    }
557
558    /// The JWK thumbprint of the key, which includes its algorithm.
559    pub fn jwk_thumbprint(&self) -> String {
560        self.pk.jwk_thumbprint(Self::jwt_alg_name())
561    }
562}
563
564//
565
566#[derive(Clone)]
567pub struct MLDSA87KeyPair {
568    key_pair: MLDSAKeyPair,
569    key_id: Option<String>,
570}
571
572impl std::fmt::Debug for MLDSA87KeyPair {
573    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
574        f.debug_struct("PKey")
575            .field("algorithm", &"ML-DSA-87")
576            .finish()
577    }
578}
579
580#[derive(Debug, Clone)]
581pub struct MLDSA87PublicKey {
582    pk: MLDSAPublicKey,
583    key_id: Option<String>,
584}
585
586impl MLDSAKeyPairLike for MLDSA87KeyPair {
587    fn jwt_alg_name() -> &'static str {
588        "ML-DSA-87"
589    }
590
591    fn key_pair(&self) -> &MLDSAKeyPair {
592        &self.key_pair
593    }
594
595    fn key_id(&self) -> &Option<String> {
596        &self.key_id
597    }
598
599    fn metadata(&self) -> &Option<KeyMetadata> {
600        &self.key_pair.metadata
601    }
602
603    fn attach_metadata(&mut self, metadata: KeyMetadata) -> Result<(), Error> {
604        self.key_pair.metadata = Some(metadata);
605        Ok(())
606    }
607}
608
609impl MLDSA87KeyPair {
610    pub fn from_bytes(raw: &[u8]) -> Result<Self, Error> {
611        Ok(MLDSA87KeyPair {
612            key_pair: MLDSAKeyPair::from_bytes(Algorithm::MlDsa87, raw)?,
613            key_id: None,
614        })
615    }
616
617    pub fn to_bytes(&self) -> Vec<u8> {
618        self.key_pair.to_bytes()
619    }
620
621    pub fn public_key(&self) -> MLDSA87PublicKey {
622        MLDSA87PublicKey {
623            pk: self.key_pair.public_key(),
624            key_id: self.key_id.clone(),
625        }
626    }
627
628    pub fn generate() -> Self {
629        MLDSA87KeyPair {
630            key_pair: MLDSAKeyPair::generate(Algorithm::MlDsa87),
631            key_id: None,
632        }
633    }
634
635    pub fn with_key_id(mut self, key_id: &str) -> Self {
636        self.key_id = Some(key_id.to_string());
637        self
638    }
639
640    /// Import a key pair from a private AKP JWK, with `alg: ML-DSA-87`.
641    ///
642    /// `priv` must contain the 32-byte seed used by `from_bytes()`.
643    /// See [strict JWK validation](crate#strict-jwk-validation) for the other rules.
644    pub fn from_jwk(jwk: &str) -> Result<Self, Error> {
645        let (key_pair, key_id) =
646            MLDSAKeyPair::from_jwk(jwk, Self::jwt_alg_name(), Algorithm::MlDsa87)?;
647        Ok(MLDSA87KeyPair { key_pair, key_id })
648    }
649
650    /// Export the key pair as a JWK, private key included.
651    ///
652    /// Use `public_key().to_jwk()` to share the public key.
653    pub fn to_jwk(&self) -> String {
654        self.key_pair
655            .to_jwk(Self::jwt_alg_name(), self.key_id.as_deref())
656    }
657
658    /// The JWK thumbprint of the public key, which includes its algorithm.
659    pub fn jwk_thumbprint(&self) -> String {
660        self.key_pair
661            .public_key()
662            .jwk_thumbprint(Self::jwt_alg_name())
663    }
664}
665
666impl MLDSAPublicKeyLike for MLDSA87PublicKey {
667    fn jwt_alg_name() -> &'static str {
668        "ML-DSA-87"
669    }
670
671    fn public_key(&self) -> &MLDSAPublicKey {
672        &self.pk
673    }
674
675    fn key_id(&self) -> &Option<String> {
676        &self.key_id
677    }
678
679    fn set_key_id(&mut self, key_id: String) {
680        self.key_id = Some(key_id);
681    }
682}
683
684impl MLDSA87PublicKey {
685    pub fn from_bytes(raw: &[u8]) -> Result<Self, Error> {
686        Ok(MLDSA87PublicKey {
687            pk: MLDSAPublicKey::from_bytes(Algorithm::MlDsa87, raw)?,
688            key_id: None,
689        })
690    }
691
692    pub fn to_bytes(&self) -> Vec<u8> {
693        self.pk.to_bytes()
694    }
695
696    pub fn with_key_id(mut self, key_id: &str) -> Self {
697        self.key_id = Some(key_id.to_string());
698        self
699    }
700
701    /// Import a public key from an AKP JWK, with `alg: ML-DSA-87`.
702    ///
703    /// See [strict JWK validation](crate#strict-jwk-validation) for the rules.
704    pub fn from_jwk(jwk: &str) -> Result<Self, Error> {
705        let (pk, key_id) = MLDSAPublicKey::from_jwk(jwk, Self::jwt_alg_name(), Algorithm::MlDsa87)?;
706        Ok(MLDSA87PublicKey { pk, key_id })
707    }
708
709    /// Export the public key as an AKP JWK, with `alg: ML-DSA-87` and `use: sig`.
710    pub fn to_jwk(&self) -> String {
711        self.pk.to_jwk(Self::jwt_alg_name(), self.key_id.as_deref())
712    }
713
714    /// The JWK thumbprint of the key, which includes its algorithm.
715    pub fn jwk_thumbprint(&self) -> String {
716        self.pk.jwk_thumbprint(Self::jwt_alg_name())
717    }
718}
719
720#[cfg(test)]
721mod jwk_tests {
722    use super::*;
723    use crate::algorithms::jwk_test_vectors;
724    use crate::algorithms::test_util::{b64, error_of, json as value, unb64, with, without};
725
726    macro_rules! rfc9964_vector {
727        ($vector:expr, $key_pair:ident, $public_key:ident) => {{
728            let vector = $vector;
729            let private_jwk = vector["jwk"].to_string();
730            let kid = vector["jwk"]["kid"].as_str().unwrap();
731            let key_pair = $key_pair::from_jwk(&private_jwk).unwrap();
732            assert_eq!(key_pair.key_id.as_deref(), Some(kid));
733            assert_eq!(key_pair.to_bytes(), [0u8; 32]);
734            assert_eq!(key_pair.jwk_thumbprint(), kid);
735            assert_eq!(value(&key_pair.to_jwk())["pub"], vector["jwk"]["pub"]);
736
737            let public_key = $public_key::from_jwk(&without(&private_jwk, &["priv"])).unwrap();
738            assert_eq!(public_key.jwk_thumbprint(), kid);
739            assert_eq!(key_pair.public_key().to_bytes(), public_key.to_bytes());
740
741            // The example has no claims, so check its signature directly.
742            let jws = vector["jws"].as_str().unwrap();
743            let (signed, signature) = jws.rsplit_once('.').unwrap();
744            public_key
745                .public_key()
746                .as_ref()
747                .verify(signed.as_bytes(), &unb64(signature))
748                .unwrap();
749            private_jwk
750        }};
751    }
752
753    #[test]
754    fn rfc9964_vectors() {
755        let vectors = value(jwk_test_vectors::RFC9964_JOSE);
756        let jwk44 = rfc9964_vector!(&vectors[0], MLDSA44KeyPair, MLDSA44PublicKey);
757        let jwk65 = rfc9964_vector!(&vectors[1], MLDSA65KeyPair, MLDSA65PublicKey);
758        let jwk87 = rfc9964_vector!(&vectors[2], MLDSA87KeyPair, MLDSA87PublicKey);
759
760        // Changing `alg` must not let a key load as a different type.
761        assert!(MLDSA65KeyPair::from_jwk(&jwk44).is_err());
762        assert!(MLDSA44KeyPair::from_jwk(&jwk65).is_err());
763        assert!(MLDSA44PublicKey::from_jwk(&without(&jwk87, &["priv"])).is_err());
764        let relabeled = with(&jwk44, "alg", "ML-DSA-65".into());
765        assert!(MLDSA65KeyPair::from_jwk(&relabeled).is_err());
766        assert!(MLDSA44KeyPair::from_jwk(&relabeled).is_err());
767    }
768
769    #[test]
770    fn akp_keys_require_their_algorithm() {
771        let key_pair = MLDSA44KeyPair::generate();
772        let private_jwk = key_pair.to_jwk();
773        let public_jwk = key_pair.public_key().to_jwk();
774        for jwk in [
775            without(&public_jwk, &["alg"]),
776            with(&public_jwk, "alg", "ML-DSA-44 ".into()),
777            with(&public_jwk, "alg", serde_json::Value::Null),
778            with(&public_jwk, "crv", "Ed25519".into()),
779            with(&public_jwk, "x", "AAAA".into()),
780            with(&public_jwk, "priv", serde_json::Value::Null),
781            with(&public_jwk, "use", "enc".into()),
782            with(&public_jwk, "kty", "OKP".into()),
783        ] {
784            assert!(
785                matches!(
786                    error_of(MLDSA44PublicKey::from_jwk(&jwk)),
787                    JWTError::InvalidPublicKey
788                ),
789                "{}",
790                jwk
791            );
792        }
793        assert!(matches!(
794            error_of(MLDSA44KeyPair::from_jwk(&without(&private_jwk, &["alg"]))),
795            JWTError::InvalidKeyPair
796        ));
797    }
798
799    #[test]
800    fn private_key_is_the_seed() {
801        let key_pair = MLDSA44KeyPair::generate();
802        let other = MLDSA44KeyPair::generate();
803        let private_jwk = key_pair.to_jwk();
804        let seed = key_pair.to_bytes();
805        let with_private = |private: &[u8]| with(&private_jwk, "priv", b64(private).into());
806        let other_public = value(&other.to_jwk())["pub"].clone();
807        for (case, jwk) in [
808            ("short seed", with_private(&seed[..31])),
809            ("long seed", with_private(&[&seed[..], &[0]].concat())),
810            // ML-DSA-44 expanded keys are 2560 bytes.
811            ("expanded key", with_private(&[0x5au8; 2560])),
812            ("empty", with_private(&[])),
813            ("other seed", with_private(&other.to_bytes())),
814            ("other public key", with(&private_jwk, "pub", other_public)),
815            ("no public key", without(&private_jwk, &["pub"])),
816        ] {
817            assert!(
818                matches!(
819                    error_of(MLDSA44KeyPair::from_jwk(&jwk)),
820                    JWTError::InvalidKeyPair
821                ),
822                "{}",
823                case
824            );
825        }
826    }
827
828    #[test]
829    fn round_trips() {
830        let key_pair = MLDSA65KeyPair::generate().with_key_id("pq");
831        let private_jwk = key_pair.to_jwk();
832        let public_jwk = key_pair.public_key().to_jwk();
833        let exported = value(&public_jwk);
834        assert_eq!(exported["kty"], "AKP");
835        assert_eq!(exported["alg"], "ML-DSA-65");
836        assert_eq!(exported["use"], "sig");
837        assert_eq!(exported["kid"], "pq");
838        assert!(exported.get("priv").is_none());
839
840        let restored = MLDSA65KeyPair::from_jwk(&private_jwk).unwrap();
841        let public_key = MLDSA65PublicKey::from_jwk(&public_jwk).unwrap();
842        assert_eq!(restored.to_jwk(), private_jwk);
843        assert_eq!(public_key.to_jwk(), public_jwk);
844        assert_eq!(restored.jwk_thumbprint(), public_key.jwk_thumbprint());
845        let token = restored
846            .sign(Claims::create(coarsetime::Duration::from_hours(1)))
847            .unwrap();
848        public_key
849            .verify_token::<NoCustomClaims>(&token, None)
850            .unwrap();
851
852        let mut public_key = public_key;
853        let key_id = public_key.create_key_id().to_string();
854        let raw_hash = hmac_sha256::Hash::hash(&public_key.to_bytes());
855        assert_eq!(key_id, b64(&raw_hash));
856    }
857
858    #[test]
859    fn webcrypto_exports() {
860        for export in jwk_test_vectors::ALL {
861            let export = value(export);
862            let keys = &export["keys"];
863            if let Some(entry) = keys.get("ML-DSA-44") {
864                let key_pair = MLDSA44KeyPair::from_jwk(&entry["private"].to_string()).unwrap();
865                let public_key = MLDSA44PublicKey::from_jwk(&entry["public"].to_string()).unwrap();
866                assert_eq!(key_pair.jwk_thumbprint(), public_key.jwk_thumbprint());
867                public_key
868                    .verify_token::<NoCustomClaims>(entry["jwt"].as_str().unwrap(), None)
869                    .unwrap();
870            }
871            if let Some(entry) = keys.get("ML-DSA-65") {
872                let key_pair = MLDSA65KeyPair::from_jwk(&entry["private"].to_string()).unwrap();
873                let public_key = MLDSA65PublicKey::from_jwk(&entry["public"].to_string()).unwrap();
874                assert_eq!(key_pair.jwk_thumbprint(), public_key.jwk_thumbprint());
875            }
876            if let Some(entry) = keys.get("ML-DSA-87") {
877                let key_pair = MLDSA87KeyPair::from_jwk(&entry["private"].to_string()).unwrap();
878                let public_key = MLDSA87PublicKey::from_jwk(&entry["public"].to_string()).unwrap();
879                assert_eq!(key_pair.jwk_thumbprint(), public_key.jwk_thumbprint());
880            }
881        }
882    }
883}
jwt_simple/algorithms/mldsa.rs

jwt_simple/algorithms/
mldsa.rs
