Skip to main content

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::claims::*;
8use crate::common::*;
9#[cfg(feature = "cwt")]
10use crate::cwt_token::*;
11use crate::error::*;
12use crate::jwt_header::*;
13use crate::token::*;
14
15#[doc(hidden)]
16#[derive(Debug, Clone)]
17pub struct MLDSAPublicKey(MlDsaPublicKey);
18
19impl AsRef<MlDsaPublicKey> for MLDSAPublicKey {
20    fn as_ref(&self) -> &MlDsaPublicKey {
21        &self.0
22    }
23}
24
25impl MLDSAPublicKey {
26    pub fn from_bytes(algorithm: Algorithm, raw: &[u8]) -> Result<Self, Error> {
27        let mldsa_pk = MlDsaPublicKey::from_slice(algorithm, raw);
28        Ok(MLDSAPublicKey(
29            mldsa_pk.map_err(|_| JWTError::InvalidPublicKey)?,
30        ))
31    }
32
33    pub fn to_bytes(&self) -> Vec<u8> {
34        self.0.to_bytes().expect("failed to serialize public key")
35    }
36}
37
38#[doc(hidden)]
39#[derive(Clone)]
40pub struct MLDSAKeyPair {
41    mldsa_sk: MlDsaPrivateKey,
42    metadata: Option<KeyMetadata>,
43}
44
45impl std::fmt::Debug for MLDSAKeyPair {
46    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47        f.debug_struct("PKey")
48            .field("algorithm", &"ML-DSA")
49            .finish()
50    }
51}
52
53impl AsRef<MlDsaPrivateKey> for MLDSAKeyPair {
54    fn as_ref(&self) -> &MlDsaPrivateKey {
55        &self.mldsa_sk
56    }
57}
58
59impl MLDSAKeyPair {
60    /// The raw representation of an ML-DSA key pair is its 32-byte seed, as
61    /// mandated for JOSE and COSE.
62    pub fn from_bytes(algorithm: Algorithm, raw: &[u8]) -> Result<Self, Error> {
63        let seed = raw.try_into().map_err(|_| JWTError::InvalidKeyPair)?;
64        let mldsa_sk =
65            MlDsaPrivateKey::from_seed(algorithm, &seed).map_err(|_| JWTError::InvalidKeyPair)?;
66        Ok(MLDSAKeyPair {
67            mldsa_sk,
68            metadata: None,
69        })
70    }
71
72    pub fn to_bytes(&self) -> Vec<u8> {
73        self.mldsa_sk.seed_bytes().to_vec()
74    }
75
76    pub fn public_key(&self) -> MLDSAPublicKey {
77        let mldsa_pk = self
78            .mldsa_sk
79            .public_key()
80            .expect("failed to create public key");
81        MLDSAPublicKey(mldsa_pk)
82    }
83
84    pub fn generate(algorithm: Algorithm) -> Self {
85        let (_, mldsa_sk) =
86            MlDsaPrivateKey::generate(algorithm).expect("failed to generate key pair");
87        MLDSAKeyPair {
88            mldsa_sk,
89            metadata: None,
90        }
91    }
92}
93
94pub trait MLDSAKeyPairLike {
95    fn jwt_alg_name() -> &'static str;
96    fn key_pair(&self) -> &MLDSAKeyPair;
97    fn key_id(&self) -> &Option<String>;
98    fn metadata(&self) -> &Option<KeyMetadata>;
99    fn attach_metadata(&mut self, metadata: KeyMetadata) -> Result<(), Error>;
100
101    fn sign<CustomClaims: Serialize>(
102        &self,
103        claims: JWTClaims<CustomClaims>,
104    ) -> Result<String, Error> {
105        self.sign_with_options(claims, &Default::default())
106    }
107
108    fn sign_with_options<CustomClaims: Serialize>(
109        &self,
110        claims: JWTClaims<CustomClaims>,
111        opts: &HeaderOptions,
112    ) -> Result<String, Error> {
113        let jwt_header = JWTHeader::new(Self::jwt_alg_name().to_string(), self.key_id().clone())
114            .with_key_metadata(self.metadata())
115            .with_options(opts);
116        Token::build(&jwt_header, claims, |authenticated| {
117            let signature = self.key_pair().as_ref().sign(authenticated.as_bytes())?;
118            Ok(signature)
119        })
120    }
121}
122
123pub trait MLDSAPublicKeyLike {
124    fn jwt_alg_name() -> &'static str;
125    fn public_key(&self) -> &MLDSAPublicKey;
126    fn key_id(&self) -> &Option<String>;
127    fn set_key_id(&mut self, key_id: String);
128
129    fn verify_token<CustomClaims: DeserializeOwned>(
130        &self,
131        token: &str,
132        options: Option<VerificationOptions>,
133    ) -> Result<JWTClaims<CustomClaims>, Error> {
134        Token::verify(
135            Self::jwt_alg_name(),
136            token,
137            options,
138            |authenticated, signature| {
139                self.public_key()
140                    .as_ref()
141                    .verify(authenticated.as_bytes(), signature)
142                    .map_err(|_| JWTError::InvalidSignature)?;
143                Ok(())
144            },
145            |_salt: Option<&[u8]>| Ok(()),
146        )
147    }
148
149    #[cfg(feature = "cwt")]
150    fn verify_cwt_token<CustomClaims: DeserializeOwned>(
151        &self,
152        token: &[u8],
153        options: Option<VerificationOptions>,
154    ) -> Result<JWTClaims<NoCustomClaims>, Error> {
155        CWTToken::verify(
156            Self::jwt_alg_name(),
157            token,
158            options,
159            |authenticated, signature| {
160                self.public_key()
161                    .as_ref()
162                    .verify(authenticated.as_bytes(), signature)
163                    .map_err(|_| JWTError::InvalidSignature)?;
164                Ok(())
165            },
166        )
167    }
168
169    /// Decode CWT token metadata that can be useful prior to signature/tag verification
170    #[cfg(feature = "cwt")]
171    fn decode_cwt_metadata(&self, token: impl AsRef<[u8]>) -> Result<TokenMetadata, Error> {
172        CWTToken::decode_metadata(token)
173    }
174
175    fn create_key_id(&mut self) -> &str {
176        self.set_key_id(
177            Base64UrlSafeNoPadding::encode_to_string(hmac_sha256::Hash::hash(
178                &self.public_key().to_bytes(),
179            ))
180            .unwrap(),
181        );
182        self.key_id().as_ref().map(|x| x.as_str()).unwrap()
183    }
184}
185
186#[derive(Clone)]
187pub struct MLDSA44KeyPair {
188    key_pair: MLDSAKeyPair,
189    key_id: Option<String>,
190}
191
192impl std::fmt::Debug for MLDSA44KeyPair {
193    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
194        f.debug_struct("PKey")
195            .field("algorithm", &"ML-DSA-44")
196            .finish()
197    }
198}
199
200#[derive(Debug, Clone)]
201pub struct MLDSA44PublicKey {
202    pk: MLDSAPublicKey,
203    key_id: Option<String>,
204}
205
206impl MLDSAKeyPairLike for MLDSA44KeyPair {
207    fn jwt_alg_name() -> &'static str {
208        "ML-DSA-44"
209    }
210
211    fn key_pair(&self) -> &MLDSAKeyPair {
212        &self.key_pair
213    }
214
215    fn key_id(&self) -> &Option<String> {
216        &self.key_id
217    }
218
219    fn metadata(&self) -> &Option<KeyMetadata> {
220        &self.key_pair.metadata
221    }
222
223    fn attach_metadata(&mut self, metadata: KeyMetadata) -> Result<(), Error> {
224        self.key_pair.metadata = Some(metadata);
225        Ok(())
226    }
227}
228
229impl MLDSA44KeyPair {
230    pub fn from_bytes(raw: &[u8]) -> Result<Self, Error> {
231        Ok(MLDSA44KeyPair {
232            key_pair: MLDSAKeyPair::from_bytes(Algorithm::MlDsa44, raw)?,
233            key_id: None,
234        })
235    }
236
237    pub fn to_bytes(&self) -> Vec<u8> {
238        self.key_pair.to_bytes()
239    }
240
241    pub fn public_key(&self) -> MLDSA44PublicKey {
242        MLDSA44PublicKey {
243            pk: self.key_pair.public_key(),
244            key_id: self.key_id.clone(),
245        }
246    }
247
248    pub fn generate() -> Self {
249        MLDSA44KeyPair {
250            key_pair: MLDSAKeyPair::generate(Algorithm::MlDsa44),
251            key_id: None,
252        }
253    }
254
255    pub fn with_key_id(mut self, key_id: &str) -> Self {
256        self.key_id = Some(key_id.to_string());
257        self
258    }
259}
260
261impl MLDSAPublicKeyLike for MLDSA44PublicKey {
262    fn jwt_alg_name() -> &'static str {
263        "ML-DSA-44"
264    }
265
266    fn public_key(&self) -> &MLDSAPublicKey {
267        &self.pk
268    }
269
270    fn key_id(&self) -> &Option<String> {
271        &self.key_id
272    }
273
274    fn set_key_id(&mut self, key_id: String) {
275        self.key_id = Some(key_id);
276    }
277}
278
279impl MLDSA44PublicKey {
280    pub fn from_bytes(raw: &[u8]) -> Result<Self, Error> {
281        Ok(MLDSA44PublicKey {
282            pk: MLDSAPublicKey::from_bytes(Algorithm::MlDsa44, raw)?,
283            key_id: None,
284        })
285    }
286
287    pub fn to_bytes(&self) -> Vec<u8> {
288        self.pk.to_bytes()
289    }
290
291    pub fn with_key_id(mut self, key_id: &str) -> Self {
292        self.key_id = Some(key_id.to_string());
293        self
294    }
295}
296
297//
298
299#[derive(Clone)]
300pub struct MLDSA65KeyPair {
301    key_pair: MLDSAKeyPair,
302    key_id: Option<String>,
303}
304
305impl std::fmt::Debug for MLDSA65KeyPair {
306    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
307        f.debug_struct("PKey")
308            .field("algorithm", &"ML-DSA-65")
309            .finish()
310    }
311}
312
313#[derive(Debug, Clone)]
314pub struct MLDSA65PublicKey {
315    pk: MLDSAPublicKey,
316    key_id: Option<String>,
317}
318
319impl MLDSAKeyPairLike for MLDSA65KeyPair {
320    fn jwt_alg_name() -> &'static str {
321        "ML-DSA-65"
322    }
323
324    fn key_pair(&self) -> &MLDSAKeyPair {
325        &self.key_pair
326    }
327
328    fn key_id(&self) -> &Option<String> {
329        &self.key_id
330    }
331
332    fn metadata(&self) -> &Option<KeyMetadata> {
333        &self.key_pair.metadata
334    }
335
336    fn attach_metadata(&mut self, metadata: KeyMetadata) -> Result<(), Error> {
337        self.key_pair.metadata = Some(metadata);
338        Ok(())
339    }
340}
341
342impl MLDSA65KeyPair {
343    pub fn from_bytes(raw: &[u8]) -> Result<Self, Error> {
344        Ok(MLDSA65KeyPair {
345            key_pair: MLDSAKeyPair::from_bytes(Algorithm::MlDsa65, raw)?,
346            key_id: None,
347        })
348    }
349
350    pub fn to_bytes(&self) -> Vec<u8> {
351        self.key_pair.to_bytes()
352    }
353
354    pub fn public_key(&self) -> MLDSA65PublicKey {
355        MLDSA65PublicKey {
356            pk: self.key_pair.public_key(),
357            key_id: self.key_id.clone(),
358        }
359    }
360
361    pub fn generate() -> Self {
362        MLDSA65KeyPair {
363            key_pair: MLDSAKeyPair::generate(Algorithm::MlDsa65),
364            key_id: None,
365        }
366    }
367
368    pub fn with_key_id(mut self, key_id: &str) -> Self {
369        self.key_id = Some(key_id.to_string());
370        self
371    }
372}
373
374impl MLDSAPublicKeyLike for MLDSA65PublicKey {
375    fn jwt_alg_name() -> &'static str {
376        "ML-DSA-65"
377    }
378
379    fn public_key(&self) -> &MLDSAPublicKey {
380        &self.pk
381    }
382
383    fn key_id(&self) -> &Option<String> {
384        &self.key_id
385    }
386
387    fn set_key_id(&mut self, key_id: String) {
388        self.key_id = Some(key_id);
389    }
390}
391
392impl MLDSA65PublicKey {
393    pub fn from_bytes(raw: &[u8]) -> Result<Self, Error> {
394        Ok(MLDSA65PublicKey {
395            pk: MLDSAPublicKey::from_bytes(Algorithm::MlDsa65, raw)?,
396            key_id: None,
397        })
398    }
399
400    pub fn to_bytes(&self) -> Vec<u8> {
401        self.pk.to_bytes()
402    }
403
404    pub fn with_key_id(mut self, key_id: &str) -> Self {
405        self.key_id = Some(key_id.to_string());
406        self
407    }
408}
409
410//
411
412#[derive(Clone)]
413pub struct MLDSA87KeyPair {
414    key_pair: MLDSAKeyPair,
415    key_id: Option<String>,
416}
417
418impl std::fmt::Debug for MLDSA87KeyPair {
419    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
420        f.debug_struct("PKey")
421            .field("algorithm", &"ML-DSA-87")
422            .finish()
423    }
424}
425
426#[derive(Debug, Clone)]
427pub struct MLDSA87PublicKey {
428    pk: MLDSAPublicKey,
429    key_id: Option<String>,
430}
431
432impl MLDSAKeyPairLike for MLDSA87KeyPair {
433    fn jwt_alg_name() -> &'static str {
434        "ML-DSA-87"
435    }
436
437    fn key_pair(&self) -> &MLDSAKeyPair {
438        &self.key_pair
439    }
440
441    fn key_id(&self) -> &Option<String> {
442        &self.key_id
443    }
444
445    fn metadata(&self) -> &Option<KeyMetadata> {
446        &self.key_pair.metadata
447    }
448
449    fn attach_metadata(&mut self, metadata: KeyMetadata) -> Result<(), Error> {
450        self.key_pair.metadata = Some(metadata);
451        Ok(())
452    }
453}
454
455impl MLDSA87KeyPair {
456    pub fn from_bytes(raw: &[u8]) -> Result<Self, Error> {
457        Ok(MLDSA87KeyPair {
458            key_pair: MLDSAKeyPair::from_bytes(Algorithm::MlDsa87, raw)?,
459            key_id: None,
460        })
461    }
462
463    pub fn to_bytes(&self) -> Vec<u8> {
464        self.key_pair.to_bytes()
465    }
466
467    pub fn public_key(&self) -> MLDSA87PublicKey {
468        MLDSA87PublicKey {
469            pk: self.key_pair.public_key(),
470            key_id: self.key_id.clone(),
471        }
472    }
473
474    pub fn generate() -> Self {
475        MLDSA87KeyPair {
476            key_pair: MLDSAKeyPair::generate(Algorithm::MlDsa87),
477            key_id: None,
478        }
479    }
480
481    pub fn with_key_id(mut self, key_id: &str) -> Self {
482        self.key_id = Some(key_id.to_string());
483        self
484    }
485}
486
487impl MLDSAPublicKeyLike for MLDSA87PublicKey {
488    fn jwt_alg_name() -> &'static str {
489        "ML-DSA-87"
490    }
491
492    fn public_key(&self) -> &MLDSAPublicKey {
493        &self.pk
494    }
495
496    fn key_id(&self) -> &Option<String> {
497        &self.key_id
498    }
499
500    fn set_key_id(&mut self, key_id: String) {
501        self.key_id = Some(key_id);
502    }
503}
504
505impl MLDSA87PublicKey {
506    pub fn from_bytes(raw: &[u8]) -> Result<Self, Error> {
507        Ok(MLDSA87PublicKey {
508            pk: MLDSAPublicKey::from_bytes(Algorithm::MlDsa87, raw)?,
509            key_id: None,
510        })
511    }
512
513    pub fn to_bytes(&self) -> Vec<u8> {
514        self.pk.to_bytes()
515    }
516
517    pub fn with_key_id(mut self, key_id: &str) -> Self {
518        self.key_id = Some(key_id.to_string());
519        self
520    }
521}