zoe_wire_protocol/inbox/
pqxdh.rs

1//! PQXDH (Post-Quantum Extended Diffie-Hellman) inbox types and cryptographic helpers
2//!
3//! This module implements the PQXDH protocol for asynchronous secure communication,
4//! based on Signal's PQXDH specification. It provides:
5//!
6//! - Prekey bundle generation and management
7//! - PQXDH key agreement protocol
8//! - Inbox types for different protocols (RPC, messaging, etc.)
9//!
10//! ## Security
11//!
12//! PQXDH provides:
13//! - Post-quantum forward secrecy via ML-KEM (using libcrux-ml-kem)
14//! - Classical security via X25519 ECDH
15//! - Authentication via ML-DSA signatures
16//! - Perfect forward secrecy through one-time prekeys
17
18use std::collections::BTreeMap;
19
20// Note: libcrux_ml_kem will be used for actual key generation in implementation
21pub mod pqxdh_crypto;
22pub use pqxdh_crypto::*;
23use serde::{Deserialize, Serialize};
24use zeroize::{Zeroize, ZeroizeOnDrop};
25
26use crate::{Signature, VerifyingKey};
27
28/// Inbox type indicating expected responsiveness and access control
29#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
30pub enum InboxType {
31    /// Only authenticated senders, recipient will likely respond
32    Private = 0,
33    /// Anyone can send, recipient may or may not respond  
34    Public = 9,
35}
36
37/// PQXDH prekey bundle containing both classical and post-quantum keys
38///
39/// This bundle contains all the cryptographic material needed for a client
40/// to initiate a PQXDH key agreement with the bundle owner.
41///
42/// ## Security Properties
43///
44/// - **Hybrid Security**: Combines X25519 (classical) and ML-KEM (post-quantum)
45/// - **Forward Secrecy**: One-time keys provide perfect forward secrecy
46/// - **Authentication**: All keys are signed by the identity key
47/// - **Key Rotation**: Signed prekeys are rotated periodically
48#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
49pub struct PqxdhPrekeyBundle {
50    // Classical ECDH keys
51    /// Medium-term X25519 public key for ECDH, rotated periodically
52    pub signed_prekey: x25519_dalek::PublicKey,
53    /// Signature over the signed prekey by the identity key
54    pub signed_prekey_signature: Signature,
55    /// Unique identifier for this signed prekey
56    pub signed_prekey_id: String,
57
58    /// One-time X25519 public keys (each used exactly once)
59    pub one_time_prekeys: BTreeMap<String, x25519_dalek::PublicKey>,
60
61    // Post-quantum KEM keys
62    /// Medium-term ML-KEM public key, rotated periodically  
63    pub pq_signed_prekey: Vec<u8>, // ML-KEM 768 public key bytes (1184 bytes)
64    /// Signature over the PQ signed prekey by the identity key
65    pub pq_signed_prekey_signature: Signature,
66    /// Unique identifier for this PQ signed prekey
67    pub pq_signed_prekey_id: String,
68
69    /// One-time ML-KEM public keys (each used exactly once)
70    pub pq_one_time_keys: BTreeMap<String, Vec<u8>>, // ML-KEM 768 public key bytes (1184 bytes each)
71    /// Signatures over each one-time PQ key by the identity key
72    pub pq_one_time_signatures: BTreeMap<String, Signature>,
73}
74
75/// PQXDH inbox for connecting to a user
76///
77#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
78pub struct PqxdhInbox {
79    /// Access control and responsiveness expectations
80    pub inbox_type: InboxType,
81    /// PQXDH prekeys for key agreement (always present)
82    pub pqxdh_prekeys: PqxdhPrekeyBundle,
83    /// Maximum echo payload size in bytes (None = unlimited)
84    pub max_echo_size: Option<u32>,
85    /// When this inbox expires (Unix timestamp)
86    pub expires_at: Option<u64>,
87}
88
89/// Private key material for PQXDH operations
90///
91/// This contains the private keys corresponding to a PqxdhPrekeyBundle.
92/// It should be stored securely and zeroized when no longer needed.
93#[derive(Clone, Serialize, Deserialize)]
94pub struct PqxdhPrivateKeys {
95    /// Private key for the signed X25519 prekey
96    pub signed_prekey_private: x25519_dalek::StaticSecret,
97    /// Private keys for one-time X25519 prekeys
98    pub one_time_prekey_privates: BTreeMap<String, x25519_dalek::StaticSecret>,
99    /// Private key for the signed ML-KEM prekey
100    pub pq_signed_prekey_private: Vec<u8>, // ML-KEM 768 private key bytes (2400 bytes)
101    /// Private keys for one-time ML-KEM prekeys  
102    pub pq_one_time_prekey_privates: BTreeMap<String, Vec<u8>>, // ML-KEM 768 private key bytes (2400 bytes each)
103}
104
105impl PartialEq for PqxdhPrivateKeys {
106    fn eq(&self, other: &Self) -> bool {
107        // For security reasons, we don't compare the actual private key values.
108        // Instead, we compare the structure (key IDs) to determine if they represent
109        // the same set of keys without exposing the private key material.
110        self.one_time_prekey_privates
111            .keys()
112            .collect::<std::collections::BTreeSet<_>>()
113            == other
114                .one_time_prekey_privates
115                .keys()
116                .collect::<std::collections::BTreeSet<_>>()
117            && self
118                .pq_one_time_prekey_privates
119                .keys()
120                .collect::<std::collections::BTreeSet<_>>()
121                == other
122                    .pq_one_time_prekey_privates
123                    .keys()
124                    .collect::<std::collections::BTreeSet<_>>()
125    }
126}
127
128/// PQXDH initial message sent to establish secure communication (Phase 2)
129///
130/// This message contains the initiator's ephemeral key, KEM ciphertext,
131/// prekey identifiers, and the initial encrypted payload. This establishes
132/// the shared secret and delivers the first message in one round-trip.
133#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
134pub struct PqxdhInitialMessage {
135    /// Initiator's identity key
136    pub initiator_identity: VerifyingKey,
137    /// Ephemeral X25519 public key generated for this session
138    pub ephemeral_key: x25519_dalek::PublicKey,
139    /// ML-KEM ciphertext encapsulating shared secret
140    pub kem_ciphertext: Vec<u8>,
141    /// ID of the signed prekey that was used
142    pub signed_prekey_id: String,
143    /// ID of the one-time prekey that was used (if any)
144    pub one_time_prekey_id: Option<String>,
145    /// ID of the PQ signed prekey that was used
146    pub pq_signed_prekey_id: String,
147    /// ID of the PQ one-time key that was used (if any)
148    pub pq_one_time_key_id: Option<String>,
149    /// Initial encrypted payload (typically the first RPC message)
150    pub encrypted_payload: Vec<u8>,
151}
152
153/// PQXDH session message for ongoing communication (Phase 3)
154///
155/// After the initial PQXDH handshake, follow-up messages use the established
156/// shared secret for AEAD encryption. This provides efficient ongoing communication.
157#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
158pub struct PqxdhSessionMessage {
159    /// Message sequence number (for replay protection)
160    pub sequence_number: u64,
161    /// AEAD encrypted payload using session keys
162    pub encrypted_payload: Vec<u8>,
163    /// AEAD authentication tag
164    pub auth_tag: [u8; 16],
165}
166
167/// Result of PQXDH key agreement
168#[derive(Debug, Clone, Zeroize, ZeroizeOnDrop, Serialize, Deserialize)]
169pub struct PqxdhSharedSecret {
170    /// 32-byte shared secret derived from PQXDH
171    pub shared_key: [u8; 32],
172    /// IDs of consumed one-time keys (to be deleted)
173    pub consumed_one_time_key_ids: Vec<String>,
174}
175
176/// Initial payload structure for PQXDH sessions
177///
178/// This structure is encrypted inside the PqxdhInitialMessage and contains
179/// both the user's initial payload and a randomized channel ID for subsequent
180/// session messages to provide unlinkability.
181#[derive(Serialize, Deserialize, Debug, Clone)]
182pub struct PqxdhInitialPayload<T> {
183    /// The actual user payload (e.g., RPC request)
184    pub user_payload: T,
185    /// randomized session id prefix to take and generate the target tags to listen for
186    pub session_channel_id_prefix: [u8; 32],
187}
188
189/// Errors that can occur during PQXDH operations
190#[derive(Debug, thiserror::Error)]
191pub enum PqxdhError {
192    #[error("Cryptographic operation failed: {0}")]
193    CryptoError(String),
194    #[error("Invalid prekey bundle: {0}")]
195    InvalidPrekeyBundle(String),
196    #[error("Missing required prekey: {0}")]
197    MissingPrekey(String),
198    #[error("Signature verification failed")]
199    SignatureVerificationFailed,
200    #[error("Key derivation failed: {0}")]
201    KeyDerivationFailed(String),
202    #[error("Serialization error: {0}")]
203    SerializationError(#[from] postcard::Error),
204}
205
206impl PqxdhPrekeyBundle {
207    /// Verify all signatures in the prekey bundle
208    ///
209    /// This checks that all prekeys are properly signed by the given identity key.
210    ///
211    /// ## Verification Process
212    ///
213    /// 1. Verifies `signed_prekey_signature` over `signed_prekey` bytes
214    /// 2. Verifies `pq_signed_prekey_signature` over `pq_signed_prekey` bytes  
215    /// 3. Verifies each `pq_one_time_signatures` over their respective `pq_one_time_keys`
216    /// 4. Ensures all PQ one-time keys have corresponding signatures (no missing or extra signatures)
217    ///
218    /// ## Security Properties
219    ///
220    /// - **Authentication**: Proves all keys were signed by the identity key holder
221    /// - **Integrity**: Detects any tampering with prekey data after signing
222    /// - **Completeness**: Ensures signature coverage matches key availability
223    ///
224    /// ## Returns
225    ///
226    /// - `Ok(())` if all signatures are valid and complete
227    /// - `Err(PqxdhError::SignatureVerificationFailed)` if any signature is invalid or missing
228    pub fn verify_signatures(
229        &self,
230        identity_key: &VerifyingKey,
231    ) -> std::result::Result<(), PqxdhError> {
232        // Verify signed prekey signature
233        identity_key
234            .verify(self.signed_prekey.as_bytes(), &self.signed_prekey_signature)
235            .map_err(|_| PqxdhError::SignatureVerificationFailed)?;
236
237        // Verify PQ signed prekey signature
238        identity_key
239            .verify(&self.pq_signed_prekey, &self.pq_signed_prekey_signature)
240            .map_err(|_| PqxdhError::SignatureVerificationFailed)?;
241
242        // Verify that every PQ one-time key has a corresponding signature
243        for (key_id, pq_key) in &self.pq_one_time_keys {
244            let signature = self.pq_one_time_signatures.get(key_id).ok_or_else(|| {
245                PqxdhError::InvalidPrekeyBundle(format!(
246                    "Missing signature for PQ one-time key: {key_id}"
247                ))
248            })?;
249
250            identity_key
251                .verify(pq_key, signature)
252                .map_err(|_| PqxdhError::SignatureVerificationFailed)?;
253        }
254
255        // Verify that every PQ one-time signature has a corresponding key (no extra signatures)
256        for key_id in self.pq_one_time_signatures.keys() {
257            if !self.pq_one_time_keys.contains_key(key_id) {
258                return Err(PqxdhError::InvalidPrekeyBundle(format!(
259                    "Extra signature found for non-existent PQ one-time key: {key_id}"
260                )));
261            }
262        }
263
264        Ok(())
265    }
266
267    /// Get the number of available one-time keys
268    pub fn one_time_key_count(&self) -> usize {
269        std::cmp::min(self.one_time_prekeys.len(), self.pq_one_time_keys.len())
270    }
271
272    /// Check if the bundle has any one-time keys available
273    pub fn has_one_time_keys(&self) -> bool {
274        !self.one_time_prekeys.is_empty() && !self.pq_one_time_keys.is_empty()
275    }
276}
277
278impl PqxdhInbox {
279    /// Create a new echo service inbox
280    pub fn new(
281        inbox_type: InboxType,
282        pqxdh_prekeys: PqxdhPrekeyBundle,
283        max_echo_size: Option<u32>,
284        expires_at: Option<u64>,
285    ) -> Self {
286        Self {
287            inbox_type,
288            pqxdh_prekeys,
289            max_echo_size,
290            expires_at,
291        }
292    }
293
294    /// Check if this inbox has expired
295    pub fn is_expired(&self, current_time: u64) -> bool {
296        self.expires_at.is_some_and(|expiry| current_time > expiry)
297    }
298
299    /// Check if a payload size is acceptable for this echo service
300    pub fn accepts_payload_size(&self, size: u32) -> bool {
301        self.max_echo_size.is_none_or(|max_size| size <= max_size)
302    }
303}
304
305#[cfg(test)]
306mod tests {
307    use super::*;
308    use crate::KeyPair;
309    use assert_matches::assert_matches;
310    use rand::rngs::OsRng;
311
312    /// Helper function to create a test keypair and identity
313    fn create_test_identity() -> (KeyPair, VerifyingKey) {
314        let keypair = KeyPair::generate_ed25519(&mut OsRng);
315        let identity = keypair.public_key();
316        (keypair, identity)
317    }
318
319    /// Helper function to create a prekey bundle with proper signatures
320    fn create_signed_prekey_bundle(identity_keypair: &KeyPair) -> PqxdhPrekeyBundle {
321        let signed_prekey = x25519_dalek::PublicKey::from([1u8; 32]);
322        let pq_signed_prekey = vec![2u8; 1184]; // ML-KEM 768 public key size
323
324        // Create proper signatures
325        let signed_prekey_signature = identity_keypair.sign(signed_prekey.as_bytes());
326        let pq_signed_prekey_signature = identity_keypair.sign(&pq_signed_prekey);
327
328        // Create one-time keys with signatures
329        let mut pq_one_time_keys = BTreeMap::new();
330        let mut pq_one_time_signatures = BTreeMap::new();
331
332        let pq_otk_1 = vec![3u8; 1184];
333        let pq_otk_2 = vec![4u8; 1184];
334
335        pq_one_time_keys.insert("pq_otk_1".to_string(), pq_otk_1.clone());
336        pq_one_time_keys.insert("pq_otk_2".to_string(), pq_otk_2.clone());
337
338        pq_one_time_signatures.insert("pq_otk_1".to_string(), identity_keypair.sign(&pq_otk_1));
339        pq_one_time_signatures.insert("pq_otk_2".to_string(), identity_keypair.sign(&pq_otk_2));
340
341        PqxdhPrekeyBundle {
342            signed_prekey,
343            signed_prekey_signature,
344            signed_prekey_id: "spk_001".to_string(),
345            one_time_prekeys: BTreeMap::new(),
346            pq_signed_prekey,
347            pq_signed_prekey_signature,
348            pq_signed_prekey_id: "pqspk_001".to_string(),
349            pq_one_time_keys,
350            pq_one_time_signatures,
351        }
352    }
353
354    /// Helper function to create a prekey bundle with invalid signatures
355    fn create_invalid_prekey_bundle() -> PqxdhPrekeyBundle {
356        // Create signatures with a different keypair (invalid)
357        let wrong_keypair = KeyPair::generate_ed25519(&mut OsRng);
358
359        let signed_prekey = x25519_dalek::PublicKey::from([1u8; 32]);
360        let pq_signed_prekey = vec![2u8; 1184];
361
362        // These signatures are from the wrong key, so they should fail verification
363        let signed_prekey_signature = wrong_keypair.sign(signed_prekey.as_bytes());
364        let pq_signed_prekey_signature = wrong_keypair.sign(&pq_signed_prekey);
365
366        let mut pq_one_time_keys = BTreeMap::new();
367        let mut pq_one_time_signatures = BTreeMap::new();
368
369        let pq_otk_1 = vec![3u8; 1184];
370        pq_one_time_keys.insert("pq_otk_1".to_string(), pq_otk_1.clone());
371        pq_one_time_signatures.insert("pq_otk_1".to_string(), wrong_keypair.sign(&pq_otk_1));
372
373        PqxdhPrekeyBundle {
374            signed_prekey,
375            signed_prekey_signature,
376            signed_prekey_id: "spk_001".to_string(),
377            one_time_prekeys: BTreeMap::new(),
378            pq_signed_prekey,
379            pq_signed_prekey_signature,
380            pq_signed_prekey_id: "pqspk_001".to_string(),
381            pq_one_time_keys,
382            pq_one_time_signatures,
383        }
384    }
385
386    #[test]
387    fn test_signature_verification_with_valid_signatures_should_pass() {
388        let (identity_keypair, identity_key) = create_test_identity();
389        let bundle = create_signed_prekey_bundle(&identity_keypair);
390
391        // This should pass when properly implemented
392        let result = bundle.verify_signatures(&identity_key);
393
394        // Currently this passes because the TODO just returns Ok(())
395        // But when implemented, it should still pass because signatures are valid
396        assert!(result.is_ok(), "Valid signatures should pass verification");
397    }
398
399    #[test]
400    fn test_signature_verification_with_invalid_signatures_should_fail() {
401        let (_identity_keypair, identity_key) = create_test_identity();
402        let bundle = create_invalid_prekey_bundle();
403
404        // This should fail because signatures are from wrong keypair
405        let result = bundle.verify_signatures(&identity_key);
406
407        assert!(
408            result.is_err(),
409            "Invalid signatures should fail verification"
410        );
411        assert_matches!(result.unwrap_err(), PqxdhError::SignatureVerificationFailed);
412    }
413
414    #[test]
415    fn test_signature_verification_with_wrong_identity_key() {
416        let (identity_keypair, _identity_key) = create_test_identity();
417        let bundle = create_signed_prekey_bundle(&identity_keypair);
418
419        // Create a different identity key
420        let (_wrong_keypair, wrong_identity_key) = create_test_identity();
421
422        // This should fail because we're using the wrong identity key
423        let result = bundle.verify_signatures(&wrong_identity_key);
424
425        assert!(
426            result.is_err(),
427            "Wrong identity key should fail verification"
428        );
429        assert!(matches!(
430            result.unwrap_err(),
431            PqxdhError::SignatureVerificationFailed
432        ));
433    }
434
435    #[test]
436    fn test_signature_verification_with_tampered_signed_prekey() {
437        let (identity_keypair, identity_key) = create_test_identity();
438        let mut bundle = create_signed_prekey_bundle(&identity_keypair);
439
440        // Tamper with the signed prekey after it was signed
441        bundle.signed_prekey = x25519_dalek::PublicKey::from([99u8; 32]);
442
443        let result = bundle.verify_signatures(&identity_key);
444
445        assert!(
446            result.is_err(),
447            "Tampered signed prekey should fail verification"
448        );
449        assert!(matches!(
450            result.unwrap_err(),
451            PqxdhError::SignatureVerificationFailed
452        ));
453    }
454
455    #[test]
456    fn test_signature_verification_with_tampered_pq_signed_prekey() {
457        let (identity_keypair, identity_key) = create_test_identity();
458        let mut bundle = create_signed_prekey_bundle(&identity_keypair);
459
460        // Tamper with the PQ signed prekey after it was signed
461        bundle.pq_signed_prekey = vec![99u8; 1184];
462
463        let result = bundle.verify_signatures(&identity_key);
464
465        assert!(
466            result.is_err(),
467            "Tampered PQ signed prekey should fail verification"
468        );
469        assert!(matches!(
470            result.unwrap_err(),
471            PqxdhError::SignatureVerificationFailed
472        ));
473    }
474
475    #[test]
476    fn test_signature_verification_with_tampered_pq_one_time_key() {
477        let (identity_keypair, identity_key) = create_test_identity();
478        let mut bundle = create_signed_prekey_bundle(&identity_keypair);
479
480        // Tamper with one of the PQ one-time keys after it was signed
481        bundle
482            .pq_one_time_keys
483            .insert("pq_otk_1".to_string(), vec![99u8; 1184]);
484
485        let result = bundle.verify_signatures(&identity_key);
486
487        assert!(
488            result.is_err(),
489            "Tampered PQ one-time key should fail verification"
490        );
491        assert!(matches!(
492            result.unwrap_err(),
493            PqxdhError::SignatureVerificationFailed
494        ));
495    }
496
497    #[test]
498    fn test_signature_verification_with_missing_pq_one_time_signature() {
499        let (identity_keypair, identity_key) = create_test_identity();
500        let mut bundle = create_signed_prekey_bundle(&identity_keypair);
501
502        // Remove a signature for an existing key
503        bundle.pq_one_time_signatures.remove("pq_otk_1");
504
505        let result = bundle.verify_signatures(&identity_key);
506
507        assert!(
508            result.is_err(),
509            "Missing PQ one-time signature should fail verification"
510        );
511        assert!(matches!(
512            result.unwrap_err(),
513            PqxdhError::InvalidPrekeyBundle(_)
514        ));
515    }
516
517    #[test]
518    fn test_signature_verification_with_extra_pq_one_time_signature() {
519        let (identity_keypair, identity_key) = create_test_identity();
520        let mut bundle = create_signed_prekey_bundle(&identity_keypair);
521
522        // Add a signature for a non-existent key
523        let fake_key = vec![88u8; 1184];
524        bundle
525            .pq_one_time_signatures
526            .insert("fake_key".to_string(), identity_keypair.sign(&fake_key));
527
528        let result = bundle.verify_signatures(&identity_key);
529
530        assert!(
531            result.is_err(),
532            "Extra PQ one-time signature should fail verification"
533        );
534        assert!(matches!(
535            result.unwrap_err(),
536            PqxdhError::InvalidPrekeyBundle(_)
537        ));
538    }
539
540    #[test]
541    fn test_signature_verification_with_empty_bundle() {
542        let (_identity_keypair, identity_key) = create_test_identity();
543
544        // Create a bundle with minimal required fields but no one-time keys
545        let bundle = PqxdhPrekeyBundle {
546            signed_prekey: x25519_dalek::PublicKey::from([0u8; 32]),
547            signed_prekey_signature: crate::Signature::Ed25519(Box::new(
548                ed25519_dalek::Signature::from_bytes(&[0u8; 64]),
549            )),
550            signed_prekey_id: "spk_001".to_string(),
551            one_time_prekeys: BTreeMap::new(),
552            pq_signed_prekey: vec![0u8; 1184],
553            pq_signed_prekey_signature: crate::Signature::Ed25519(Box::new(
554                ed25519_dalek::Signature::from_bytes(&[0u8; 64]),
555            )),
556            pq_signed_prekey_id: "pqspk_001".to_string(),
557            pq_one_time_keys: BTreeMap::new(),
558            pq_one_time_signatures: BTreeMap::new(),
559        };
560
561        let result = bundle.verify_signatures(&identity_key);
562
563        // This should fail because the signatures are dummy values (all zeros)
564        assert!(
565            result.is_err(),
566            "Invalid dummy signatures should fail verification"
567        );
568        assert!(matches!(
569            result.unwrap_err(),
570            PqxdhError::SignatureVerificationFailed
571        ));
572    }
573
574    #[test]
575    fn test_signature_verification_with_mismatched_signature_types() {
576        let (identity_keypair, identity_key) = create_test_identity();
577        let mut bundle = create_signed_prekey_bundle(&identity_keypair);
578
579        // Create a different keypair with ML-DSA and try to mix signature types
580        let mldsa_keypair = KeyPair::generate_ml_dsa65(&mut OsRng);
581        let mldsa_signature = mldsa_keypair.sign(bundle.signed_prekey.as_bytes());
582
583        // Replace Ed25519 signature with ML-DSA signature (type mismatch)
584        bundle.signed_prekey_signature = mldsa_signature;
585
586        let result = bundle.verify_signatures(&identity_key);
587
588        // This should fail because signature type doesn't match identity key type
589        // The VerifyingKey::verify method returns Ok(false) for mismatched types
590        assert!(
591            result.is_err(),
592            "Mismatched signature types should fail verification"
593        );
594        assert!(matches!(
595            result.unwrap_err(),
596            PqxdhError::SignatureVerificationFailed
597        ));
598    }
599
600    #[test]
601    fn test_inbox_type_serialization() {
602        let inbox_types = vec![InboxType::Private, InboxType::Public];
603
604        // Test postcard serialization
605        let serialized = postcard::to_stdvec(&inbox_types).unwrap();
606        let deserialized: Vec<InboxType> = postcard::from_bytes(&serialized).unwrap();
607
608        assert_eq!(inbox_types, deserialized);
609    }
610
611    #[test]
612    fn test_inbox_type_values() {
613        // Test that enum values match expected discriminants
614        assert_eq!(InboxType::Private as u8, 0);
615        assert_eq!(InboxType::Public as u8, 9);
616    }
617
618    #[test]
619    fn test_pqxdh_echo_service_inbox_creation() {
620        // Create a minimal prekey bundle for testing
621        let prekey_bundle = PqxdhPrekeyBundle {
622            signed_prekey: x25519_dalek::PublicKey::from([0u8; 32]),
623            signed_prekey_signature: crate::Signature::Ed25519(Box::new(
624                ed25519_dalek::Signature::from_bytes(&[0u8; 64]),
625            )),
626            signed_prekey_id: "spk_001".to_string(),
627            one_time_prekeys: BTreeMap::new(),
628            pq_signed_prekey: vec![0u8; 1184], // Placeholder ML-KEM 768 public key,
629            pq_signed_prekey_signature: crate::Signature::Ed25519(Box::new(
630                ed25519_dalek::Signature::from_bytes(&[0u8; 64]),
631            )),
632            pq_signed_prekey_id: "pqspk_001".to_string(),
633            pq_one_time_keys: BTreeMap::new(),
634            pq_one_time_signatures: BTreeMap::new(),
635        };
636
637        let inbox = PqxdhInbox::new(
638            InboxType::Public,
639            prekey_bundle,
640            Some(1024),
641            Some(1640995200),
642        );
643
644        assert_eq!(inbox.inbox_type, InboxType::Public);
645        assert_eq!(inbox.max_echo_size, Some(1024));
646        assert_eq!(inbox.expires_at, Some(1640995200));
647        assert!(inbox.accepts_payload_size(512));
648        assert!(!inbox.accepts_payload_size(2048));
649        assert!(inbox.is_expired(1640995201));
650        assert!(!inbox.is_expired(1640995199));
651    }
652
653    #[test]
654    fn test_prekey_bundle_one_time_key_count() {
655        let mut prekey_bundle = PqxdhPrekeyBundle {
656            signed_prekey: x25519_dalek::PublicKey::from([0u8; 32]),
657            signed_prekey_signature: crate::Signature::Ed25519(Box::new(
658                ed25519_dalek::Signature::from_bytes(&[0u8; 64]),
659            )),
660            signed_prekey_id: "spk_001".to_string(),
661            one_time_prekeys: BTreeMap::new(),
662            pq_signed_prekey: vec![0u8; 1184], // Placeholder ML-KEM 768 public key,
663            pq_signed_prekey_signature: crate::Signature::Ed25519(Box::new(
664                ed25519_dalek::Signature::from_bytes(&[0u8; 64]),
665            )),
666            pq_signed_prekey_id: "pqspk_001".to_string(),
667            pq_one_time_keys: BTreeMap::new(),
668            pq_one_time_signatures: BTreeMap::new(),
669        };
670
671        assert_eq!(prekey_bundle.one_time_key_count(), 0);
672        assert!(!prekey_bundle.has_one_time_keys());
673
674        // Add one X25519 key but no PQ key
675        prekey_bundle.one_time_prekeys.insert(
676            "otk_001".to_string(),
677            x25519_dalek::PublicKey::from([1u8; 32]),
678        );
679        assert_eq!(prekey_bundle.one_time_key_count(), 0); // Still 0 because we need both
680        assert!(!prekey_bundle.has_one_time_keys());
681
682        // Add one PQ key
683        prekey_bundle.pq_one_time_keys.insert(
684            "pqotk_001".to_string(),
685            vec![1u8; 1184], // Placeholder ML-KEM 768 public key
686        );
687        assert_eq!(prekey_bundle.one_time_key_count(), 1); // Now we have a pair
688        assert!(prekey_bundle.has_one_time_keys());
689    }
690
691    #[test]
692    fn test_pqxdh_structures_serialization() {
693        // Test that our main structures can be serialized with postcard
694        let prekey_bundle = PqxdhPrekeyBundle {
695            signed_prekey: x25519_dalek::PublicKey::from([0u8; 32]),
696            signed_prekey_signature: crate::Signature::Ed25519(Box::new(
697                ed25519_dalek::Signature::from_bytes(&[0u8; 64]),
698            )),
699            signed_prekey_id: "spk_001".to_string(),
700            one_time_prekeys: BTreeMap::new(),
701            pq_signed_prekey: vec![0u8; 1184], // Placeholder ML-KEM 768 public key,
702            pq_signed_prekey_signature: crate::Signature::Ed25519(Box::new(
703                ed25519_dalek::Signature::from_bytes(&[0u8; 64]),
704            )),
705            pq_signed_prekey_id: "pqspk_001".to_string(),
706            pq_one_time_keys: BTreeMap::new(),
707            pq_one_time_signatures: BTreeMap::new(),
708        };
709
710        let inbox = PqxdhInbox::new(InboxType::Private, prekey_bundle, None, None);
711
712        // Test serialization round-trip
713        let serialized = postcard::to_stdvec(&inbox).unwrap();
714        let deserialized: PqxdhInbox = postcard::from_bytes(&serialized).unwrap();
715
716        assert_eq!(inbox, deserialized);
717    }
718
719    #[test]
720    fn test_pqxdh_private_keys_random_serialization() {
721        use rand::RngCore;
722
723        // Generate random private keys
724        let mut rng = rand::thread_rng();
725
726        // Create random X25519 keys
727        let signed_prekey_private = x25519_dalek::StaticSecret::random_from_rng(&mut rng);
728        let mut one_time_prekey_privates = BTreeMap::new();
729        for i in 0..3 {
730            let key_id = format!("otk_{i}");
731            let private_key = x25519_dalek::StaticSecret::random_from_rng(&mut rng);
732            one_time_prekey_privates.insert(key_id, private_key);
733        }
734
735        // Create random ML-KEM keys (just random bytes for testing)
736        let mut pq_signed_prekey_private = vec![0u8; 2400]; // ML-KEM 768 private key size
737        rng.fill_bytes(&mut pq_signed_prekey_private);
738
739        let mut pq_one_time_prekey_privates = BTreeMap::new();
740        for i in 0..2 {
741            let key_id = format!("pq_otk_{i}");
742            let mut private_key = vec![0u8; 2400];
743            rng.fill_bytes(&mut private_key);
744            pq_one_time_prekey_privates.insert(key_id, private_key);
745        }
746
747        let private_keys = PqxdhPrivateKeys {
748            signed_prekey_private,
749            one_time_prekey_privates,
750            pq_signed_prekey_private,
751            pq_one_time_prekey_privates,
752        };
753
754        // Test serialization round-trip
755        let serialized = postcard::to_stdvec(&private_keys).unwrap();
756        let deserialized: PqxdhPrivateKeys = postcard::from_bytes(&serialized).unwrap();
757
758        // Verify the data is preserved
759        assert_eq!(
760            private_keys.signed_prekey_private.to_bytes(),
761            deserialized.signed_prekey_private.to_bytes()
762        );
763        assert_eq!(
764            private_keys.one_time_prekey_privates.len(),
765            deserialized.one_time_prekey_privates.len()
766        );
767        assert_eq!(
768            private_keys.pq_signed_prekey_private,
769            deserialized.pq_signed_prekey_private
770        );
771        assert_eq!(
772            private_keys.pq_one_time_prekey_privates,
773            deserialized.pq_one_time_prekey_privates
774        );
775
776        // Verify one-time keys
777        for (key_id, original_key) in &private_keys.one_time_prekey_privates {
778            let deserialized_key = &deserialized.one_time_prekey_privates[key_id];
779            assert_eq!(original_key.to_bytes(), deserialized_key.to_bytes());
780        }
781
782        // Test that we can generate different random keys
783        let signed_prekey_private2 = x25519_dalek::StaticSecret::random_from_rng(&mut rng);
784        assert_ne!(
785            private_keys.signed_prekey_private.to_bytes(),
786            signed_prekey_private2.to_bytes(),
787            "Random keys should be different"
788        );
789    }
790}
zoe_wire_protocol/inbox/pqxdh.rs

zoe_wire_protocol/inbox/
pqxdh.rs
