lib

Core libraries for Radroots
git clone https://radroots.dev/git/lib.git
Log | Files | Refs | README | LICENSE

ratchet.rs (57744B)


      1 use crate::error::RadrootsSimplexSmpCryptoError;
      2 use crate::message::{
      3     RADROOTS_SIMPLEX_SMP_NONCE_LENGTH, RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH, decrypt_padded,
      4     encrypt_padded,
      5 };
      6 use crate::official_ratchet::{
      7     RADROOTS_SIMPLEX_OFFICIAL_AES_IV_LENGTH, RADROOTS_SIMPLEX_OFFICIAL_E2E_CURRENT_VERSION,
      8     RADROOTS_SIMPLEX_OFFICIAL_SNTRUP761_PRIVATE_KEY_LENGTH,
      9     RADROOTS_SIMPLEX_OFFICIAL_SNTRUP761_PUBLIC_KEY_LENGTH,
     10     RADROOTS_SIMPLEX_OFFICIAL_X448_KEY_LENGTH, RadrootsSimplexOfficialAesGcmPayload,
     11     RadrootsSimplexOfficialEncryptedHeader, RadrootsSimplexOfficialEncryptedMessage,
     12     RadrootsSimplexOfficialMsgHeader, RadrootsSimplexOfficialSntrup761Keypair,
     13     RadrootsSimplexOfficialX3dhInit, decapsulate_official_sntrup761,
     14     decode_official_encrypted_header, decode_official_encrypted_message,
     15     decode_official_msg_header, derive_official_x448_shared_secret, encapsulate_official_sntrup761,
     16     encode_official_encrypted_header, encode_official_encrypted_message,
     17     encode_official_msg_header, generate_official_sntrup761_keypair,
     18     generate_official_x448_keypair, official_aes_gcm_decrypt_padded,
     19     official_aes_gcm_encrypt_padded, official_chain_kdf, official_ratchet_header_len,
     20     official_root_kdf,
     21 };
     22 use alloc::vec::Vec;
     23 use hkdf::Hkdf;
     24 use sha2::Sha512;
     25 
     26 const RADROOTS_SIMPLEX_AGENT_RATCHET_INFO: &[u8] = b"SimpleXAgentRatchetMessage";
     27 const RADROOTS_SIMPLEX_AGENT_RATCHET_OUTPUT_LENGTH: usize =
     28     RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH + RADROOTS_SIMPLEX_SMP_NONCE_LENGTH;
     29 const RADROOTS_SIMPLEX_OFFICIAL_MAX_SKIPPED_MESSAGES: u32 = 512;
     30 
     31 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
     32 pub enum RadrootsSimplexSmpRatchetRole {
     33     Initiator,
     34     Responder,
     35 }
     36 
     37 #[derive(Debug, Clone, PartialEq, Eq)]
     38 pub struct RadrootsSimplexSmpRatchetHeader {
     39     pub previous_sending_chain_length: u32,
     40     pub message_number: u32,
     41     pub dh_public_key: Vec<u8>,
     42     pub pq_public_key: Option<Vec<u8>>,
     43     pub pq_ciphertext: Option<Vec<u8>>,
     44 }
     45 
     46 impl RadrootsSimplexSmpRatchetHeader {
     47     pub fn validate(&self) -> Result<(), RadrootsSimplexSmpCryptoError> {
     48         if self.dh_public_key.is_empty() {
     49             return Err(RadrootsSimplexSmpCryptoError::MissingRatchetKey(
     50                 "dh_public_key",
     51             ));
     52         }
     53         if self.pq_ciphertext.is_some() && self.pq_public_key.is_none() {
     54             return Err(RadrootsSimplexSmpCryptoError::IncompletePqHeader);
     55         }
     56         Ok(())
     57     }
     58 }
     59 
     60 #[derive(Debug, Clone, PartialEq, Eq)]
     61 pub struct RadrootsSimplexSmpSkippedMessageKey {
     62     pub header_key: Vec<u8>,
     63     pub message_number: u32,
     64     pub message_key: Vec<u8>,
     65     pub message_iv: [u8; RADROOTS_SIMPLEX_OFFICIAL_AES_IV_LENGTH],
     66 }
     67 
     68 #[derive(Debug, Clone, PartialEq, Eq)]
     69 pub struct RadrootsSimplexSmpRatchetState {
     70     pub role: RadrootsSimplexSmpRatchetRole,
     71     pub root_epoch: u64,
     72     pub previous_sending_chain_length: u32,
     73     pub sending_chain_length: u32,
     74     pub receiving_chain_length: u32,
     75     pub local_dh_public_key: Vec<u8>,
     76     pub remote_dh_public_key: Vec<u8>,
     77     pub current_pq_public_key: Option<Vec<u8>>,
     78     pub remote_pq_public_key: Option<Vec<u8>>,
     79     pub pending_outbound_pq_ciphertext: Option<Vec<u8>>,
     80     pub pending_inbound_pq_ciphertext: Option<Vec<u8>>,
     81     pub current_pq_shared_secret: Option<Vec<u8>>,
     82     pub local_pq_private_key: Option<Vec<u8>>,
     83     pub local_dh_private_key: Option<Vec<u8>>,
     84     pub official_associated_data: Option<Vec<u8>>,
     85     pub official_root_key: Option<Vec<u8>>,
     86     pub official_sending_chain_key: Option<Vec<u8>>,
     87     pub official_receiving_chain_key: Option<Vec<u8>>,
     88     pub official_sending_header_key: Option<Vec<u8>>,
     89     pub official_receiving_header_key: Option<Vec<u8>>,
     90     pub official_next_sending_header_key: Option<Vec<u8>>,
     91     pub official_next_receiving_header_key: Option<Vec<u8>>,
     92     pub official_skipped_message_keys: Vec<RadrootsSimplexSmpSkippedMessageKey>,
     93 }
     94 
     95 impl RadrootsSimplexSmpRatchetState {
     96     pub fn initiator(
     97         local_dh_public_key: Vec<u8>,
     98         remote_dh_public_key: Vec<u8>,
     99         remote_pq_public_key: Option<Vec<u8>>,
    100     ) -> Result<Self, RadrootsSimplexSmpCryptoError> {
    101         validate_public_key(&local_dh_public_key)?;
    102         validate_public_key(&remote_dh_public_key)?;
    103         if let Some(key) = remote_pq_public_key.as_deref() {
    104             validate_public_key(key)?;
    105         }
    106 
    107         Ok(Self {
    108             role: RadrootsSimplexSmpRatchetRole::Initiator,
    109             root_epoch: 0,
    110             previous_sending_chain_length: 0,
    111             sending_chain_length: 0,
    112             receiving_chain_length: 0,
    113             local_dh_public_key,
    114             remote_dh_public_key,
    115             current_pq_public_key: None,
    116             remote_pq_public_key,
    117             pending_outbound_pq_ciphertext: None,
    118             pending_inbound_pq_ciphertext: None,
    119             current_pq_shared_secret: None,
    120             local_pq_private_key: None,
    121             local_dh_private_key: None,
    122             official_associated_data: None,
    123             official_root_key: None,
    124             official_sending_chain_key: None,
    125             official_receiving_chain_key: None,
    126             official_sending_header_key: None,
    127             official_receiving_header_key: None,
    128             official_next_sending_header_key: None,
    129             official_next_receiving_header_key: None,
    130             official_skipped_message_keys: Vec::new(),
    131         })
    132     }
    133 
    134     pub fn responder(
    135         local_dh_public_key: Vec<u8>,
    136         remote_dh_public_key: Vec<u8>,
    137         local_pq_public_key: Option<Vec<u8>>,
    138     ) -> Result<Self, RadrootsSimplexSmpCryptoError> {
    139         validate_public_key(&local_dh_public_key)?;
    140         validate_public_key(&remote_dh_public_key)?;
    141         if let Some(key) = local_pq_public_key.as_deref() {
    142             validate_public_key(key)?;
    143         }
    144 
    145         Ok(Self {
    146             role: RadrootsSimplexSmpRatchetRole::Responder,
    147             root_epoch: 0,
    148             previous_sending_chain_length: 0,
    149             sending_chain_length: 0,
    150             receiving_chain_length: 0,
    151             local_dh_public_key,
    152             remote_dh_public_key,
    153             current_pq_public_key: local_pq_public_key,
    154             remote_pq_public_key: None,
    155             pending_outbound_pq_ciphertext: None,
    156             pending_inbound_pq_ciphertext: None,
    157             current_pq_shared_secret: None,
    158             local_pq_private_key: None,
    159             local_dh_private_key: None,
    160             official_associated_data: None,
    161             official_root_key: None,
    162             official_sending_chain_key: None,
    163             official_receiving_chain_key: None,
    164             official_sending_header_key: None,
    165             official_receiving_header_key: None,
    166             official_next_sending_header_key: None,
    167             official_next_receiving_header_key: None,
    168             official_skipped_message_keys: Vec::new(),
    169         })
    170     }
    171 
    172     pub fn initialize_official_sender(
    173         &mut self,
    174         local_dh_private_key: Vec<u8>,
    175         init: RadrootsSimplexOfficialX3dhInit,
    176     ) -> Result<(), RadrootsSimplexSmpCryptoError> {
    177         validate_official_private_key(&local_dh_private_key)?;
    178         let root_dh =
    179             derive_official_x448_shared_secret(&local_dh_private_key, &self.remote_dh_public_key)?;
    180         let root = official_root_kdf(
    181             &init.ratchet_key,
    182             &root_dh,
    183             init.accepted_pq_shared_secret.as_deref(),
    184         )?;
    185         self.local_dh_private_key = Some(local_dh_private_key);
    186         self.official_associated_data = Some(init.associated_data);
    187         self.official_root_key = Some(root.root_key);
    188         self.official_sending_chain_key = Some(root.chain_key);
    189         self.official_receiving_chain_key = None;
    190         self.current_pq_shared_secret = init.accepted_pq_shared_secret;
    191         self.official_sending_header_key = Some(init.sending_header_key);
    192         self.official_receiving_header_key = None;
    193         self.official_next_sending_header_key = Some(root.next_header_key);
    194         self.official_next_receiving_header_key = Some(init.receiving_next_header_key);
    195         self.previous_sending_chain_length = 0;
    196         self.sending_chain_length = 0;
    197         self.receiving_chain_length = 0;
    198         self.root_epoch = 0;
    199         Ok(())
    200     }
    201 
    202     pub fn initialize_official_receiver(
    203         &mut self,
    204         local_dh_private_key: Vec<u8>,
    205         init: RadrootsSimplexOfficialX3dhInit,
    206     ) -> Result<(), RadrootsSimplexSmpCryptoError> {
    207         validate_official_private_key(&local_dh_private_key)?;
    208         self.local_dh_private_key = Some(local_dh_private_key);
    209         self.official_associated_data = Some(init.associated_data);
    210         self.official_root_key = Some(init.ratchet_key);
    211         self.current_pq_shared_secret = init.accepted_pq_shared_secret;
    212         self.official_sending_chain_key = None;
    213         self.official_receiving_chain_key = None;
    214         self.official_sending_header_key = None;
    215         self.official_receiving_header_key = None;
    216         self.official_next_sending_header_key = Some(init.receiving_next_header_key);
    217         self.official_next_receiving_header_key = Some(init.sending_header_key);
    218         self.previous_sending_chain_length = 0;
    219         self.sending_chain_length = 0;
    220         self.receiving_chain_length = 0;
    221         self.root_epoch = 0;
    222         Ok(())
    223     }
    224 
    225     pub fn stage_outbound_pq_step(
    226         &mut self,
    227         pq_public_key: Vec<u8>,
    228         pq_ciphertext: Vec<u8>,
    229         shared_secret: Vec<u8>,
    230     ) -> Result<(), RadrootsSimplexSmpCryptoError> {
    231         validate_public_key(&pq_public_key)?;
    232         if pq_ciphertext.is_empty() {
    233             return Err(RadrootsSimplexSmpCryptoError::InvalidCiphertextLength(0));
    234         }
    235         if shared_secret.is_empty() {
    236             return Err(RadrootsSimplexSmpCryptoError::InvalidSharedSecretLength(0));
    237         }
    238 
    239         self.current_pq_public_key = Some(pq_public_key);
    240         self.pending_outbound_pq_ciphertext = Some(pq_ciphertext);
    241         self.current_pq_shared_secret = Some(shared_secret);
    242         self.root_epoch = self.root_epoch.saturating_add(1);
    243         Ok(())
    244     }
    245 
    246     pub fn next_outbound_header(
    247         &mut self,
    248     ) -> Result<RadrootsSimplexSmpRatchetHeader, RadrootsSimplexSmpCryptoError> {
    249         validate_public_key(&self.local_dh_public_key)?;
    250         let header = RadrootsSimplexSmpRatchetHeader {
    251             previous_sending_chain_length: self.previous_sending_chain_length,
    252             message_number: self.sending_chain_length,
    253             dh_public_key: self.local_dh_public_key.clone(),
    254             pq_public_key: self.current_pq_public_key.clone(),
    255             pq_ciphertext: self.pending_outbound_pq_ciphertext.clone(),
    256         };
    257         header.validate()?;
    258         self.sending_chain_length = self.sending_chain_length.saturating_add(1);
    259         Ok(header)
    260     }
    261 
    262     pub fn apply_inbound_header(
    263         &mut self,
    264         header: &RadrootsSimplexSmpRatchetHeader,
    265         next_local_dh_public_key: Option<Vec<u8>>,
    266     ) -> Result<bool, RadrootsSimplexSmpCryptoError> {
    267         header.validate()?;
    268         let dh_advanced = header.dh_public_key != self.remote_dh_public_key;
    269 
    270         if dh_advanced {
    271             self.previous_sending_chain_length = self.sending_chain_length;
    272             self.sending_chain_length = 0;
    273             self.remote_dh_public_key = header.dh_public_key.clone();
    274             if let Some(next_local_key) = next_local_dh_public_key {
    275                 validate_public_key(&next_local_key)?;
    276                 self.local_dh_public_key = next_local_key;
    277             }
    278             self.root_epoch = self.root_epoch.saturating_add(1);
    279         } else if header.message_number < self.receiving_chain_length {
    280             return Err(RadrootsSimplexSmpCryptoError::RatchetMessageRegression {
    281                 received: header.message_number,
    282                 current: self.receiving_chain_length,
    283             });
    284         }
    285 
    286         self.receiving_chain_length = header.message_number.saturating_add(1);
    287         if let Some(public_key) = header.pq_public_key.as_ref() {
    288             self.remote_pq_public_key = Some(public_key.clone());
    289         }
    290         if let Some(ciphertext) = header.pq_ciphertext.as_ref() {
    291             self.pending_inbound_pq_ciphertext = Some(ciphertext.clone());
    292         }
    293 
    294         Ok(dh_advanced)
    295     }
    296 
    297     pub fn complete_inbound_pq_step(
    298         &mut self,
    299         shared_secret: Vec<u8>,
    300     ) -> Result<(), RadrootsSimplexSmpCryptoError> {
    301         if shared_secret.is_empty() {
    302             return Err(RadrootsSimplexSmpCryptoError::InvalidSharedSecretLength(0));
    303         }
    304         self.current_pq_shared_secret = Some(shared_secret);
    305         self.pending_inbound_pq_ciphertext = None;
    306         self.root_epoch = self.root_epoch.saturating_add(1);
    307         Ok(())
    308     }
    309 
    310     pub fn encrypt_payload(
    311         &mut self,
    312         shared_secret: &[u8],
    313         plaintext: &[u8],
    314         padded_len: usize,
    315     ) -> Result<(RadrootsSimplexSmpRatchetHeader, Vec<u8>), RadrootsSimplexSmpCryptoError> {
    316         let header = self.next_outbound_header()?;
    317         let associated_data = ratchet_header_associated_data(&header)?;
    318         let (message_key, nonce) = derive_ratchet_message_key(
    319             shared_secret,
    320             self.current_pq_shared_secret.as_deref(),
    321             self.root_epoch,
    322             &associated_data,
    323         )?;
    324         let ciphertext = encrypt_padded(&message_key, &nonce, plaintext, padded_len)?;
    325         Ok((header, ciphertext))
    326     }
    327 
    328     pub fn decrypt_payload(
    329         &mut self,
    330         shared_secret: &[u8],
    331         header: &RadrootsSimplexSmpRatchetHeader,
    332         ciphertext: &[u8],
    333     ) -> Result<Vec<u8>, RadrootsSimplexSmpCryptoError> {
    334         header.validate()?;
    335         if header.message_number < self.receiving_chain_length {
    336             return Err(RadrootsSimplexSmpCryptoError::RatchetMessageRegression {
    337                 received: header.message_number,
    338                 current: self.receiving_chain_length,
    339             });
    340         }
    341         let associated_data = ratchet_header_associated_data(header)?;
    342         let (message_key, nonce) = derive_ratchet_message_key(
    343             shared_secret,
    344             self.current_pq_shared_secret.as_deref(),
    345             self.root_epoch,
    346             &associated_data,
    347         )?;
    348         let plaintext = decrypt_padded(&message_key, &nonce, ciphertext)?;
    349         self.apply_inbound_header(header, None)?;
    350         Ok(plaintext)
    351     }
    352 
    353     pub fn encrypt_official_payload(
    354         &mut self,
    355         _shared_secret: &[u8],
    356         plaintext: &[u8],
    357         padded_len: usize,
    358     ) -> Result<Vec<u8>, RadrootsSimplexSmpCryptoError> {
    359         let message_number = self.sending_chain_length;
    360         let header = RadrootsSimplexSmpRatchetHeader {
    361             previous_sending_chain_length: self.previous_sending_chain_length,
    362             message_number,
    363             dh_public_key: self.local_dh_public_key.clone(),
    364             pq_public_key: self.current_pq_public_key.clone(),
    365             pq_ciphertext: self.pending_outbound_pq_ciphertext.clone(),
    366         };
    367         header.validate()?;
    368         let header_plaintext = encode_official_msg_header(
    369             RADROOTS_SIMPLEX_OFFICIAL_E2E_CURRENT_VERSION,
    370             &official_msg_header_from_ratchet_header(&header),
    371         )?;
    372         let ratchet_ad = self.official_associated_data.clone().ok_or(
    373             RadrootsSimplexSmpCryptoError::MissingRatchetKey("official_associated_data"),
    374         )?;
    375         let sending_header_key = self.official_sending_header_key.clone().ok_or(
    376             RadrootsSimplexSmpCryptoError::MissingRatchetKey("official_sending_header_key"),
    377         )?;
    378         let sending_chain_key = self.official_sending_chain_key.clone().ok_or(
    379             RadrootsSimplexSmpCryptoError::MissingRatchetKey("official_sending_chain_key"),
    380         )?;
    381         let chain = official_chain_kdf(&sending_chain_key)?;
    382         let header_payload = official_aes_gcm_encrypt_padded(
    383             &sending_header_key,
    384             &chain.header_iv,
    385             &header_plaintext,
    386             official_ratchet_header_len(
    387                 RADROOTS_SIMPLEX_OFFICIAL_E2E_CURRENT_VERSION,
    388                 self.pq_enabled(),
    389             )?,
    390             &ratchet_ad,
    391         )?;
    392         let encrypted_header = encode_official_encrypted_header(&official_encrypted_header(
    393             chain.header_iv,
    394             header_payload,
    395         )?)?;
    396         let message_ad = official_message_associated_data(&ratchet_ad, &encrypted_header);
    397         let message_payload = official_aes_gcm_encrypt_padded(
    398             &chain.message_key,
    399             &chain.message_iv,
    400             plaintext,
    401             padded_len,
    402             &message_ad,
    403         )?;
    404         self.official_sending_chain_key = Some(chain.chain_key);
    405         self.sending_chain_length = self.sending_chain_length.saturating_add(1);
    406         encode_official_encrypted_message(
    407             RADROOTS_SIMPLEX_OFFICIAL_E2E_CURRENT_VERSION,
    408             &RadrootsSimplexOfficialEncryptedMessage {
    409                 encrypted_header,
    410                 auth_tag: message_payload.auth_tag,
    411                 body: message_payload.ciphertext,
    412             },
    413         )
    414     }
    415 
    416     pub fn decrypt_official_payload(
    417         &mut self,
    418         _shared_secret: &[u8],
    419         encrypted_message: &[u8],
    420     ) -> Result<Vec<u8>, RadrootsSimplexSmpCryptoError> {
    421         let message = decode_official_encrypted_message(encrypted_message)?;
    422         let header = decode_official_encrypted_header(&message.encrypted_header)?;
    423         let ratchet_ad = self.official_associated_data.clone().ok_or(
    424             RadrootsSimplexSmpCryptoError::MissingRatchetKey("official_associated_data"),
    425         )?;
    426         if let Some(plaintext) =
    427             self.decrypt_official_skipped_payload(&header, &message, &ratchet_ad)?
    428         {
    429             return Ok(plaintext);
    430         }
    431         let (ratchet_step, ratchet_header) = self.decrypt_official_header(&header, &ratchet_ad)?;
    432         if ratchet_header.message_number < self.receiving_chain_length {
    433             return Err(RadrootsSimplexSmpCryptoError::RatchetMessageRegression {
    434                 received: ratchet_header.message_number,
    435                 current: self.receiving_chain_length,
    436             });
    437         }
    438         if ratchet_step == OfficialRatchetStep::Advance {
    439             self.skip_official_receiving_messages_until(
    440                 ratchet_header.previous_sending_chain_length,
    441             )?;
    442             self.advance_official_receiving_ratchet(&ratchet_header)?;
    443         }
    444         self.skip_official_receiving_messages_until(ratchet_header.message_number)?;
    445         let receiving_chain_key = self.official_receiving_chain_key.clone().ok_or(
    446             RadrootsSimplexSmpCryptoError::MissingRatchetKey("official_receiving_chain_key"),
    447         )?;
    448         let chain = official_chain_kdf(&receiving_chain_key)?;
    449         let message_ad = official_message_associated_data(&ratchet_ad, &message.encrypted_header);
    450         let plaintext = official_aes_gcm_decrypt_padded(
    451             &chain.message_key,
    452             &chain.message_iv,
    453             &RadrootsSimplexOfficialAesGcmPayload {
    454                 auth_tag: message.auth_tag,
    455                 ciphertext: message.body,
    456             },
    457             &message_ad,
    458         )?;
    459         self.official_receiving_chain_key = Some(chain.chain_key);
    460         self.apply_inbound_header(&ratchet_header, None)?;
    461         Ok(plaintext)
    462     }
    463 
    464     pub fn is_official_payload_replay(
    465         &self,
    466         encrypted_message: &[u8],
    467     ) -> Result<bool, RadrootsSimplexSmpCryptoError> {
    468         let message = decode_official_encrypted_message(encrypted_message)?;
    469         let header = decode_official_encrypted_header(&message.encrypted_header)?;
    470         let ratchet_ad = self.official_associated_data.clone().ok_or(
    471             RadrootsSimplexSmpCryptoError::MissingRatchetKey("official_associated_data"),
    472         )?;
    473         for skipped in &self.official_skipped_message_keys {
    474             if let Ok(ratchet_header) =
    475                 decrypt_official_header_with_key(&header, &skipped.header_key, &ratchet_ad)
    476                 && ratchet_header.message_number == skipped.message_number
    477             {
    478                 return Ok(false);
    479             }
    480         }
    481         if let Some(receiving_header_key) = self.official_receiving_header_key.as_ref()
    482             && let Ok(ratchet_header) =
    483                 decrypt_official_header_with_key(&header, receiving_header_key, &ratchet_ad)
    484         {
    485             return Ok(ratchet_header.message_number < self.receiving_chain_length);
    486         }
    487         if let Some(next_receiving_header_key) = self.official_next_receiving_header_key.as_ref()
    488             && let Ok(ratchet_header) =
    489                 decrypt_official_header_with_key(&header, next_receiving_header_key, &ratchet_ad)
    490         {
    491             return Ok(ratchet_header.message_number < self.receiving_chain_length
    492                 && ratchet_header.previous_sending_chain_length < self.receiving_chain_length);
    493         }
    494         Ok(false)
    495     }
    496 
    497     fn decrypt_official_skipped_payload(
    498         &mut self,
    499         header: &RadrootsSimplexOfficialEncryptedHeader,
    500         message: &RadrootsSimplexOfficialEncryptedMessage,
    501         ratchet_ad: &[u8],
    502     ) -> Result<Option<Vec<u8>>, RadrootsSimplexSmpCryptoError> {
    503         for skipped in self.official_skipped_message_keys.clone() {
    504             let Ok(ratchet_header) =
    505                 decrypt_official_header_with_key(header, &skipped.header_key, ratchet_ad)
    506             else {
    507                 continue;
    508             };
    509             if ratchet_header.message_number != skipped.message_number {
    510                 continue;
    511             }
    512             let position = self
    513                 .official_skipped_message_keys
    514                 .iter()
    515                 .position(|entry| {
    516                     entry.header_key == skipped.header_key
    517                         && entry.message_number == skipped.message_number
    518                 })
    519                 .ok_or(RadrootsSimplexSmpCryptoError::RatchetMessageRegression {
    520                     received: ratchet_header.message_number,
    521                     current: self.receiving_chain_length,
    522                 })?;
    523             let skipped = self.official_skipped_message_keys.remove(position);
    524             let message_ad =
    525                 official_message_associated_data(ratchet_ad, &message.encrypted_header);
    526             let plaintext = official_aes_gcm_decrypt_padded(
    527                 &skipped.message_key,
    528                 &skipped.message_iv,
    529                 &RadrootsSimplexOfficialAesGcmPayload {
    530                     auth_tag: message.auth_tag.clone(),
    531                     ciphertext: message.body.clone(),
    532                 },
    533                 &message_ad,
    534             )?;
    535             return Ok(Some(plaintext));
    536         }
    537         Ok(None)
    538     }
    539 
    540     fn skip_official_receiving_messages_until(
    541         &mut self,
    542         until_message_number: u32,
    543     ) -> Result<(), RadrootsSimplexSmpCryptoError> {
    544         if self.receiving_chain_length > until_message_number {
    545             return Err(RadrootsSimplexSmpCryptoError::RatchetMessageRegression {
    546                 received: until_message_number,
    547                 current: self.receiving_chain_length,
    548             });
    549         }
    550         let skipped = until_message_number.saturating_sub(self.receiving_chain_length);
    551         if skipped > RADROOTS_SIMPLEX_OFFICIAL_MAX_SKIPPED_MESSAGES {
    552             return Err(RadrootsSimplexSmpCryptoError::RatchetTooManySkipped {
    553                 skipped,
    554                 max: RADROOTS_SIMPLEX_OFFICIAL_MAX_SKIPPED_MESSAGES,
    555             });
    556         }
    557         if skipped == 0 {
    558             return Ok(());
    559         }
    560         let mut receiving_chain_key = self.official_receiving_chain_key.clone().ok_or(
    561             RadrootsSimplexSmpCryptoError::MissingRatchetKey("official_receiving_chain_key"),
    562         )?;
    563         let receiving_header_key = self.official_receiving_header_key.clone().ok_or(
    564             RadrootsSimplexSmpCryptoError::MissingRatchetKey("official_receiving_header_key"),
    565         )?;
    566         while self.receiving_chain_length < until_message_number {
    567             let chain = official_chain_kdf(&receiving_chain_key)?;
    568             self.official_skipped_message_keys
    569                 .push(RadrootsSimplexSmpSkippedMessageKey {
    570                     header_key: receiving_header_key.clone(),
    571                     message_number: self.receiving_chain_length,
    572                     message_key: chain.message_key,
    573                     message_iv: chain.message_iv,
    574                 });
    575             receiving_chain_key = chain.chain_key;
    576             self.receiving_chain_length = self.receiving_chain_length.saturating_add(1);
    577         }
    578         self.official_receiving_chain_key = Some(receiving_chain_key);
    579         Ok(())
    580     }
    581 
    582     fn decrypt_official_header(
    583         &self,
    584         header: &RadrootsSimplexOfficialEncryptedHeader,
    585         ratchet_ad: &[u8],
    586     ) -> Result<(OfficialRatchetStep, RadrootsSimplexSmpRatchetHeader), RadrootsSimplexSmpCryptoError>
    587     {
    588         if let Some(receiving_header_key) = self.official_receiving_header_key.as_ref() {
    589             if let Ok(ratchet_header) =
    590                 decrypt_official_header_with_key(header, receiving_header_key, ratchet_ad)
    591             {
    592                 return Ok((OfficialRatchetStep::Same, ratchet_header));
    593             }
    594         }
    595         let next_receiving_header_key = self.official_next_receiving_header_key.as_ref().ok_or(
    596             RadrootsSimplexSmpCryptoError::MissingRatchetKey("official_next_receiving_header_key"),
    597         )?;
    598         decrypt_official_header_with_key(header, next_receiving_header_key, ratchet_ad)
    599             .map(|ratchet_header| (OfficialRatchetStep::Advance, ratchet_header))
    600     }
    601 
    602     fn advance_official_receiving_ratchet(
    603         &mut self,
    604         header: &RadrootsSimplexSmpRatchetHeader,
    605     ) -> Result<(), RadrootsSimplexSmpCryptoError> {
    606         let pq_step = self.official_pq_receiving_step(header)?;
    607         let local_private_key = self.local_dh_private_key.clone().ok_or(
    608             RadrootsSimplexSmpCryptoError::MissingRatchetKey("local_dh_private_key"),
    609         )?;
    610         let root_key = self.official_root_key.clone().ok_or(
    611             RadrootsSimplexSmpCryptoError::MissingRatchetKey("official_root_key"),
    612         )?;
    613         let receiving_dh =
    614             derive_official_x448_shared_secret(&local_private_key, &header.dh_public_key)?;
    615         let receiving_root = official_root_kdf(
    616             &root_key,
    617             &receiving_dh,
    618             pq_step.receiving_shared_secret.as_deref(),
    619         )?;
    620         let next_local_keypair = generate_official_x448_keypair()?;
    621         let sending_dh = derive_official_x448_shared_secret(
    622             &next_local_keypair.private_key,
    623             &header.dh_public_key,
    624         )?;
    625         let sending_root = official_root_kdf(
    626             &receiving_root.root_key,
    627             &sending_dh,
    628             pq_step.sending_shared_secret.as_deref(),
    629         )?;
    630         self.previous_sending_chain_length = self.sending_chain_length;
    631         self.sending_chain_length = 0;
    632         self.receiving_chain_length = 0;
    633         self.remote_dh_public_key = header.dh_public_key.clone();
    634         self.remote_pq_public_key = header.pq_public_key.clone();
    635         self.pending_inbound_pq_ciphertext = header.pq_ciphertext.clone();
    636         if let Some(next_pq_keypair) = pq_step.next_local_keypair {
    637             self.current_pq_public_key = Some(next_pq_keypair.public_key);
    638             self.local_pq_private_key = Some(next_pq_keypair.private_key);
    639             self.pending_outbound_pq_ciphertext = pq_step.pending_outbound_pq_ciphertext;
    640             self.current_pq_shared_secret = pq_step.sending_shared_secret;
    641         } else if pq_step.receiving_shared_secret.is_some() {
    642             self.current_pq_shared_secret = pq_step.receiving_shared_secret;
    643         }
    644         self.local_dh_public_key = next_local_keypair.public_key;
    645         self.local_dh_private_key = Some(next_local_keypair.private_key);
    646         self.official_root_key = Some(sending_root.root_key);
    647         self.official_receiving_chain_key = Some(receiving_root.chain_key);
    648         self.official_receiving_header_key = self.official_next_receiving_header_key.take();
    649         self.official_next_receiving_header_key = Some(receiving_root.next_header_key);
    650         self.official_sending_chain_key = Some(sending_root.chain_key);
    651         self.official_sending_header_key = self.official_next_sending_header_key.take();
    652         self.official_next_sending_header_key = Some(sending_root.next_header_key);
    653         Ok(())
    654     }
    655 
    656     fn official_pq_receiving_step(
    657         &self,
    658         header: &RadrootsSimplexSmpRatchetHeader,
    659     ) -> Result<OfficialPqReceivingStep, RadrootsSimplexSmpCryptoError> {
    660         let Some(remote_pq_public_key) = header.pq_public_key.as_deref() else {
    661             return Ok(OfficialPqReceivingStep::default());
    662         };
    663         validate_official_pq_public_key(remote_pq_public_key)?;
    664         let receiving_shared_secret = match header.pq_ciphertext.as_deref() {
    665             Some(ciphertext) => {
    666                 let local_pq_private_key = self.local_pq_private_key.as_deref().ok_or(
    667                     RadrootsSimplexSmpCryptoError::MissingRatchetKey("local_pq_private_key"),
    668                 )?;
    669                 validate_official_pq_private_key(local_pq_private_key)?;
    670                 Some(decapsulate_official_sntrup761(
    671                     local_pq_private_key,
    672                     ciphertext,
    673                 )?)
    674             }
    675             None => None,
    676         };
    677         if !self.pq_enabled() {
    678             return Ok(OfficialPqReceivingStep {
    679                 receiving_shared_secret,
    680                 ..OfficialPqReceivingStep::default()
    681             });
    682         }
    683         let next_local_keypair = generate_official_sntrup761_keypair()?;
    684         let seed = random_official_pq_seed()?;
    685         let (pending_outbound_pq_ciphertext, sending_shared_secret) =
    686             encapsulate_official_sntrup761(remote_pq_public_key, &seed)?;
    687         Ok(OfficialPqReceivingStep {
    688             receiving_shared_secret,
    689             next_local_keypair: Some(next_local_keypair),
    690             pending_outbound_pq_ciphertext: Some(pending_outbound_pq_ciphertext),
    691             sending_shared_secret: Some(sending_shared_secret),
    692         })
    693     }
    694 
    695     fn pq_enabled(&self) -> bool {
    696         self.current_pq_public_key.is_some()
    697             || self.remote_pq_public_key.is_some()
    698             || self.current_pq_shared_secret.is_some()
    699             || self.local_pq_private_key.is_some()
    700     }
    701 }
    702 
    703 #[derive(Debug, Default, Clone, PartialEq, Eq)]
    704 struct OfficialPqReceivingStep {
    705     receiving_shared_secret: Option<Vec<u8>>,
    706     next_local_keypair: Option<RadrootsSimplexOfficialSntrup761Keypair>,
    707     pending_outbound_pq_ciphertext: Option<Vec<u8>>,
    708     sending_shared_secret: Option<Vec<u8>>,
    709 }
    710 
    711 fn validate_public_key(value: &[u8]) -> Result<(), RadrootsSimplexSmpCryptoError> {
    712     if value.is_empty() {
    713         return Err(RadrootsSimplexSmpCryptoError::InvalidPublicKeyLength(0));
    714     }
    715     Ok(())
    716 }
    717 
    718 fn validate_official_pq_public_key(value: &[u8]) -> Result<(), RadrootsSimplexSmpCryptoError> {
    719     if value.len() != RADROOTS_SIMPLEX_OFFICIAL_SNTRUP761_PUBLIC_KEY_LENGTH {
    720         return Err(RadrootsSimplexSmpCryptoError::InvalidPqKeyLength(
    721             value.len(),
    722         ));
    723     }
    724     Ok(())
    725 }
    726 
    727 fn validate_official_pq_private_key(value: &[u8]) -> Result<(), RadrootsSimplexSmpCryptoError> {
    728     if value.len() != RADROOTS_SIMPLEX_OFFICIAL_SNTRUP761_PRIVATE_KEY_LENGTH {
    729         return Err(RadrootsSimplexSmpCryptoError::InvalidPrivateKeyLength(
    730             value.len(),
    731         ));
    732     }
    733     Ok(())
    734 }
    735 
    736 fn random_official_pq_seed() -> Result<[u8; 32], RadrootsSimplexSmpCryptoError> {
    737     let mut seed = [0_u8; 32];
    738     getrandom::getrandom(&mut seed)
    739         .map_err(|_| RadrootsSimplexSmpCryptoError::EntropyUnavailable)?;
    740     Ok(seed)
    741 }
    742 
    743 fn validate_official_private_key(value: &[u8]) -> Result<(), RadrootsSimplexSmpCryptoError> {
    744     if value.len() != RADROOTS_SIMPLEX_OFFICIAL_X448_KEY_LENGTH {
    745         return Err(RadrootsSimplexSmpCryptoError::InvalidPrivateKeyLength(
    746             value.len(),
    747         ));
    748     }
    749     Ok(())
    750 }
    751 
    752 fn derive_ratchet_message_key(
    753     shared_secret: &[u8],
    754     pq_shared_secret: Option<&[u8]>,
    755     root_epoch: u64,
    756     associated_data: &[u8],
    757 ) -> Result<(Vec<u8>, [u8; RADROOTS_SIMPLEX_SMP_NONCE_LENGTH]), RadrootsSimplexSmpCryptoError> {
    758     let mut ikm = Vec::with_capacity(shared_secret.len() + pq_shared_secret.map_or(0, <[u8]>::len));
    759     ikm.extend_from_slice(shared_secret);
    760     if let Some(secret) = pq_shared_secret {
    761         ikm.extend_from_slice(secret);
    762     }
    763     let mut salt = Vec::with_capacity(8 + associated_data.len());
    764     salt.extend_from_slice(&root_epoch.to_be_bytes());
    765     salt.extend_from_slice(associated_data);
    766     let hkdf = Hkdf::<Sha512>::new(Some(&salt), &ikm);
    767     let mut output = [0_u8; RADROOTS_SIMPLEX_AGENT_RATCHET_OUTPUT_LENGTH];
    768     hkdf.expand(RADROOTS_SIMPLEX_AGENT_RATCHET_INFO, &mut output)
    769         .map_err(|_| {
    770             RadrootsSimplexSmpCryptoError::InvalidKeyDerivationLength(
    771                 RADROOTS_SIMPLEX_AGENT_RATCHET_OUTPUT_LENGTH,
    772             )
    773         })?;
    774     let mut nonce = [0_u8; RADROOTS_SIMPLEX_SMP_NONCE_LENGTH];
    775     nonce.copy_from_slice(&output[RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH..]);
    776     Ok((
    777         output[..RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH].to_vec(),
    778         nonce,
    779     ))
    780 }
    781 
    782 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    783 enum OfficialRatchetStep {
    784     Same,
    785     Advance,
    786 }
    787 
    788 fn decrypt_official_header_with_key(
    789     header: &RadrootsSimplexOfficialEncryptedHeader,
    790     header_key: &[u8],
    791     ratchet_ad: &[u8],
    792 ) -> Result<RadrootsSimplexSmpRatchetHeader, RadrootsSimplexSmpCryptoError> {
    793     let header_plaintext = official_aes_gcm_decrypt_padded(
    794         header_key,
    795         &header.iv,
    796         &RadrootsSimplexOfficialAesGcmPayload {
    797             auth_tag: header.auth_tag.clone(),
    798             ciphertext: header.body.clone(),
    799         },
    800         ratchet_ad,
    801     )?;
    802     Ok(ratchet_header_from_official_msg_header(
    803         decode_official_msg_header(header.version, &header_plaintext)?,
    804     ))
    805 }
    806 
    807 fn official_encrypted_header(
    808     iv: [u8; crate::official_ratchet::RADROOTS_SIMPLEX_OFFICIAL_AES_IV_LENGTH],
    809     payload: RadrootsSimplexOfficialAesGcmPayload,
    810 ) -> Result<RadrootsSimplexOfficialEncryptedHeader, RadrootsSimplexSmpCryptoError> {
    811     Ok(RadrootsSimplexOfficialEncryptedHeader {
    812         version: RADROOTS_SIMPLEX_OFFICIAL_E2E_CURRENT_VERSION,
    813         iv,
    814         auth_tag: payload.auth_tag,
    815         body: payload.ciphertext,
    816     })
    817 }
    818 
    819 fn official_message_associated_data(ratchet_ad: &[u8], encrypted_header: &[u8]) -> Vec<u8> {
    820     let mut associated_data = Vec::with_capacity(ratchet_ad.len() + encrypted_header.len());
    821     associated_data.extend_from_slice(ratchet_ad);
    822     associated_data.extend_from_slice(encrypted_header);
    823     associated_data
    824 }
    825 
    826 fn official_msg_header_from_ratchet_header(
    827     header: &RadrootsSimplexSmpRatchetHeader,
    828 ) -> RadrootsSimplexOfficialMsgHeader {
    829     RadrootsSimplexOfficialMsgHeader {
    830         max_version: RADROOTS_SIMPLEX_OFFICIAL_E2E_CURRENT_VERSION,
    831         dh_public_key: header.dh_public_key.clone(),
    832         pq_public_key: header.pq_public_key.clone(),
    833         pq_ciphertext: header.pq_ciphertext.clone(),
    834         previous_sending_chain_length: header.previous_sending_chain_length,
    835         message_number: header.message_number,
    836     }
    837 }
    838 
    839 fn ratchet_header_from_official_msg_header(
    840     header: RadrootsSimplexOfficialMsgHeader,
    841 ) -> RadrootsSimplexSmpRatchetHeader {
    842     RadrootsSimplexSmpRatchetHeader {
    843         previous_sending_chain_length: header.previous_sending_chain_length,
    844         message_number: header.message_number,
    845         dh_public_key: header.dh_public_key,
    846         pq_public_key: header.pq_public_key,
    847         pq_ciphertext: header.pq_ciphertext,
    848     }
    849 }
    850 
    851 fn ratchet_header_associated_data(
    852     header: &RadrootsSimplexSmpRatchetHeader,
    853 ) -> Result<Vec<u8>, RadrootsSimplexSmpCryptoError> {
    854     let mut buffer = Vec::new();
    855     buffer.extend_from_slice(&header.previous_sending_chain_length.to_be_bytes());
    856     buffer.extend_from_slice(&header.message_number.to_be_bytes());
    857     push_large_bytes(&mut buffer, &header.dh_public_key)?;
    858     push_maybe_large_bytes(&mut buffer, header.pq_public_key.as_deref())?;
    859     push_maybe_large_bytes(&mut buffer, header.pq_ciphertext.as_deref())?;
    860     Ok(buffer)
    861 }
    862 
    863 fn push_maybe_large_bytes(
    864     buffer: &mut Vec<u8>,
    865     value: Option<&[u8]>,
    866 ) -> Result<(), RadrootsSimplexSmpCryptoError> {
    867     match value {
    868         Some(value) => {
    869             buffer.push(1);
    870             push_large_bytes(buffer, value)
    871         }
    872         None => {
    873             buffer.push(0);
    874             Ok(())
    875         }
    876     }
    877 }
    878 
    879 fn push_large_bytes(
    880     buffer: &mut Vec<u8>,
    881     value: &[u8],
    882 ) -> Result<(), RadrootsSimplexSmpCryptoError> {
    883     if value.len() > u16::MAX as usize {
    884         return Err(RadrootsSimplexSmpCryptoError::InvalidMessageLength {
    885             actual: value.len(),
    886             padded: u16::MAX as usize,
    887         });
    888     }
    889     buffer.extend_from_slice(&(value.len() as u16).to_be_bytes());
    890     buffer.extend_from_slice(value);
    891     Ok(())
    892 }
    893 
    894 #[cfg(test)]
    895 mod tests {
    896     use super::*;
    897     use crate::official_ratchet::{
    898         RADROOTS_SIMPLEX_OFFICIAL_E2E_CURRENT_VERSION, RADROOTS_SIMPLEX_OFFICIAL_E2E_KDF_VERSION,
    899         RadrootsSimplexOfficialX3dhParams, decode_official_encrypted_header,
    900         decode_official_encrypted_message, official_sntrup761_keypair_from_seed,
    901         official_x3dh_receiver_init, official_x3dh_receiver_init_accepting_pq,
    902         official_x3dh_sender_init, official_x3dh_sender_init_accepting_pq,
    903         official_x448_keypair_from_seed,
    904     };
    905     use radroots_simplex_smp_proto::prelude::RadrootsSimplexSmpVersionRange;
    906 
    907     fn official_sender_receiver_ratchets() -> (
    908         RadrootsSimplexSmpRatchetState,
    909         RadrootsSimplexSmpRatchetState,
    910     ) {
    911         let receiver_key_1 = official_x448_keypair_from_seed(b"rr-synth-ratchet-rcv-1");
    912         let receiver_key_2 = official_x448_keypair_from_seed(b"rr-synth-ratchet-rcv-2");
    913         let sender_key_1 = official_x448_keypair_from_seed(b"rr-synth-ratchet-snd-1");
    914         let sender_key_2 = official_x448_keypair_from_seed(b"rr-synth-ratchet-snd-2");
    915         let receiver_params = RadrootsSimplexOfficialX3dhParams {
    916             version_range: RadrootsSimplexSmpVersionRange::new(
    917                 RADROOTS_SIMPLEX_OFFICIAL_E2E_KDF_VERSION,
    918                 RADROOTS_SIMPLEX_OFFICIAL_E2E_CURRENT_VERSION,
    919             )
    920             .unwrap(),
    921             key_1: receiver_key_1.public_key.clone(),
    922             key_2: receiver_key_2.public_key.clone(),
    923             pq_public_key: None,
    924             pq_ciphertext: None,
    925         };
    926         let sender_params = RadrootsSimplexOfficialX3dhParams {
    927             version_range: receiver_params.version_range,
    928             key_1: sender_key_1.public_key.clone(),
    929             key_2: sender_key_2.public_key.clone(),
    930             pq_public_key: None,
    931             pq_ciphertext: None,
    932         };
    933         let sender_init =
    934             official_x3dh_sender_init(&sender_key_1, &sender_key_2, &receiver_params).unwrap();
    935         let receiver_init =
    936             official_x3dh_receiver_init(&receiver_key_1, &receiver_key_2, &sender_params).unwrap();
    937         let mut sender = RadrootsSimplexSmpRatchetState::responder(
    938             sender_key_2.public_key.clone(),
    939             receiver_key_2.public_key.clone(),
    940             None,
    941         )
    942         .unwrap();
    943         sender
    944             .initialize_official_sender(sender_key_2.private_key, sender_init)
    945             .unwrap();
    946         let mut receiver = RadrootsSimplexSmpRatchetState::initiator(
    947             receiver_key_2.public_key.clone(),
    948             receiver_key_1.public_key.clone(),
    949             None,
    950         )
    951         .unwrap();
    952         receiver
    953             .initialize_official_receiver(receiver_key_2.private_key, receiver_init)
    954             .unwrap();
    955         (sender, receiver)
    956     }
    957 
    958     fn official_pq_sender_receiver_ratchets() -> (
    959         RadrootsSimplexSmpRatchetState,
    960         RadrootsSimplexSmpRatchetState,
    961     ) {
    962         let receiver_key_1 = official_x448_keypair_from_seed(b"rr-synth-pq-ratchet-rcv-1");
    963         let receiver_key_2 = official_x448_keypair_from_seed(b"rr-synth-pq-ratchet-rcv-2");
    964         let receiver_pq_keypair = official_sntrup761_keypair_from_seed(b"rr-synth-pq-rcv-kem");
    965         let sender_key_1 = official_x448_keypair_from_seed(b"rr-synth-pq-ratchet-snd-1");
    966         let sender_key_2 = official_x448_keypair_from_seed(b"rr-synth-pq-ratchet-snd-2");
    967         let sender_pq_keypair = official_sntrup761_keypair_from_seed(b"rr-synth-pq-snd-kem");
    968         let receiver_params = RadrootsSimplexOfficialX3dhParams {
    969             version_range: RadrootsSimplexSmpVersionRange::new(
    970                 RADROOTS_SIMPLEX_OFFICIAL_E2E_KDF_VERSION,
    971                 RADROOTS_SIMPLEX_OFFICIAL_E2E_CURRENT_VERSION,
    972             )
    973             .unwrap(),
    974             key_1: receiver_key_1.public_key.clone(),
    975             key_2: receiver_key_2.public_key.clone(),
    976             pq_public_key: Some(receiver_pq_keypair.public_key.clone()),
    977             pq_ciphertext: None,
    978         };
    979         let sender_init = official_x3dh_sender_init_accepting_pq(
    980             &sender_key_1,
    981             &sender_key_2,
    982             sender_pq_keypair,
    983             &receiver_params,
    984             b"rr-synth-pq-x3dh-accept",
    985         )
    986         .unwrap();
    987         let receiver_init = official_x3dh_receiver_init_accepting_pq(
    988             &receiver_key_1,
    989             &receiver_key_2,
    990             &receiver_pq_keypair,
    991             &sender_init.sender_params,
    992         )
    993         .unwrap();
    994         let mut sender = RadrootsSimplexSmpRatchetState::responder(
    995             sender_key_2.public_key.clone(),
    996             receiver_key_2.public_key.clone(),
    997             sender_init.sender_params.pq_public_key.clone(),
    998         )
    999         .unwrap();
   1000         sender
   1001             .initialize_official_sender(sender_key_2.private_key, sender_init.init)
   1002             .unwrap();
   1003         sender.current_pq_public_key = sender_init.sender_params.pq_public_key.clone();
   1004         sender.pending_outbound_pq_ciphertext = sender_init.sender_params.pq_ciphertext.clone();
   1005         sender.local_pq_private_key = Some(sender_init.local_pq_keypair.private_key);
   1006         let mut receiver = RadrootsSimplexSmpRatchetState::initiator(
   1007             receiver_key_2.public_key.clone(),
   1008             receiver_key_1.public_key.clone(),
   1009             None,
   1010         )
   1011         .unwrap();
   1012         receiver.current_pq_public_key = Some(receiver_pq_keypair.public_key);
   1013         receiver.local_pq_private_key = Some(receiver_pq_keypair.private_key);
   1014         receiver
   1015             .initialize_official_receiver(receiver_key_2.private_key, receiver_init.init)
   1016             .unwrap();
   1017         (sender, receiver)
   1018     }
   1019 
   1020     #[test]
   1021     fn stages_outbound_pq_state_and_emits_header() {
   1022         let mut state = RadrootsSimplexSmpRatchetState::responder(
   1023             b"bob-dh".to_vec(),
   1024             b"alice-dh".to_vec(),
   1025             Some(b"bob-pq".to_vec()),
   1026         )
   1027         .unwrap();
   1028         state
   1029             .stage_outbound_pq_step(
   1030                 b"bob-pq-next".to_vec(),
   1031                 b"ciphertext".to_vec(),
   1032                 b"shared-secret".to_vec(),
   1033             )
   1034             .unwrap();
   1035 
   1036         let header = state.next_outbound_header().unwrap();
   1037         assert_eq!(header.message_number, 0);
   1038         assert_eq!(header.pq_public_key, Some(b"bob-pq-next".to_vec()));
   1039         assert_eq!(header.pq_ciphertext, Some(b"ciphertext".to_vec()));
   1040         assert_eq!(state.sending_chain_length, 1);
   1041     }
   1042 
   1043     #[test]
   1044     fn applies_inbound_dh_and_pq_transition() {
   1045         let mut state = RadrootsSimplexSmpRatchetState::initiator(
   1046             b"alice-dh".to_vec(),
   1047             b"bob-dh".to_vec(),
   1048             Some(b"bob-pq".to_vec()),
   1049         )
   1050         .unwrap();
   1051         state.sending_chain_length = 4;
   1052 
   1053         let advanced = state
   1054             .apply_inbound_header(
   1055                 &RadrootsSimplexSmpRatchetHeader {
   1056                     previous_sending_chain_length: 2,
   1057                     message_number: 0,
   1058                     dh_public_key: b"bob-dh-next".to_vec(),
   1059                     pq_public_key: Some(b"bob-pq-next".to_vec()),
   1060                     pq_ciphertext: Some(b"ciphertext".to_vec()),
   1061                 },
   1062                 Some(b"alice-dh-next".to_vec()),
   1063             )
   1064             .unwrap();
   1065 
   1066         assert!(advanced);
   1067         assert_eq!(state.previous_sending_chain_length, 4);
   1068         assert_eq!(state.sending_chain_length, 0);
   1069         assert_eq!(state.receiving_chain_length, 1);
   1070         assert_eq!(state.remote_pq_public_key, Some(b"bob-pq-next".to_vec()));
   1071         assert_eq!(
   1072             state.pending_inbound_pq_ciphertext,
   1073             Some(b"ciphertext".to_vec())
   1074         );
   1075 
   1076         state
   1077             .complete_inbound_pq_step(b"shared-secret".to_vec())
   1078             .unwrap();
   1079         assert_eq!(
   1080             state.current_pq_shared_secret,
   1081             Some(b"shared-secret".to_vec())
   1082         );
   1083         assert_eq!(state.pending_inbound_pq_ciphertext, None);
   1084     }
   1085 
   1086     #[test]
   1087     fn rejects_incomplete_pq_header() {
   1088         let header = RadrootsSimplexSmpRatchetHeader {
   1089             previous_sending_chain_length: 0,
   1090             message_number: 0,
   1091             dh_public_key: b"dh".to_vec(),
   1092             pq_public_key: None,
   1093             pq_ciphertext: Some(b"ciphertext".to_vec()),
   1094         };
   1095 
   1096         let error = header.validate().unwrap_err();
   1097         assert_eq!(error, RadrootsSimplexSmpCryptoError::IncompletePqHeader);
   1098     }
   1099 
   1100     #[test]
   1101     fn encrypts_payload_and_advances_receive_state() {
   1102         let mut sender =
   1103             RadrootsSimplexSmpRatchetState::initiator(vec![1_u8; 56], vec![2_u8; 56], None)
   1104                 .unwrap();
   1105         let mut receiver =
   1106             RadrootsSimplexSmpRatchetState::responder(vec![2_u8; 56], vec![1_u8; 56], None)
   1107                 .unwrap();
   1108         let shared_secret = [7_u8; RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH];
   1109 
   1110         let (header, ciphertext) = sender
   1111             .encrypt_payload(&shared_secret, b"agent body", 64)
   1112             .unwrap();
   1113 
   1114         assert_ne!(ciphertext, b"agent body");
   1115         let plaintext = receiver
   1116             .decrypt_payload(&shared_secret, &header, &ciphertext)
   1117             .unwrap();
   1118         assert_eq!(plaintext, b"agent body");
   1119         assert_eq!(receiver.receiving_chain_length, 1);
   1120     }
   1121 
   1122     #[test]
   1123     fn rejects_tampered_ratchet_header() {
   1124         let mut sender =
   1125             RadrootsSimplexSmpRatchetState::initiator(vec![1_u8; 56], vec![2_u8; 56], None)
   1126                 .unwrap();
   1127         let mut receiver =
   1128             RadrootsSimplexSmpRatchetState::responder(vec![2_u8; 56], vec![1_u8; 56], None)
   1129                 .unwrap();
   1130         let shared_secret = [9_u8; RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH];
   1131         let (mut header, ciphertext) = sender
   1132             .encrypt_payload(&shared_secret, b"agent body", 64)
   1133             .unwrap();
   1134         header.message_number = header.message_number.saturating_add(1);
   1135 
   1136         let error = receiver
   1137             .decrypt_payload(&shared_secret, &header, &ciphertext)
   1138             .unwrap_err();
   1139         assert!(matches!(
   1140             error,
   1141             RadrootsSimplexSmpCryptoError::InvalidCiphertextLength(_)
   1142         ));
   1143     }
   1144 
   1145     #[test]
   1146     fn stages_large_pq_material_in_header() {
   1147         let mut sender = RadrootsSimplexSmpRatchetState::initiator(
   1148             b"alice-dh".to_vec(),
   1149             b"bob-dh".to_vec(),
   1150             None,
   1151         )
   1152         .unwrap();
   1153         sender
   1154             .stage_outbound_pq_step(vec![1_u8; 1158], vec![2_u8; 1039], vec![3_u8; 32])
   1155             .unwrap();
   1156 
   1157         let header = sender.next_outbound_header().unwrap();
   1158         assert_eq!(header.pq_public_key.as_ref().unwrap().len(), 1158);
   1159         assert_eq!(header.pq_ciphertext.as_ref().unwrap().len(), 1039);
   1160         assert!(ratchet_header_associated_data(&header).unwrap().len() > 2200);
   1161     }
   1162 
   1163     #[test]
   1164     fn encrypts_official_payload_as_opaque_message() {
   1165         let (mut sender, mut receiver) = official_sender_receiver_ratchets();
   1166         let shared_secret = [11_u8; RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH];
   1167 
   1168         let encrypted = sender
   1169             .encrypt_official_payload(&shared_secret, b"official agent body", 96)
   1170             .unwrap();
   1171         assert_ne!(encrypted, b"official agent body");
   1172         assert_eq!(encrypted.len(), 2 + 124 + 16 + 96);
   1173 
   1174         let plaintext = receiver
   1175             .decrypt_official_payload(&shared_secret, &encrypted)
   1176             .unwrap();
   1177         assert_eq!(plaintext, b"official agent body");
   1178         assert_eq!(receiver.receiving_chain_length, 1);
   1179     }
   1180 
   1181     #[test]
   1182     fn detects_official_payload_replay_without_consuming_skipped_messages() {
   1183         let (mut sender, mut receiver) = official_sender_receiver_ratchets();
   1184         let shared_secret = [14_u8; RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH];
   1185         let first = sender
   1186             .encrypt_official_payload(&shared_secret, b"first", 96)
   1187             .unwrap();
   1188         let second = sender
   1189             .encrypt_official_payload(&shared_secret, b"second", 96)
   1190             .unwrap();
   1191         let third = sender
   1192             .encrypt_official_payload(&shared_secret, b"third", 96)
   1193             .unwrap();
   1194 
   1195         assert_eq!(
   1196             receiver
   1197                 .decrypt_official_payload(&shared_secret, &second)
   1198                 .unwrap(),
   1199             b"second"
   1200         );
   1201         assert!(!receiver.is_official_payload_replay(&first).unwrap());
   1202         assert_eq!(
   1203             receiver
   1204                 .decrypt_official_payload(&shared_secret, &third)
   1205                 .unwrap(),
   1206             b"third"
   1207         );
   1208         assert_eq!(
   1209             receiver
   1210                 .decrypt_official_payload(&shared_secret, &first)
   1211                 .unwrap(),
   1212             b"first"
   1213         );
   1214         assert!(receiver.is_official_payload_replay(&first).unwrap());
   1215     }
   1216 
   1217     #[test]
   1218     fn advances_official_pq_ratchet_in_both_directions() {
   1219         let (mut sender, mut receiver) = official_pq_sender_receiver_ratchets();
   1220         let shared_secret = [21_u8; RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH];
   1221 
   1222         let encrypted = sender
   1223             .encrypt_official_payload(&shared_secret, b"pq first", 96)
   1224             .unwrap();
   1225         let plaintext = receiver
   1226             .decrypt_official_payload(&shared_secret, &encrypted)
   1227             .unwrap();
   1228         assert_eq!(plaintext, b"pq first");
   1229         assert!(receiver.pending_outbound_pq_ciphertext.is_some());
   1230         assert!(receiver.local_pq_private_key.is_some());
   1231 
   1232         let reply = receiver
   1233             .encrypt_official_payload(&shared_secret, b"pq reply", 96)
   1234             .unwrap();
   1235         let reply_plaintext = sender
   1236             .decrypt_official_payload(&shared_secret, &reply)
   1237             .unwrap();
   1238         assert_eq!(reply_plaintext, b"pq reply");
   1239         assert!(sender.pending_outbound_pq_ciphertext.is_some());
   1240         assert!(sender.local_pq_private_key.is_some());
   1241     }
   1242 
   1243     #[test]
   1244     fn retains_accepted_pq_ciphertext_across_same_sending_chain_headers() {
   1245         let (mut sender, _) = official_pq_sender_receiver_ratchets();
   1246         let shared_secret = [22_u8; RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH];
   1247         let header_key = sender.official_sending_header_key.clone().unwrap();
   1248         let ratchet_ad = sender.official_associated_data.clone().unwrap();
   1249         let first = sender
   1250             .encrypt_official_payload(&shared_secret, b"pq first", 96)
   1251             .unwrap();
   1252         let second = sender
   1253             .encrypt_official_payload(&shared_secret, b"pq second", 96)
   1254             .unwrap();
   1255 
   1256         let first_message = decode_official_encrypted_message(&first).unwrap();
   1257         let first_header = decrypt_official_header_with_key(
   1258             &decode_official_encrypted_header(&first_message.encrypted_header).unwrap(),
   1259             &header_key,
   1260             &ratchet_ad,
   1261         )
   1262         .unwrap();
   1263         let second_message = decode_official_encrypted_message(&second).unwrap();
   1264         let second_header = decrypt_official_header_with_key(
   1265             &decode_official_encrypted_header(&second_message.encrypted_header).unwrap(),
   1266             &header_key,
   1267             &ratchet_ad,
   1268         )
   1269         .unwrap();
   1270 
   1271         assert!(first_header.pq_ciphertext.is_some());
   1272         assert_eq!(first_header.pq_ciphertext, second_header.pq_ciphertext);
   1273         assert_eq!(second_header.message_number, 1);
   1274     }
   1275 
   1276     #[test]
   1277     fn decrypts_official_skipped_messages_once() {
   1278         let (mut sender, mut receiver) = official_sender_receiver_ratchets();
   1279         let shared_secret = [12_u8; RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH];
   1280         let first = sender
   1281             .encrypt_official_payload(&shared_secret, b"first", 96)
   1282             .unwrap();
   1283         let second = sender
   1284             .encrypt_official_payload(&shared_secret, b"second", 96)
   1285             .unwrap();
   1286         let third = sender
   1287             .encrypt_official_payload(&shared_secret, b"third", 96)
   1288             .unwrap();
   1289 
   1290         assert_eq!(
   1291             receiver
   1292                 .decrypt_official_payload(&shared_secret, &third)
   1293                 .unwrap(),
   1294             b"third"
   1295         );
   1296         assert_eq!(receiver.receiving_chain_length, 3);
   1297         assert_eq!(receiver.official_skipped_message_keys.len(), 2);
   1298         assert_eq!(
   1299             receiver
   1300                 .decrypt_official_payload(&shared_secret, &first)
   1301                 .unwrap(),
   1302             b"first"
   1303         );
   1304         assert_eq!(
   1305             receiver
   1306                 .decrypt_official_payload(&shared_secret, &second)
   1307                 .unwrap(),
   1308             b"second"
   1309         );
   1310         assert!(receiver.official_skipped_message_keys.is_empty());
   1311 
   1312         let replay = receiver
   1313             .decrypt_official_payload(&shared_secret, &first)
   1314             .unwrap_err();
   1315         assert!(matches!(
   1316             replay,
   1317             RadrootsSimplexSmpCryptoError::RatchetMessageRegression {
   1318                 received: 0,
   1319                 current: 3
   1320             }
   1321         ));
   1322     }
   1323 
   1324     #[test]
   1325     fn rejects_too_many_official_skipped_messages() {
   1326         let (mut sender, mut receiver) = official_sender_receiver_ratchets();
   1327         let shared_secret = [13_u8; RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH];
   1328         let mut encrypted = Vec::new();
   1329         for index in 0..=RADROOTS_SIMPLEX_OFFICIAL_MAX_SKIPPED_MESSAGES + 1 {
   1330             encrypted = sender
   1331                 .encrypt_official_payload(&shared_secret, &index.to_be_bytes(), 96)
   1332                 .unwrap();
   1333         }
   1334 
   1335         let error = receiver
   1336             .decrypt_official_payload(&shared_secret, &encrypted)
   1337             .unwrap_err();
   1338         assert_eq!(
   1339             error,
   1340             RadrootsSimplexSmpCryptoError::RatchetTooManySkipped {
   1341                 skipped: RADROOTS_SIMPLEX_OFFICIAL_MAX_SKIPPED_MESSAGES + 1,
   1342                 max: RADROOTS_SIMPLEX_OFFICIAL_MAX_SKIPPED_MESSAGES
   1343             }
   1344         );
   1345         assert_eq!(
   1346             error.to_string(),
   1347             "SMP ratchet skipped 513 messages, exceeding maximum 512"
   1348         );
   1349     }
   1350 
   1351     #[test]
   1352     fn rejects_tampered_official_payload_body() {
   1353         let (mut sender, mut receiver) = official_sender_receiver_ratchets();
   1354         let shared_secret = [12_u8; RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH];
   1355         let mut encrypted = sender
   1356             .encrypt_official_payload(&shared_secret, b"official agent body", 96)
   1357             .unwrap();
   1358         let last = encrypted.len() - 1;
   1359         encrypted[last] ^= 1;
   1360 
   1361         let error = receiver
   1362             .decrypt_official_payload(&shared_secret, &encrypted)
   1363             .unwrap_err();
   1364         assert_eq!(
   1365             error,
   1366             RadrootsSimplexSmpCryptoError::AesGcmAuthenticationFailed
   1367         );
   1368     }
   1369 
   1370     #[test]
   1371     fn decrypts_official_payloads_received_out_of_order() {
   1372         let (mut sender, mut receiver) = official_sender_receiver_ratchets();
   1373         let shared_secret = [13_u8; RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH];
   1374         let encrypted_0 = sender
   1375             .encrypt_official_payload(&shared_secret, b"first official body", 96)
   1376             .unwrap();
   1377         let encrypted_1 = sender
   1378             .encrypt_official_payload(&shared_secret, b"second official body", 96)
   1379             .unwrap();
   1380 
   1381         let plaintext_1 = receiver
   1382             .decrypt_official_payload(&shared_secret, &encrypted_1)
   1383             .unwrap();
   1384         assert_eq!(plaintext_1, b"second official body");
   1385         assert_eq!(receiver.receiving_chain_length, 2);
   1386         assert_eq!(receiver.official_skipped_message_keys.len(), 1);
   1387 
   1388         let plaintext_0 = receiver
   1389             .decrypt_official_payload(&shared_secret, &encrypted_0)
   1390             .unwrap();
   1391         assert_eq!(plaintext_0, b"first official body");
   1392         assert!(receiver.official_skipped_message_keys.is_empty());
   1393         assert_eq!(receiver.receiving_chain_length, 2);
   1394     }
   1395 
   1396     #[test]
   1397     fn rejects_too_many_official_skipped_payloads() {
   1398         let (mut sender, mut receiver) = official_sender_receiver_ratchets();
   1399         let shared_secret = [14_u8; RADROOTS_SIMPLEX_SMP_SHARED_SECRET_LENGTH];
   1400         let mut encrypted = Vec::new();
   1401         for index in 0..514 {
   1402             encrypted = sender
   1403                 .encrypt_official_payload(&shared_secret, &[index as u8], 96)
   1404                 .unwrap();
   1405         }
   1406 
   1407         let error = receiver
   1408             .decrypt_official_payload(&shared_secret, &encrypted)
   1409             .unwrap_err();
   1410         assert_eq!(
   1411             error,
   1412             RadrootsSimplexSmpCryptoError::RatchetTooManySkipped {
   1413                 skipped: 513,
   1414                 max: 512
   1415             }
   1416         );
   1417     }
   1418 }