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 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 #[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#[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#[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}