qcoin_crypto/
lib.rs

1use serde::{
2    de::Error as DeError, ser::Error as SerError, Deserialize, Deserializer, Serialize, Serializer,
3};
4use std::collections::HashMap;
5use std::fmt::Display;
6
7use pqcrypto_dilithium::dilithium2;
8use pqcrypto_falcon::falcon512;
9use pqcrypto_traits::sign::{
10    DetachedSignature, PublicKey as PqPublicKeyTrait, SecretKey as PqSecretKeyTrait,
11};
12use zeroize::Zeroizing;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
15pub enum SignatureSchemeId {
16    Dilithium2,
17    Falcon512,
18    Unknown(u16),
19}
20
21#[derive(Debug, thiserror::Error, PartialEq, Eq)]
22pub enum CryptoError {
23    #[error("signature scheme mismatch")]
24    WrongScheme,
25    #[error("invalid public key")]
26    InvalidPublicKey,
27    #[error("invalid secret key")]
28    InvalidSecretKey,
29    #[error("invalid signature")]
30    InvalidSignature,
31    #[error("unsupported scheme {0}")]
32    UnsupportedScheme(SignatureSchemeId),
33}
34
35impl SignatureSchemeId {
36    pub const DILITHIUM2_ID: u16 = 0x01;
37    pub const FALCON512_ID: u16 = 0x02;
38
39    pub fn from_u16(id: u16) -> Self {
40        match id {
41            Self::DILITHIUM2_ID => SignatureSchemeId::Dilithium2,
42            Self::FALCON512_ID => SignatureSchemeId::Falcon512,
43            other => SignatureSchemeId::Unknown(other),
44        }
45    }
46
47    pub fn to_u16(self) -> u16 {
48        match self {
49            SignatureSchemeId::Dilithium2 => Self::DILITHIUM2_ID,
50            SignatureSchemeId::Falcon512 => Self::FALCON512_ID,
51            SignatureSchemeId::Unknown(value) => value,
52        }
53    }
54}
55
56impl Display for SignatureSchemeId {
57    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58        match self {
59            SignatureSchemeId::Dilithium2 => write!(f, "dilithium2"),
60            SignatureSchemeId::Falcon512 => write!(f, "falcon512"),
61            SignatureSchemeId::Unknown(value) => write!(f, "unknown-{}", value),
62        }
63    }
64}
65
66#[derive(Debug, Clone, PartialEq, Eq)]
67pub struct PublicKey {
68    pub scheme: SignatureSchemeId,
69    pub bytes: Vec<u8>,
70}
71
72#[derive(Debug, PartialEq, Eq)]
73pub struct PrivateKey {
74    pub scheme: SignatureSchemeId,
75    pub bytes: Zeroizing<Vec<u8>>,
76}
77
78#[derive(Debug, Clone, PartialEq, Eq)]
79pub struct Signature {
80    pub scheme: SignatureSchemeId,
81    pub bytes: Vec<u8>,
82}
83
84impl PublicKey {
85    pub fn new(scheme: SignatureSchemeId, bytes: Vec<u8>) -> Result<Self, CryptoError> {
86        validate_key_size(
87            public_key_size,
88            scheme,
89            bytes.len(),
90            false,
91            CryptoError::InvalidPublicKey,
92        )?;
93        Ok(Self { scheme, bytes })
94    }
95
96    pub fn to_bytes(&self) -> Result<Vec<u8>, CryptoError> {
97        encode_with_scheme(
98            self.scheme,
99            public_key_size,
100            &self.bytes,
101            false,
102            CryptoError::InvalidPublicKey,
103        )
104    }
105
106    pub fn from_bytes(encoded: &[u8]) -> Result<Self, CryptoError> {
107        let (scheme, payload) = decode_with_scheme(encoded, CryptoError::InvalidPublicKey)?;
108        validate_key_size(
109            public_key_size,
110            scheme,
111            payload.len(),
112            false,
113            CryptoError::InvalidPublicKey,
114        )?;
115        Ok(Self {
116            scheme,
117            bytes: payload,
118        })
119    }
120}
121
122impl PrivateKey {
123    pub fn new(scheme: SignatureSchemeId, bytes: Vec<u8>) -> Result<Self, CryptoError> {
124        validate_key_size(
125            secret_key_size,
126            scheme,
127            bytes.len(),
128            false,
129            CryptoError::InvalidSecretKey,
130        )?;
131        Ok(Self {
132            scheme,
133            bytes: Zeroizing::new(bytes),
134        })
135    }
136
137    pub fn to_bytes(&self) -> Result<Vec<u8>, CryptoError> {
138        encode_with_scheme(
139            self.scheme,
140            secret_key_size,
141            &self.bytes,
142            false,
143            CryptoError::InvalidSecretKey,
144        )
145    }
146
147    pub fn from_bytes(encoded: &[u8]) -> Result<Self, CryptoError> {
148        let (scheme, payload) = decode_with_scheme(encoded, CryptoError::InvalidSecretKey)?;
149        validate_key_size(
150            secret_key_size,
151            scheme,
152            payload.len(),
153            false,
154            CryptoError::InvalidSecretKey,
155        )?;
156        Ok(Self {
157            scheme,
158            bytes: Zeroizing::new(payload),
159        })
160    }
161}
162
163impl Signature {
164    pub fn new(scheme: SignatureSchemeId, bytes: Vec<u8>) -> Result<Self, CryptoError> {
165        validate_key_size(
166            signature_size,
167            scheme,
168            bytes.len(),
169            true,
170            CryptoError::InvalidSignature,
171        )?;
172        Ok(Self { scheme, bytes })
173    }
174
175    pub fn to_bytes(&self) -> Result<Vec<u8>, CryptoError> {
176        encode_with_scheme(
177            self.scheme,
178            signature_size,
179            &self.bytes,
180            true,
181            CryptoError::InvalidSignature,
182        )
183    }
184
185    pub fn from_bytes(encoded: &[u8]) -> Result<Self, CryptoError> {
186        let (scheme, payload) = decode_with_scheme(encoded, CryptoError::InvalidSignature)?;
187        validate_key_size(
188            signature_size,
189            scheme,
190            payload.len(),
191            true,
192            CryptoError::InvalidSignature,
193        )?;
194        Ok(Self {
195            scheme,
196            bytes: payload,
197        })
198    }
199}
200
201impl Serialize for PublicKey {
202    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
203    where
204        S: Serializer,
205    {
206        let bytes = self.to_bytes().map_err(SerError::custom)?;
207        serializer.serialize_bytes(&bytes)
208    }
209}
210
211impl<'de> Deserialize<'de> for PublicKey {
212    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
213    where
214        D: Deserializer<'de>,
215    {
216        struct PublicKeyVisitor;
217
218        impl<'de> serde::de::Visitor<'de> for PublicKeyVisitor {
219            type Value = PublicKey;
220
221            fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
222                write!(f, "byte-encoded public key")
223            }
224
225            fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
226            where
227                E: DeError,
228            {
229                PublicKey::from_bytes(v).map_err(DeError::custom)
230            }
231        }
232
233        deserializer.deserialize_bytes(PublicKeyVisitor)
234    }
235}
236
237impl Serialize for PrivateKey {
238    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
239    where
240        S: Serializer,
241    {
242        let bytes = self.to_bytes().map_err(SerError::custom)?;
243        serializer.serialize_bytes(&bytes)
244    }
245}
246
247impl<'de> Deserialize<'de> for PrivateKey {
248    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
249    where
250        D: Deserializer<'de>,
251    {
252        struct PrivateKeyVisitor;
253
254        impl<'de> serde::de::Visitor<'de> for PrivateKeyVisitor {
255            type Value = PrivateKey;
256
257            fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
258                write!(f, "byte-encoded private key")
259            }
260
261            fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
262            where
263                E: DeError,
264            {
265                PrivateKey::from_bytes(v).map_err(DeError::custom)
266            }
267        }
268
269        deserializer.deserialize_bytes(PrivateKeyVisitor)
270    }
271}
272
273impl Serialize for Signature {
274    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
275    where
276        S: Serializer,
277    {
278        let bytes = self.to_bytes().map_err(SerError::custom)?;
279        serializer.serialize_bytes(&bytes)
280    }
281}
282
283impl<'de> Deserialize<'de> for Signature {
284    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
285    where
286        D: Deserializer<'de>,
287    {
288        struct SignatureVisitor;
289
290        impl<'de> serde::de::Visitor<'de> for SignatureVisitor {
291            type Value = Signature;
292
293            fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
294                write!(f, "byte-encoded signature")
295            }
296
297            fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
298            where
299                E: DeError,
300            {
301                Signature::from_bytes(v).map_err(DeError::custom)
302            }
303        }
304
305        deserializer.deserialize_bytes(SignatureVisitor)
306    }
307}
308
309fn validate_key_size(
310    size_fn: fn(SignatureSchemeId) -> Option<usize>,
311    scheme: SignatureSchemeId,
312    len: usize,
313    allow_shorter: bool,
314    mismatch_error: CryptoError,
315) -> Result<(), CryptoError> {
316    match size_fn(scheme) {
317        Some(expected)
318            if (!allow_shorter && expected == len) || (allow_shorter && len <= expected) =>
319        {
320            Ok(())
321        }
322        Some(_) => Err(mismatch_error),
323        None => Err(CryptoError::UnsupportedScheme(scheme)),
324    }
325}
326
327fn encode_with_scheme(
328    scheme: SignatureSchemeId,
329    size_fn: fn(SignatureSchemeId) -> Option<usize>,
330    bytes: &[u8],
331    allow_shorter: bool,
332    mismatch_error: CryptoError,
333) -> Result<Vec<u8>, CryptoError> {
334    validate_key_size(size_fn, scheme, bytes.len(), allow_shorter, mismatch_error)?;
335
336    let len = bytes.len() as u32;
337    let mut encoded = Vec::with_capacity(2 + 4 + bytes.len());
338    encoded.extend_from_slice(&scheme.to_u16().to_le_bytes());
339    encoded.extend_from_slice(&len.to_le_bytes());
340    encoded.extend_from_slice(bytes);
341    Ok(encoded)
342}
343
344fn decode_with_scheme(
345    encoded: &[u8],
346    mismatch_error: CryptoError,
347) -> Result<(SignatureSchemeId, Vec<u8>), CryptoError> {
348    const PREFIX_LEN: usize = 2 + 4;
349    if encoded.len() < PREFIX_LEN {
350        return Err(mismatch_error);
351    }
352
353    let scheme = SignatureSchemeId::from_u16(u16::from_le_bytes([encoded[0], encoded[1]]));
354    let len = u32::from_le_bytes([encoded[2], encoded[3], encoded[4], encoded[5]]) as usize;
355
356    if encoded.len() != PREFIX_LEN + len {
357        return Err(mismatch_error);
358    }
359
360    let payload = encoded[PREFIX_LEN..].to_vec();
361    Ok((scheme, payload))
362}
363
364fn public_key_size(scheme: SignatureSchemeId) -> Option<usize> {
365    match scheme {
366        SignatureSchemeId::Dilithium2 => Some(dilithium2::public_key_bytes()),
367        SignatureSchemeId::Falcon512 => Some(falcon512::public_key_bytes()),
368        SignatureSchemeId::Unknown(_) => None,
369    }
370}
371
372fn secret_key_size(scheme: SignatureSchemeId) -> Option<usize> {
373    match scheme {
374        SignatureSchemeId::Dilithium2 => Some(dilithium2::secret_key_bytes()),
375        SignatureSchemeId::Falcon512 => Some(falcon512::secret_key_bytes()),
376        SignatureSchemeId::Unknown(_) => None,
377    }
378}
379
380fn signature_size(scheme: SignatureSchemeId) -> Option<usize> {
381    match scheme {
382        SignatureSchemeId::Dilithium2 => Some(dilithium2::signature_bytes()),
383        SignatureSchemeId::Falcon512 => Some(falcon512::signature_bytes()),
384        SignatureSchemeId::Unknown(_) => None,
385    }
386}
387
388pub trait PqSignatureScheme: Send + Sync {
389    fn id(&self) -> SignatureSchemeId;
390    fn keygen(&self) -> Result<(PublicKey, PrivateKey), CryptoError>;
391    fn sign(&self, sk: &PrivateKey, msg: &[u8]) -> Result<Signature, CryptoError>;
392    fn verify(&self, pk: &PublicKey, msg: &[u8], sig: &Signature) -> Result<(), CryptoError>;
393}
394
395pub trait PqSchemeRegistry {
396    fn get(&self, id: &SignatureSchemeId) -> Option<&dyn PqSignatureScheme>;
397}
398
399pub struct InMemoryRegistry {
400    schemes: HashMap<SignatureSchemeId, Box<dyn PqSignatureScheme>>,
401}
402
403impl InMemoryRegistry {
404    pub fn new() -> Self {
405        Self {
406            schemes: HashMap::new(),
407        }
408    }
409
410    pub fn with_scheme(mut self, scheme: Box<dyn PqSignatureScheme>) -> Self {
411        let id = scheme.id();
412        self.schemes.insert(id, scheme);
413        self
414    }
415
416    pub fn add_scheme(&mut self, scheme: Box<dyn PqSignatureScheme>) {
417        let id = scheme.id();
418        self.schemes.insert(id, scheme);
419    }
420}
421
422impl PqSchemeRegistry for InMemoryRegistry {
423    fn get(&self, id: &SignatureSchemeId) -> Option<&dyn PqSignatureScheme> {
424        self.schemes.get(id).map(|boxed| boxed.as_ref())
425    }
426}
427
428pub struct Dilithium2Scheme;
429
430impl PqSignatureScheme for Dilithium2Scheme {
431    fn id(&self) -> SignatureSchemeId {
432        SignatureSchemeId::Dilithium2
433    }
434
435    fn keygen(&self) -> Result<(PublicKey, PrivateKey), CryptoError> {
436        let (pk, sk) = dilithium2::keypair();
437        Ok((
438            PublicKey::new(self.id(), pk.as_bytes().to_vec())?,
439            PrivateKey::new(self.id(), sk.as_bytes().to_vec())?,
440        ))
441    }
442
443    fn sign(&self, sk: &PrivateKey, msg: &[u8]) -> Result<Signature, CryptoError> {
444        if sk.scheme != self.id() {
445            return Err(CryptoError::WrongScheme);
446        }
447
448        let sk = dilithium2::SecretKey::from_bytes(&sk.bytes)
449            .map_err(|_| CryptoError::InvalidSecretKey)?;
450
451        let signature = dilithium2::detached_sign(msg, &sk);
452        Signature::new(self.id(), signature.as_bytes().to_vec())
453    }
454
455    fn verify(&self, pk: &PublicKey, msg: &[u8], sig: &Signature) -> Result<(), CryptoError> {
456        if pk.scheme != self.id() || sig.scheme != self.id() {
457            return Err(CryptoError::WrongScheme);
458        }
459
460        let pk = dilithium2::PublicKey::from_bytes(&pk.bytes)
461            .map_err(|_| CryptoError::InvalidPublicKey)?;
462        let sig = dilithium2::DetachedSignature::from_bytes(&sig.bytes)
463            .map_err(|_| CryptoError::InvalidSignature)?;
464
465        dilithium2::verify_detached_signature(&sig, msg, &pk)
466            .map_err(|_| CryptoError::InvalidSignature)
467    }
468}
469
470pub struct Falcon512Scheme;
471
472impl PqSignatureScheme for Falcon512Scheme {
473    fn id(&self) -> SignatureSchemeId {
474        SignatureSchemeId::Falcon512
475    }
476
477    fn keygen(&self) -> Result<(PublicKey, PrivateKey), CryptoError> {
478        let (pk, sk) = falcon512::keypair();
479        Ok((
480            PublicKey::new(self.id(), pk.as_bytes().to_vec())?,
481            PrivateKey::new(self.id(), sk.as_bytes().to_vec())?,
482        ))
483    }
484
485    fn sign(&self, sk: &PrivateKey, msg: &[u8]) -> Result<Signature, CryptoError> {
486        if sk.scheme != self.id() {
487            return Err(CryptoError::WrongScheme);
488        }
489
490        let sk = falcon512::SecretKey::from_bytes(&sk.bytes)
491            .map_err(|_| CryptoError::InvalidSecretKey)?;
492
493        let signature = falcon512::detached_sign(msg, &sk);
494        Signature::new(self.id(), signature.as_bytes().to_vec())
495    }
496
497    fn verify(&self, pk: &PublicKey, msg: &[u8], sig: &Signature) -> Result<(), CryptoError> {
498        if pk.scheme != self.id() || sig.scheme != self.id() {
499            return Err(CryptoError::WrongScheme);
500        }
501
502        let pk = falcon512::PublicKey::from_bytes(&pk.bytes)
503            .map_err(|_| CryptoError::InvalidPublicKey)?;
504        let sig = falcon512::DetachedSignature::from_bytes(&sig.bytes)
505            .map_err(|_| CryptoError::InvalidSignature)?;
506
507        falcon512::verify_detached_signature(&sig, msg, &pk)
508            .map_err(|_| CryptoError::InvalidSignature)
509    }
510}
511
512pub fn default_registry() -> InMemoryRegistry {
513    InMemoryRegistry::new()
514        .with_scheme(Box::new(Dilithium2Scheme))
515        .with_scheme(Box::new(Falcon512Scheme))
516}
517
518#[cfg(test)]
519mod tests {
520    use super::*;
521
522    const MESSAGE: &[u8] = b"qcoin-pq-test";
523
524    #[test]
525    fn dilithium_roundtrip() {
526        let scheme = Dilithium2Scheme;
527        let (pk, sk) = scheme.keygen().expect("keygen should succeed");
528        let sig = scheme.sign(&sk, MESSAGE).expect("signing should succeed");
529
530        assert!(scheme.verify(&pk, MESSAGE, &sig).is_ok());
531    }
532
533    #[test]
534    fn falcon_roundtrip() {
535        let scheme = Falcon512Scheme;
536        let (pk, sk) = scheme.keygen().expect("keygen should succeed");
537        let sig = scheme.sign(&sk, MESSAGE).expect("signing should succeed");
538
539        assert!(scheme.verify(&pk, MESSAGE, &sig).is_ok());
540    }
541
542    #[test]
543    fn signing_rejects_wrong_scheme_key() {
544        let dilithium = Dilithium2Scheme;
545        let falcon = Falcon512Scheme;
546
547        let (_, sk) = dilithium.keygen().expect("keygen should succeed");
548        let result = falcon.sign(&sk, MESSAGE);
549
550        assert!(matches!(result, Err(CryptoError::WrongScheme)));
551    }
552
553    #[test]
554    fn verify_rejects_wrong_scheme_signature() {
555        let dilithium = Dilithium2Scheme;
556        let falcon = Falcon512Scheme;
557
558        let (dili_pk, _dili_sk) = dilithium.keygen().expect("keygen should succeed");
559        let falcon_sig = falcon
560            .sign(
561                &falcon.keygen().expect("falcon keygen should succeed").1,
562                MESSAGE,
563            )
564            .expect("falcon signing should succeed");
565
566        let result = dilithium.verify(&dili_pk, MESSAGE, &falcon_sig);
567        assert!(matches!(result, Err(CryptoError::WrongScheme)));
568
569        let sig_from_wrong_scheme = Signature {
570            scheme: SignatureSchemeId::Dilithium2,
571            bytes: falcon_sig.bytes.clone(),
572        };
573        let result = dilithium.verify(&dili_pk, MESSAGE, &sig_from_wrong_scheme);
574        assert!(matches!(result, Err(CryptoError::InvalidSignature)));
575    }
576
577    #[test]
578    fn verify_rejects_corrupted_public_key() {
579        let scheme = Dilithium2Scheme;
580        let (pk, sk) = scheme.keygen().expect("keygen should succeed");
581        let sig = scheme.sign(&sk, MESSAGE).expect("signing should succeed");
582
583        let mut corrupted_bytes = pk.bytes.clone();
584        corrupted_bytes.push(0);
585        let corrupted_pk = PublicKey {
586            scheme: pk.scheme,
587            bytes: corrupted_bytes,
588        };
589
590        let result = scheme.verify(&corrupted_pk, MESSAGE, &sig);
591        assert!(matches!(result, Err(CryptoError::InvalidPublicKey)));
592    }
593
594    #[test]
595    fn verify_rejects_corrupted_signature() {
596        let scheme = Falcon512Scheme;
597        let (pk, sk) = scheme.keygen().expect("keygen should succeed");
598        let mut sig = scheme.sign(&sk, MESSAGE).expect("signing should succeed");
599        sig.bytes.push(1);
600
601        let result = scheme.verify(&pk, MESSAGE, &sig);
602        assert!(matches!(result, Err(CryptoError::InvalidSignature)));
603    }
604
605    #[test]
606    fn verify_rejects_corrupted_message() {
607        let scheme = Dilithium2Scheme;
608        let (pk, sk) = scheme.keygen().expect("keygen should succeed");
609        let sig = scheme.sign(&sk, MESSAGE).expect("signing should succeed");
610
611        let result = scheme.verify(&pk, b"tampered", &sig);
612        assert!(matches!(result, Err(CryptoError::InvalidSignature)));
613    }
614
615    #[test]
616    fn roundtrip_serialization_is_canonical() {
617        let scheme = Falcon512Scheme;
618        let (pk, sk) = scheme.keygen().expect("keygen should succeed");
619        let sig = scheme.sign(&sk, MESSAGE).expect("signing should succeed");
620
621        let encoded_pk = pk.to_bytes().expect("public key encoding");
622        let decoded_pk = PublicKey::from_bytes(&encoded_pk).expect("public key decoding");
623        assert_eq!(pk, decoded_pk);
624
625        let encoded_sig = sig.to_bytes().expect("signature encoding");
626        let decoded_sig = Signature::from_bytes(&encoded_sig).expect("signature decoding");
627        assert_eq!(sig, decoded_sig);
628    }
629}
qcoin_crypto/lib.rs

qcoin_crypto/
lib.rs
