tholos_pq/
crypto.rs

1//! Cryptographic operations for encryption, decryption, and key generation.
2
3use std::time::{SystemTime, UNIX_EPOCH};
4
5use chacha20poly1305::aead::{Aead, KeyInit};
6use chacha20poly1305::{XChaCha20Poly1305, XNonce};
7use dilithium::{DilithiumKeyPair, DilithiumSignature, ML_DSA_65};
8use hkdf::Hkdf;
9use ml_kem::kem::{Decapsulate, Encapsulate};
10use ml_kem::{Ciphertext, EncodedSizeUser, KemCore, MlKem1024};
11use rand::rngs::OsRng;
12use rand::{CryptoRng, RngCore};
13use sha2::Sha256;
14use zeroize::{Zeroize, Zeroizing};
15
16use crate::errors::TholosError;
17use crate::types::*;
18
19/// Recipient private key material.
20///
21/// The ML-KEM decapsulation key is zeroized on drop (`ml-kem` `zeroize` feature).
22pub struct RecipientPriv {
23    /// Recipient identifier.
24    pub kid: String,
25    /// ML-KEM-1024 decapsulation (secret) key.
26    pub sk_kyber: <MlKem1024 as KemCore>::DecapsulationKey,
27}
28
29/// Sender keypair for signing messages.
30///
31/// Private key material is held inside [`DilithiumKeyPair`], which zeroizes on drop.
32pub struct SenderKeypair {
33    /// Sender identifier.
34    pub sid: String,
35    keypair: DilithiumKeyPair,
36}
37
38impl Drop for SenderKeypair {
39    fn drop(&mut self) {
40        self.sid.zeroize();
41    }
42}
43
44impl SenderKeypair {
45    /// ML-DSA-65 public key bytes.
46    pub fn public_key_bytes(&self) -> &[u8] {
47        self.keypair.public_key()
48    }
49
50    /// ML-DSA-65 private key bytes (for tests / vectors).
51    pub fn private_key_bytes(&self) -> &[u8] {
52        self.keypair.private_key()
53    }
54}
55
56/// Options for deterministic / injectable encryption parameters.
57pub struct EncryptOptions<'a, R: RngCore + CryptoRng> {
58    /// RNG used for CEK, nonces, ML-KEM encapsulation, and signing randomness.
59    pub rng: &'a mut R,
60    /// Message ID written into the header (normally a UUID).
61    pub msg_id: String,
62    /// Unix timestamp written into the header.
63    pub timestamp_unix: u64,
64}
65
66/// Generate a new recipient keypair.
67pub fn gen_recipient_keypair(kid: &str) -> (RecipientPub, RecipientPriv) {
68    let mut rng = OsRng;
69    gen_recipient_keypair_with(kid, &mut rng)
70}
71
72/// Generate a recipient keypair using the provided RNG (for tests / vectors).
73pub fn gen_recipient_keypair_with<R: RngCore + CryptoRng>(
74    kid: &str,
75    rng: &mut R,
76) -> (RecipientPub, RecipientPriv) {
77    let (sk, pk) = MlKem1024::generate(rng);
78    let pub_bytes = pk.as_bytes().to_vec();
79    (
80        RecipientPub {
81            kid: kid.to_string(),
82            pk_kyber: pub_bytes,
83        },
84        RecipientPriv {
85            kid: kid.to_string(),
86            sk_kyber: sk,
87        },
88    )
89}
90
91/// Generate a new sender keypair (ML-DSA-65).
92pub fn gen_sender_keypair(sid: &str) -> SenderKeypair {
93    #[allow(clippy::expect_used)]
94    let keypair = DilithiumKeyPair::generate(ML_DSA_65).expect("ML-DSA-65 key generation failed");
95    SenderKeypair {
96        sid: sid.to_string(),
97        keypair,
98    }
99}
100
101/// Generate a sender keypair deterministically from a 32-byte seed (tests / vectors).
102pub fn gen_sender_keypair_deterministic(sid: &str, seed: &[u8; 32]) -> SenderKeypair {
103    SenderKeypair {
104        sid: sid.to_string(),
105        keypair: DilithiumKeyPair::generate_deterministic(ML_DSA_65, seed),
106    }
107}
108
109/// Reconstruct a sender keypair from raw ML-DSA-65 key bytes (for tests / vectors).
110pub fn sender_keypair_from_bytes(
111    sid: &str,
112    pk_bytes: &[u8],
113    sk_bytes: &[u8],
114) -> Result<SenderKeypair, TholosError> {
115    let keypair = DilithiumKeyPair::from_keys(sk_bytes, pk_bytes, ML_DSA_65)
116        .map_err(|_| TholosError::Malformed("ml-dsa keys"))?;
117    Ok(SenderKeypair {
118        sid: sid.to_string(),
119        keypair,
120    })
121}
122
123/// Extract the public key information from a sender keypair.
124pub fn sender_pub(sender: &SenderKeypair) -> SenderPub {
125    SenderPub {
126        sid: sender.sid.clone(),
127        pk_dilithium: sender.public_key_bytes().to_vec(),
128    }
129}
130
131fn unix_timestamp_secs() -> Result<u64, TholosError> {
132    SystemTime::now()
133        .duration_since(UNIX_EPOCH)
134        .map_err(|_| TholosError::Malformed("system clock before unix epoch"))
135        .map(|d| d.as_secs())
136}
137
138fn hkdf32(shared: &[u8], kid: &str, header_cbor: &[u8]) -> Zeroizing<[u8; 32]> {
139    let hk = Hkdf::<Sha256>::new(Some(kid.as_bytes()), shared);
140    let mut okm = Zeroizing::new([0u8; 32]);
141    #[allow(clippy::expect_used)]
142    hk.expand(header_cbor, okm.as_mut())
143        .expect("HKDF expand failed - this should never happen with 32-byte output");
144    okm
145}
146
147fn aead_enc(
148    key: &[u8; 32],
149    nonce24: &[u8; 24],
150    aad: &[u8],
151    pt: &[u8],
152) -> Result<Vec<u8>, TholosError> {
153    let cipher = XChaCha20Poly1305::new(key.into());
154    let nonce = XNonce::from(*nonce24);
155    cipher
156        .encrypt(&nonce, chacha20poly1305::aead::Payload { msg: pt, aad })
157        .map_err(|_| TholosError::Aead)
158}
159
160fn aead_dec(
161    key: &[u8; 32],
162    nonce24: &[u8; 24],
163    aad: &[u8],
164    ct: &[u8],
165) -> Result<Vec<u8>, TholosError> {
166    let cipher = XChaCha20Poly1305::new(key.into());
167    let nonce = XNonce::from(*nonce24);
168    cipher
169        .decrypt(&nonce, chacha20poly1305::aead::Payload { msg: ct, aad })
170        .map_err(|_| TholosError::Aead)
171}
172
173fn expect_self_describe_tag(wire_cbor: &[u8]) -> Result<(), TholosError> {
174    if !has_self_describe_tag(wire_cbor) {
175        return Err(TholosError::Malformed("wire self-describe tag"));
176    }
177    Ok(())
178}
179
180fn sign_inner<R: RngCore>(
181    inner_cbor: &[u8],
182    sender: &SenderKeypair,
183    rng: &mut R,
184) -> Result<Vec<u8>, TholosError> {
185    let mut rnd = [0u8; 32];
186    rng.fill_bytes(&mut rnd);
187    let sig = sender
188        .keypair
189        .sign_deterministic(inner_cbor, b"", &rnd)
190        .map_err(|_| TholosError::Malformed("ml-dsa sign"))?;
191    rnd.zeroize();
192    Ok(sig.as_bytes().to_vec())
193}
194
195fn validate_inner(inner: &BundleUnsigned) -> Result<(), TholosError> {
196    if inner.header.v != 1 || inner.header.suite != SUITE_V1 {
197        return Err(TholosError::UnsupportedSuite {
198            v: inner.header.v,
199            suite: inner.header.suite.clone(),
200        });
201    }
202
203    if inner.header.recipients.len() != inner.recipients.len() {
204        return Err(TholosError::Malformed("recipient list mismatch"));
205    }
206
207    let mut seen = std::collections::HashSet::new();
208    for (header_kid, env) in inner.header.recipients.iter().zip(&inner.recipients) {
209        if header_kid != &env.kid {
210            return Err(TholosError::Malformed("recipient order mismatch"));
211        }
212        if !seen.insert(env.kid.clone()) {
213            return Err(TholosError::Malformed("duplicate recipient kid"));
214        }
215    }
216
217    Ok(())
218}
219
220fn verify_signed_bundle(
221    bundle: &BundleSigned,
222    allowed_senders: &[(String, Vec<u8>)],
223) -> Result<BundleUnsigned, TholosError> {
224    if bundle.inner.is_empty() {
225        return Err(TholosError::Malformed("empty inner bundle"));
226    }
227
228    let inner: BundleUnsigned = from_cbor(&bundle.inner)?;
229    validate_inner(&inner)?;
230
231    let sender_sid = &inner.header.sender;
232    let Some((_, pk_bytes)) = allowed_senders.iter().find(|(sid, _)| sid == sender_sid) else {
233        return Err(TholosError::BadSignature);
234    };
235
236    if pk_bytes.len() != DILITHIUM3_PK_LEN {
237        return Err(TholosError::Malformed("dilithium pk"));
238    }
239    if bundle.sig_dilithium.len() != DILITHIUM3_SIG_LEN {
240        return Err(TholosError::Malformed("signature"));
241    }
242
243    let sig = DilithiumSignature::from_slice(&bundle.sig_dilithium);
244    if !DilithiumKeyPair::verify(pk_bytes, &sig, &bundle.inner, b"", ML_DSA_65) {
245        return Err(TholosError::BadSignature);
246    }
247
248    Ok(inner)
249}
250
251/// Verify a wire bundle and return its authenticated header without decrypting.
252pub fn verify_header(
253    wire_cbor: &[u8],
254    allowed_senders: &[(String, Vec<u8>)],
255) -> Result<Header, TholosError> {
256    expect_self_describe_tag(wire_cbor)?;
257    let bundle: BundleSigned = from_cbor(wire_cbor)?;
258    let inner = verify_signed_bundle(&bundle, allowed_senders)?;
259    Ok(inner.header)
260}
261
262/// Encrypt a message for multiple recipients and sign it with the sender's key.
263pub fn encrypt(
264    plaintext: &[u8],
265    sender: &SenderKeypair,
266    recipients: &[RecipientPub],
267) -> Result<Vec<u8>, TholosError> {
268    let mut rng = OsRng;
269    encrypt_with(
270        plaintext,
271        sender,
272        recipients,
273        EncryptOptions {
274            rng: &mut rng,
275            msg_id: uuid::Uuid::new_v4().to_string(),
276            timestamp_unix: unix_timestamp_secs()?,
277        },
278    )
279}
280
281/// Encrypt with caller-supplied RNG, message ID, and timestamp.
282pub fn encrypt_with<R: RngCore + CryptoRng>(
283    plaintext: &[u8],
284    sender: &SenderKeypair,
285    recipients: &[RecipientPub],
286    opts: EncryptOptions<'_, R>,
287) -> Result<Vec<u8>, TholosError> {
288    if recipients.is_empty() {
289        return Err(TholosError::NoRecipients);
290    }
291
292    let header = Header {
293        v: 1,
294        suite: SUITE_V1.to_string(),
295        sender: sender.sid.clone(),
296        recipients: recipients.iter().map(|r| r.kid.clone()).collect(),
297        msg_id: opts.msg_id,
298        timestamp_unix: opts.timestamp_unix,
299    };
300    let header_cbor = to_cbor(&header)?;
301
302    let rng = opts.rng;
303    let mut cek = Zeroizing::new([0u8; 32]);
304    rng.fill_bytes(cek.as_mut());
305
306    let mut pay_nonce = [0u8; 24];
307    rng.fill_bytes(&mut pay_nonce);
308    let ciphertext = aead_enc(&cek, &pay_nonce, &header_cbor, plaintext)?;
309
310    let mut envs = Vec::with_capacity(recipients.len());
311    for r in recipients {
312        let pk_bytes: &[u8] = &r.pk_kyber;
313        let pk = <MlKem1024 as KemCore>::EncapsulationKey::from_bytes(
314            &pk_bytes
315                .try_into()
316                .map_err(|_| TholosError::Malformed("ml-kem pk"))?,
317        );
318        let (kem_ct, shared) = pk
319            .encapsulate(rng)
320            .map_err(|_| TholosError::Malformed("encapsulation"))?;
321
322        let kek = hkdf32(shared.as_slice(), &r.kid, &header_cbor);
323
324        let mut wrap_nonce = [0u8; 24];
325        rng.fill_bytes(&mut wrap_nonce);
326        let wrapped_cek = aead_enc(&kek, &wrap_nonce, &header_cbor, cek.as_ref())?;
327
328        envs.push(RecipientEnvelope {
329            kid: r.kid.clone(),
330            kem_ct: kem_ct.as_slice().to_vec(),
331            wrap_nonce: wrap_nonce.to_vec(),
332            wrapped_cek,
333        });
334    }
335
336    let inner = BundleUnsigned {
337        header,
338        pay_nonce: pay_nonce.to_vec(),
339        ciphertext,
340        recipients: envs,
341    };
342
343    let inner_cbor = to_cbor(&inner)?;
344    let sig_dilithium = sign_inner(&inner_cbor, sender, rng)?;
345
346    let bundle = BundleSigned {
347        inner: inner_cbor,
348        sig_dilithium,
349    };
350
351    to_cbor(&bundle)
352}
353
354/// Decrypt a message as a recipient and verify the sender's signature.
355pub fn decrypt(
356    wire_cbor: &[u8],
357    my_kid: &str,
358    my_sk: &<MlKem1024 as KemCore>::DecapsulationKey,
359    allowed_senders: &[(String, Vec<u8>)],
360) -> Result<Vec<u8>, TholosError> {
361    expect_self_describe_tag(wire_cbor)?;
362    let bundle: BundleSigned = from_cbor(wire_cbor)?;
363    let inner = verify_signed_bundle(&bundle, allowed_senders)?;
364
365    let env = inner
366        .recipients
367        .iter()
368        .find(|e| e.kid == my_kid)
369        .ok_or_else(|| TholosError::MissingEnvelope(my_kid.to_string()))?;
370
371    if env.wrap_nonce.len() != 24 {
372        return Err(TholosError::Malformed("wrap nonce"));
373    }
374    let kem_ct_bytes: &[u8] = &env.kem_ct;
375    let kem_ct: Ciphertext<MlKem1024> = kem_ct_bytes
376        .try_into()
377        .map_err(|_| TholosError::Malformed("kem_ct"))?;
378    let shared = my_sk
379        .decapsulate(&kem_ct)
380        .map_err(|_| TholosError::Malformed("decapsulation"))?;
381
382    let header_cbor = to_cbor(&inner.header)?;
383    let kek = hkdf32(shared.as_slice(), my_kid, &header_cbor);
384
385    let mut wrap_nonce = [0u8; 24];
386    wrap_nonce.copy_from_slice(&env.wrap_nonce);
387    let cek = aead_dec(&kek, &wrap_nonce, &header_cbor, &env.wrapped_cek)?;
388
389    if cek.len() != 32 {
390        return Err(TholosError::Malformed("cek length"));
391    }
392    let mut cek_arr = Zeroizing::new([0u8; 32]);
393    cek_arr.copy_from_slice(&cek);
394
395    if inner.pay_nonce.len() != 24 {
396        return Err(TholosError::Malformed("pay nonce"));
397    }
398    let mut pay_nonce = [0u8; 24];
399    pay_nonce.copy_from_slice(&inner.pay_nonce);
400
401    aead_dec(&cek_arr, &pay_nonce, &header_cbor, &inner.ciphertext)
402}
403
404pub(crate) mod test_support {
405    use super::*;
406
407    pub fn resign_bundle(
408        inner: &BundleUnsigned,
409        sender: &SenderKeypair,
410    ) -> Result<BundleSigned, TholosError> {
411        let inner_cbor = to_cbor(inner)?;
412        let mut rng = OsRng;
413        let sig_dilithium = sign_inner(&inner_cbor, sender, &mut rng)?;
414        Ok(BundleSigned {
415            inner: inner_cbor,
416            sig_dilithium,
417        })
418    }
419
420    pub fn encode_bundle(bundle: &BundleSigned) -> Result<Vec<u8>, TholosError> {
421        to_cbor(bundle)
422    }
423
424    pub fn decode_bundle(wire: &[u8]) -> Result<BundleSigned, TholosError> {
425        from_cbor(wire)
426    }
427
428    pub fn decode_inner(bytes: &[u8]) -> Result<BundleUnsigned, TholosError> {
429        from_cbor(bytes)
430    }
431
432    pub fn encode_cbor<T: serde::Serialize>(v: &T) -> Result<Vec<u8>, TholosError> {
433        to_cbor(v)
434    }
435
436    pub fn decode_cbor<T: serde::de::DeserializeOwned>(data: &[u8]) -> Result<T, TholosError> {
437        from_cbor(data)
438    }
439}
tholos_pq/crypto.rs

tholos_pq/
crypto.rs
