diff --git a/Cargo.lock b/Cargo.lock index 6aec426d6..ab5bc5516 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -183,6 +183,9 @@ name = "bytes" version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b35204fbdc0b3f4446b89fc1ac2cf84a8a68971995d0bf2e925ec7cd960f9cb3" +dependencies = [ + "serde", +] [[package]] name = "cbc" diff --git a/Cargo.toml b/Cargo.toml index 258bf5f4a..e680a776e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,7 +30,7 @@ anyhow = { version = "1.0", default-features = false } async-channel = { version = "2.5.0", default-features = false } async-trait = "0.1.89" base64 = { version = "0.22.1", default-features = false, features = ["alloc"] } -bytes = { version = "1.5", default-features = false } +bytes = { version = "1.5", default-features = false, features = ["serde"] } chrono = { version = "0.4", default-features = false } ctr = { version = "0.9", default-features = false } flate2 = { version = "1.1.5", default-features = false, features = ["zlib-rs"] } diff --git a/src/message.rs b/src/message.rs index ef7733d59..4b8c475c7 100644 --- a/src/message.rs +++ b/src/message.rs @@ -1235,13 +1235,24 @@ impl Client { ); return; }; + let chain_key_arr: [u8; 32] = match chain_key.as_slice().try_into() { + Ok(arr) => arr, + Err(_) => { + log::error!( + "Invalid chain_key length {} from Go SKDM from {}", + chain_key.len(), + sender_jid + ); + return; + } + }; match SignalPublicKey::from_djb_public_key_bytes(signing_key) { Ok(pub_key) => { match SenderKeyDistributionMessage::new( SENDERKEY_MESSAGE_CURRENT_VERSION, id, iteration, - chain_key.clone(), + chain_key_arr, pub_key, ) { Ok(skdm) => skdm, diff --git a/src/store/signal.rs b/src/store/signal.rs index 87b2bbf03..22e351738 100644 --- a/src/store/signal.rs +++ b/src/store/signal.rs @@ -191,10 +191,7 @@ macro_rules! impl_store_wrapper { #[async_trait] impl IdentityKeyStore for Device { async fn get_identity_key_pair(&self) -> SignalResult { - let private_key_bytes = self.identity_key.private_key; - let private_key = PrivateKey::deserialize(&private_key_bytes.serialize())?; - let ikp = IdentityKeyPair::try_from(private_key)?; - Ok(ikp) + Ok(self.identity_key.into()) } async fn get_local_registration_id(&self) -> SignalResult { @@ -333,19 +330,9 @@ impl SignedPreKeyStore for Device { signed_prekey_id: u32, ) -> Result, StoreError> { if signed_prekey_id == self.signed_pre_key_id { - use wacore::libsignal::protocol::{KeyPair, PrivateKey, PublicKey}; - - let public_key = PublicKey::from_djb_public_key_bytes( - self.signed_pre_key.public_key.public_key_bytes(), - ) - .map_err(|e| Box::new(e) as StoreError)?; - let private_key = PrivateKey::deserialize(&self.signed_pre_key.private_key.serialize()) - .map_err(|e| Box::new(e) as StoreError)?; - let key_pair = KeyPair::new(public_key, private_key); - let record = wacore::libsignal::store::record_helpers::new_signed_pre_key_record( self.signed_pre_key_id, - &key_pair, + &self.signed_pre_key, self.signed_pre_key_signature, chrono::Utc::now(), ); diff --git a/storages/sqlite-storage/src/sqlite_store.rs b/storages/sqlite-storage/src/sqlite_store.rs index 6b5d66adf..50dfd18b8 100644 --- a/storages/sqlite-storage/src/sqlite_store.rs +++ b/storages/sqlite-storage/src/sqlite_store.rs @@ -148,7 +148,7 @@ impl SqliteStore { fn serialize_keypair(&self, key_pair: &KeyPair) -> Result> { let mut bytes = Vec::with_capacity(64); - bytes.extend_from_slice(&key_pair.private_key.serialize()); + bytes.extend_from_slice(key_pair.private_key.serialize()); bytes.extend_from_slice(key_pair.public_key.public_key_bytes()); Ok(bytes) } @@ -272,19 +272,19 @@ impl SqliteStore { let noise_key_data = { let mut bytes = Vec::with_capacity(64); - bytes.extend_from_slice(&new_device.noise_key.private_key.serialize()); + bytes.extend_from_slice(new_device.noise_key.private_key.serialize()); bytes.extend_from_slice(new_device.noise_key.public_key.public_key_bytes()); bytes }; let identity_key_data = { let mut bytes = Vec::with_capacity(64); - bytes.extend_from_slice(&new_device.identity_key.private_key.serialize()); + bytes.extend_from_slice(new_device.identity_key.private_key.serialize()); bytes.extend_from_slice(new_device.identity_key.public_key.public_key_bytes()); bytes }; let signed_pre_key_data = { let mut bytes = Vec::with_capacity(64); - bytes.extend_from_slice(&new_device.signed_pre_key.private_key.serialize()); + bytes.extend_from_slice(new_device.signed_pre_key.private_key.serialize()); bytes.extend_from_slice(new_device.signed_pre_key.public_key.public_key_bytes()); bytes }; diff --git a/wacore/libsignal/src/core/curve.rs b/wacore/libsignal/src/core/curve.rs index fbe5cdefc..f0b4553ed 100644 --- a/wacore/libsignal/src/core/curve.rs +++ b/wacore/libsignal/src/core/curve.rs @@ -105,16 +105,14 @@ impl PublicKey { } } - pub fn serialize(&self) -> Box<[u8]> { - let value_len = match &self.key { - PublicKeyData::DjbPublicKey(v) => v.len(), - }; - let mut result = Vec::with_capacity(1 + value_len); - result.push(self.key_type().value()); + /// Serialize the public key to a fixed-size array (1 type byte + 32 key bytes). + pub fn serialize(&self) -> [u8; 33] { + let mut result = [0u8; 33]; + result[0] = self.key_type().value(); match &self.key { - PublicKeyData::DjbPublicKey(v) => result.extend_from_slice(v), + PublicKeyData::DjbPublicKey(v) => result[1..].copy_from_slice(v), } - result.into_boxed_slice() + result } pub fn verify_signature(&self, message: &[u8], signature: &[u8]) -> bool { @@ -228,9 +226,9 @@ impl PrivateKey { } } - pub fn serialize(&self) -> Vec { + pub fn serialize(&self) -> &[u8; 32] { match &self.key { - PrivateKeyData::DjbPrivateKey(v) => v.to_vec(), + PrivateKeyData::DjbPrivateKey(v) => v, } } @@ -254,7 +252,7 @@ impl PrivateKey { &self, message: &[u8], csprng: &mut R, - ) -> Result, CurveError> { + ) -> Result<[u8; 64], CurveError> { self.calculate_signature_for_multipart_message(&[message], csprng) } @@ -262,20 +260,20 @@ impl PrivateKey { &self, message: &[&[u8]], csprng: &mut R, - ) -> Result, CurveError> { + ) -> Result<[u8; 64], CurveError> { match self.key { PrivateKeyData::DjbPrivateKey(k) => { let private_key = curve25519::PrivateKey::from(k); - Ok(Box::new(private_key.calculate_signature(csprng, message))) + Ok(private_key.calculate_signature(csprng, message)) } } } - pub fn calculate_agreement(&self, their_key: &PublicKey) -> Result, CurveError> { + pub fn calculate_agreement(&self, their_key: &PublicKey) -> Result<[u8; 32], CurveError> { match (self.key, their_key.key) { (PrivateKeyData::DjbPrivateKey(priv_key), PublicKeyData::DjbPublicKey(pub_key)) => { let private_key = curve25519::PrivateKey::from(priv_key); - Ok(Box::new(private_key.calculate_agreement(&pub_key))) + Ok(private_key.calculate_agreement(&pub_key)) } } } @@ -335,11 +333,11 @@ impl KeyPair { &self, message: &[u8], csprng: &mut R, - ) -> Result, CurveError> { + ) -> Result<[u8; 64], CurveError> { self.private_key.calculate_signature(message, csprng) } - pub fn calculate_agreement(&self, their_key: &PublicKey) -> Result, CurveError> { + pub fn calculate_agreement(&self, their_key: &PublicKey) -> Result<[u8; 32], CurveError> { self.private_key.calculate_agreement(their_key) } } diff --git a/wacore/libsignal/src/protocol/group_cipher.rs b/wacore/libsignal/src/protocol/group_cipher.rs index 78cb7aa94..f75be9cf9 100644 --- a/wacore/libsignal/src/protocol/group_cipher.rs +++ b/wacore/libsignal/src/protocol/group_cipher.rs @@ -17,11 +17,13 @@ use crate::protocol::{ }; use crate::store::sender_key_name::SenderKeyName; -struct EncryptionBuffer { +/// Reusable buffer for cryptographic operations (encryption and decryption). +/// Named generically since it's used for both ENCRYPTION_BUFFER and DECRYPTION_BUFFER. +struct CryptoBuffer { buffer: Vec, } -impl EncryptionBuffer { +impl CryptoBuffer { const INITIAL_CAPACITY: usize = 1024; fn new() -> Self { @@ -29,14 +31,23 @@ impl EncryptionBuffer { buffer: Vec::with_capacity(Self::INITIAL_CAPACITY), } } + + /// Clears the buffer and returns a mutable reference for writing. fn get_buffer(&mut self) -> &mut Vec { self.buffer.clear(); &mut self.buffer } + + /// Takes ownership of the buffer contents, replacing with a fresh pre-allocated buffer. + /// More efficient than `mem::take` + `reserve` since we swap with an already-allocated buffer. + fn take_buffer(&mut self) -> Vec { + std::mem::replace(&mut self.buffer, Vec::with_capacity(Self::INITIAL_CAPACITY)) + } } thread_local! { - static ENCRYPTION_BUFFER: RefCell = RefCell::new(EncryptionBuffer::new()); + static ENCRYPTION_BUFFER: RefCell = RefCell::new(CryptoBuffer::new()); + static DECRYPTION_BUFFER: RefCell = RefCell::new(CryptoBuffer::new()); } pub async fn group_encrypt( @@ -69,7 +80,7 @@ pub async fn group_encrypt( .sender_chain_key() .ok_or(SignalProtocolError::InvalidSenderKeySession)?; - let message_keys = sender_chain_key.sender_message_key(); + let (message_keys, next_sender_chain_key) = sender_chain_key.step_with_message_key()?; let ciphertext = ENCRYPTION_BUFFER.with(|buffer| { let mut buf_wrapper = buffer.borrow_mut(); @@ -79,10 +90,7 @@ pub async fn group_encrypt( log::error!("outgoing sender key state corrupt for distribution"); SignalProtocolError::InvalidSenderKeySession })?; - let result = std::mem::take(buf); - // Restore buffer capacity for next use (take() leaves empty Vec with 0 capacity) - buf.reserve(EncryptionBuffer::INITIAL_CAPACITY); - Ok::, SignalProtocolError>(result) + Ok::, SignalProtocolError>(buf_wrapper.take_buffer()) })?; let signing_key = sender_key_state @@ -98,7 +106,7 @@ pub async fn group_encrypt( &signing_key, )?; - sender_key_state.set_sender_chain_key(sender_chain_key.next()?); + sender_key_state.set_sender_chain_key(next_sender_chain_key); sender_key_store .store_sender_key(sender_key_name, &record) @@ -141,12 +149,14 @@ fn get_sender_key(state: &mut SenderKeyState, iteration: u32) -> Result { - log::error!( - "incoming sender key state corrupt for group {} sender {} (chain ID {chain_id})", - sender_key_name.group_id(), - sender_key_name.sender_id() - ); - return Err(SignalProtocolError::InvalidSenderKeySession); - } - DecryptionErrorCrypto::BadCiphertext(msg) => { - log::error!("sender key decryption failed: {msg}"); - return Err(SignalProtocolError::InvalidMessage( - CiphertextMessageType::SenderKey, - "decryption failed", - )); + let plaintext = DECRYPTION_BUFFER.with(|buffer| { + let mut buf_wrapper = buffer.borrow_mut(); + let buf = buf_wrapper.get_buffer(); + if let Err(e) = aes_256_cbc_decrypt_into( + skm.ciphertext(), + sender_key.cipher_key(), + sender_key.iv(), + buf, + ) { + match e { + DecryptionErrorCrypto::BadKeyOrIv => { + log::error!( + "incoming sender key state corrupt for group {} sender {} (chain ID {chain_id})", + sender_key_name.group_id(), + sender_key_name.sender_id() + ); + return Err(SignalProtocolError::InvalidSenderKeySession); + } + DecryptionErrorCrypto::BadCiphertext(msg) => { + log::error!("sender key decryption failed: {msg}"); + return Err(SignalProtocolError::InvalidMessage( + CiphertextMessageType::SenderKey, + "decryption failed", + )); + } } } - } + Ok::, SignalProtocolError>(buf_wrapper.take_buffer()) + })?; sender_key_store .store_sender_key(sender_key_name, &record) @@ -243,7 +257,7 @@ pub async fn process_sender_key_distribution_message( "Processing SenderKey distribution for group {} from sender {} with chain ID {}", sender_key_name.group_id(), sender_key_name.sender_id(), - skdm.chain_id()? + skdm.chain_id() ); let mut sender_key_record = sender_key_store @@ -253,10 +267,10 @@ pub async fn process_sender_key_distribution_message( sender_key_record.add_sender_key_state( skdm.message_version(), - skdm.chain_id()?, - skdm.iteration()?, - skdm.chain_key()?, - *skdm.signing_key()?, + skdm.chain_id(), + skdm.iteration(), + skdm.chain_key(), + *skdm.signing_key(), None, ); sender_key_store @@ -313,7 +327,7 @@ pub async fn create_sender_key_distribution_message( message_version, state.chain_id(), sender_chain_key.iteration(), - sender_chain_key.seed().to_vec(), + *sender_chain_key.seed(), state .signing_key_public() .map_err(|_| SignalProtocolError::InvalidSenderKeySession)?, diff --git a/wacore/libsignal/src/protocol/identity_key.rs b/wacore/libsignal/src/protocol/identity_key.rs index effe745f6..ffaf616d5 100644 --- a/wacore/libsignal/src/protocol/identity_key.rs +++ b/wacore/libsignal/src/protocol/identity_key.rs @@ -40,9 +40,9 @@ impl IdentityKey { &self.public_key } - /// Return an owned byte slice which can be deserialized with [`Self::decode`]. + /// Serialize the identity key to a fixed-size array (1 type byte + 32 key bytes). #[inline] - pub fn serialize(&self) -> Box<[u8]> { + pub fn serialize(&self) -> [u8; 33] { self.public_key.serialize() } @@ -138,7 +138,7 @@ impl IdentityKeyPair { &self, other: &IdentityKey, rng: &mut R, - ) -> Result> { + ) -> Result<[u8; 64]> { Ok(self.private_key.calculate_signature_for_multipart_message( &[ ALTERNATE_IDENTITY_SIGNATURE_PREFIX_1, diff --git a/wacore/libsignal/src/protocol/protocol.rs b/wacore/libsignal/src/protocol/protocol.rs index dc738a3d0..ea05b2f61 100644 --- a/wacore/libsignal/src/protocol/protocol.rs +++ b/wacore/libsignal/src/protocol/protocol.rs @@ -84,7 +84,7 @@ impl SignalMessage { receiver_identity_key: &IdentityKey, ) -> Result { let message = waproto::whatsapp::SignalMessage { - ratchet_key: Some(sender_ratchet_key.serialize().into_vec()), + ratchet_key: Some(sender_ratchet_key.serialize().to_vec()), counter: Some(counter), previous_counter: Some(previous_counter), ciphertext: Some(Vec::::from(ciphertext)), @@ -259,8 +259,8 @@ impl PreKeySignalMessage { registration_id: Some(registration_id), pre_key_id: pre_key_id.map(|id| id.into()), signed_pre_key_id: Some(signed_pre_key_id.into()), - base_key: Some(base_key.serialize().into_vec()), - identity_key: Some(identity_key.serialize().into_vec()), + base_key: Some(base_key.serialize().to_vec()), + identity_key: Some(identity_key.serialize().to_vec()), message: Some(Vec::from(message.as_ref())), }; let mut serialized = Vec::with_capacity(1 + proto_message.encoded_len()); @@ -397,26 +397,23 @@ impl SenderKeyMessage { let proto_message = waproto::whatsapp::SenderKeyMessage { id: Some(chain_id), iteration: Some(iteration), - ciphertext: Some(ciphertext.to_vec()), + ciphertext: Some(Vec::from(ciphertext.as_ref())), }; - let proto_bytes = proto_message.encode_to_vec(); - - // The signature must cover the version byte concatenated with the protobuf - // payload. Other clients (e.g. baileys, libsignal-go) compute the signature - // over [shifted_version || proto_bytes]. Signing only the proto bytes causes - // verification to fail on recipients. + // Build serialized buffer directly: [version_byte || proto || signature] + // Sign over [version_byte || proto], then append signature let shifted_version = (message_version << 4) | 3u8; - let mut data_to_sign = Vec::with_capacity(1 + proto_bytes.len()); - data_to_sign.push(shifted_version); - data_to_sign.extend_from_slice(&proto_bytes); + let proto_len = proto_message.encoded_len(); + let mut serialized = Vec::with_capacity(1 + proto_len + Self::SIGNATURE_LEN); + serialized.push(shifted_version); + proto_message + .encode(&mut serialized) + .expect("can always append to a buffer"); + // Sign the data we've built so far (version + proto) let signature = signature_key - .calculate_signature(&data_to_sign, csprng) + .calculate_signature(&serialized, csprng) .map_err(|_| SignalProtocolError::SignatureValidationFailed)?; - - let mut serialized = vec![shifted_version]; - serialized.extend_from_slice(&proto_bytes); serialized.extend_from_slice(&signature); Ok(Self { @@ -424,7 +421,7 @@ impl SenderKeyMessage { chain_id, iteration, ciphertext, - serialized: Box::from(serialized), + serialized: serialized.into_boxed_slice(), }) } @@ -518,7 +515,7 @@ pub struct SenderKeyDistributionMessage { message_version: u8, chain_id: u32, iteration: u32, - chain_key: Vec, + chain_key: [u8; 32], signing_key: PublicKey, serialized: Box<[u8]>, } @@ -528,13 +525,13 @@ impl SenderKeyDistributionMessage { message_version: u8, chain_id: u32, iteration: u32, - chain_key: Vec, + chain_key: [u8; 32], signing_key: PublicKey, ) -> Result { let proto_message = waproto::whatsapp::SenderKeyDistributionMessage { id: Some(chain_id), iteration: Some(iteration), - chain_key: Some(chain_key.clone()), + chain_key: Some(chain_key.to_vec()), signing_key: Some(signing_key.serialize().to_vec()), }; let mut serialized = Vec::with_capacity(1 + proto_message.encoded_len()); @@ -559,23 +556,23 @@ impl SenderKeyDistributionMessage { } #[inline] - pub fn chain_id(&self) -> Result { - Ok(self.chain_id) + pub fn chain_id(&self) -> u32 { + self.chain_id } #[inline] - pub fn iteration(&self) -> Result { - Ok(self.iteration) + pub fn iteration(&self) -> u32 { + self.iteration } #[inline] - pub fn chain_key(&self) -> Result<&[u8]> { - Ok(&self.chain_key) + pub fn chain_key(&self) -> &[u8; 32] { + &self.chain_key } #[inline] - pub fn signing_key(&self) -> Result<&PublicKey> { - Ok(&self.signing_key) + pub fn signing_key(&self) -> &PublicKey { + &self.signing_key } #[inline] @@ -621,17 +618,20 @@ impl TryFrom<&[u8]> for SenderKeyDistributionMessage { let iteration = proto_structure .iteration .ok_or(SignalProtocolError::InvalidProtobufEncoding)?; - let chain_key = proto_structure + let chain_key_vec = proto_structure .chain_key .ok_or(SignalProtocolError::InvalidProtobufEncoding)?; let signing_key = proto_structure .signing_key .ok_or(SignalProtocolError::InvalidProtobufEncoding)?; - if chain_key.len() != 32 || signing_key.len() != 33 { + if chain_key_vec.len() != 32 || signing_key.len() != 33 { return Err(SignalProtocolError::InvalidProtobufEncoding); } + let chain_key: [u8; 32] = chain_key_vec + .try_into() + .map_err(|_| SignalProtocolError::InvalidProtobufEncoding)?; let signing_key = PublicKey::deserialize(&signing_key)?; Ok(SenderKeyDistributionMessage { diff --git a/wacore/libsignal/src/protocol/ratchet.rs b/wacore/libsignal/src/protocol/ratchet.rs index d43ee74cd..bce8bbdf1 100644 --- a/wacore/libsignal/src/protocol/ratchet.rs +++ b/wacore/libsignal/src/protocol/ratchet.rs @@ -46,33 +46,41 @@ pub fn initialize_alice_session( let sending_ratchet_key = KeyPair::generate(&mut csprng); - let mut secrets = Vec::with_capacity(32 * 5); + // Stack-allocated buffer for up to 5 shared secrets (160 bytes max) + let mut secrets = [0u8; 160]; + let mut secrets_len = 0usize; - secrets.extend_from_slice(&[0xFFu8; 32]); // "discontinuity bytes" + // "discontinuity bytes" + secrets[..32].copy_from_slice(&[0xFFu8; 32]); + secrets_len += 32; let our_base_private_key = parameters.our_base_key_pair().private_key; - secrets.extend_from_slice( - ¶meters - .our_identity_key_pair() - .private_key() - .calculate_agreement(parameters.their_signed_pre_key())?, - ); + // Each agreement is 32 bytes. We have: discontinuity (32) + up to 4 agreements (128) = 160 max. + // The buffer is [u8; 160], so bounds are statically guaranteed. + let agreement = parameters + .our_identity_key_pair() + .private_key() + .calculate_agreement(parameters.their_signed_pre_key())?; + secrets[secrets_len..secrets_len + 32].copy_from_slice(&agreement); + secrets_len += 32; - secrets.extend_from_slice( - &our_base_private_key.calculate_agreement(parameters.their_identity_key().public_key())?, - ); + let agreement = + our_base_private_key.calculate_agreement(parameters.their_identity_key().public_key())?; + secrets[secrets_len..secrets_len + 32].copy_from_slice(&agreement); + secrets_len += 32; - secrets.extend_from_slice( - &our_base_private_key.calculate_agreement(parameters.their_signed_pre_key())?, - ); + let agreement = our_base_private_key.calculate_agreement(parameters.their_signed_pre_key())?; + secrets[secrets_len..secrets_len + 32].copy_from_slice(&agreement); + secrets_len += 32; if let Some(their_one_time_prekey) = parameters.their_one_time_pre_key() { - secrets - .extend_from_slice(&our_base_private_key.calculate_agreement(their_one_time_prekey)?); + let agreement = our_base_private_key.calculate_agreement(their_one_time_prekey)?; + secrets[secrets_len..secrets_len + 32].copy_from_slice(&agreement); + secrets_len += 32; } - let (root_key, chain_key, _) = derive_keys(&secrets); + let (root_key, chain_key, _) = derive_keys(&secrets[..secrets_len]); let (sending_chain_root_key, sending_chain_chain_key) = root_key.create_chain( parameters.their_ratchet_key(), @@ -95,40 +103,46 @@ pub fn initialize_alice_session( pub fn initialize_bob_session(parameters: &BobSignalProtocolParameters) -> Result { let local_identity = parameters.our_identity_key_pair().identity_key(); - let mut secrets = Vec::with_capacity(32 * 5); - - secrets.extend_from_slice(&[0xFFu8; 32]); // "discontinuity bytes" - - secrets.extend_from_slice( - ¶meters - .our_signed_pre_key_pair() - .private_key - .calculate_agreement(parameters.their_identity_key().public_key())?, - ); - - secrets.extend_from_slice( - ¶meters - .our_identity_key_pair() - .private_key() - .calculate_agreement(parameters.their_base_key())?, - ); - - secrets.extend_from_slice( - ¶meters - .our_signed_pre_key_pair() - .private_key - .calculate_agreement(parameters.their_base_key())?, - ); + // Stack-allocated buffer for up to 5 shared secrets (160 bytes max) + let mut secrets = [0u8; 160]; + let mut secrets_len = 0usize; + + // "discontinuity bytes" + secrets[..32].copy_from_slice(&[0xFFu8; 32]); + secrets_len += 32; + + // Each agreement is 32 bytes. We have: discontinuity (32) + up to 4 agreements (128) = 160 max. + // The buffer is [u8; 160], so bounds are statically guaranteed. + let agreement = parameters + .our_signed_pre_key_pair() + .private_key + .calculate_agreement(parameters.their_identity_key().public_key())?; + secrets[secrets_len..secrets_len + 32].copy_from_slice(&agreement); + secrets_len += 32; + + let agreement = parameters + .our_identity_key_pair() + .private_key() + .calculate_agreement(parameters.their_base_key())?; + secrets[secrets_len..secrets_len + 32].copy_from_slice(&agreement); + secrets_len += 32; + + let agreement = parameters + .our_signed_pre_key_pair() + .private_key + .calculate_agreement(parameters.their_base_key())?; + secrets[secrets_len..secrets_len + 32].copy_from_slice(&agreement); + secrets_len += 32; if let Some(our_one_time_pre_key_pair) = parameters.our_one_time_pre_key_pair() { - secrets.extend_from_slice( - &our_one_time_pre_key_pair - .private_key - .calculate_agreement(parameters.their_base_key())?, - ); + let agreement = our_one_time_pre_key_pair + .private_key + .calculate_agreement(parameters.their_base_key())?; + secrets[secrets_len..secrets_len + 32].copy_from_slice(&agreement); + secrets_len += 32; } - let (root_key, chain_key, _) = derive_keys(&secrets); + let (root_key, chain_key, _) = derive_keys(&secrets[..secrets_len]); let session = SessionState::new( message_version(), diff --git a/wacore/libsignal/src/protocol/ratchet/keys.rs b/wacore/libsignal/src/protocol/ratchet/keys.rs index cc960c379..d7e23c54c 100644 --- a/wacore/libsignal/src/protocol/ratchet/keys.rs +++ b/wacore/libsignal/src/protocol/ratchet/keys.rs @@ -7,56 +7,113 @@ use std::fmt; use arrayref::array_ref; +use hmac::{Hmac, Mac}; +use sha2::Sha256; + use crate::protocol::{PrivateKey, PublicKey, Result, crypto, stores::session_structure}; +/// Lazy message key generator that defers key derivation and avoids re-serialization. +/// +/// This enum enables two optimizations: +/// 1. **Lazy derivation**: Keys are only derived from seed when actually needed for encryption +/// 2. **Zero-cost round-trip**: Keys loaded from protobuf are kept in serialized form and +/// returned as-is when saving, avoiding unnecessary deserialization and re-serialization pub enum MessageKeyGenerator { + /// Native computed keys - from encryption operations Keys(MessageKeys), - Seed((Vec, u32)), + /// Seed for lazy derivation - keys derived on demand + Seed(([u8; 32], u32)), + /// Original protobuf - zero-cost pass-through on save + Serialized(session_structure::chain::MessageKey), } impl MessageKeyGenerator { - pub fn new_from_seed(seed: &[u8], counter: u32) -> Self { - Self::Seed((seed.to_vec(), counter)) + #[inline] + pub fn new_from_seed(seed: &[u8; 32], counter: u32) -> Self { + Self::Seed((*seed, counter)) } + + /// Generate the actual MessageKeys, deriving them if necessary. + /// This is called when keys are needed for actual encryption/decryption. pub fn generate_keys(self) -> MessageKeys { match self { Self::Seed((seed, counter)) => MessageKeys::derive_keys(&seed, None, counter), Self::Keys(k) => k, + Self::Serialized(pb) => { + // Parse on demand - only when keys are actually needed. + // Note: from_pb() validates field lengths before creating Serialized, + // so these conversions should always succeed. The unwrap_or fallbacks + // exist only as a defensive measure; in debug builds we assert to + // catch any invariant violations. + let cipher_key = pb + .cipher_key + .as_deref() + .and_then(|b| <[u8; 32]>::try_from(b).ok()); + let mac_key = pb + .mac_key + .as_deref() + .and_then(|b| <[u8; 32]>::try_from(b).ok()); + let iv = pb.iv.as_deref().and_then(|b| <[u8; 16]>::try_from(b).ok()); + + debug_assert!( + cipher_key.is_some() && mac_key.is_some() && iv.is_some(), + "Serialized MessageKeyGenerator has invalid field lengths - from_pb should have rejected this" + ); + + MessageKeys { + cipher_key: cipher_key.unwrap_or([0u8; 32]), + mac_key: mac_key.unwrap_or([0u8; 32]), + iv: iv.unwrap_or([0u8; 16]), + counter: pb.index.unwrap_or(0), + } + } } } + + /// Convert to protobuf format for storage. + /// Zero-cost for Serialized variant (pass-through), allocates for others. pub fn into_pb(self) -> session_structure::chain::MessageKey { - let keys = self.generate_keys(); - session_structure::chain::MessageKey { - cipher_key: Some(keys.cipher_key().to_vec()), - mac_key: Some(keys.mac_key().to_vec()), - iv: Some(keys.iv().to_vec()), - index: Some(keys.counter()), + match self { + // Zero-cost pass-through: return original protobuf unchanged + Self::Serialized(pb) => pb, + // Need to serialize: derive keys and convert + Self::Seed(_) | Self::Keys(_) => { + use prost::bytes::Bytes; + let keys = self.generate_keys(); + session_structure::chain::MessageKey { + cipher_key: Some(Bytes::copy_from_slice(keys.cipher_key())), + mac_key: Some(Bytes::copy_from_slice(keys.mac_key())), + iv: Some(Bytes::copy_from_slice(keys.iv())), + index: Some(keys.counter()), + } + } } } + + /// Load from protobuf without parsing - keeps original bytes for zero-cost round-trip. pub fn from_pb( pb: session_structure::chain::MessageKey, ) -> std::result::Result { - Ok(Self::Keys(MessageKeys { - cipher_key: pb - .cipher_key - .as_deref() - .ok_or("missing cipher key")? - .try_into() - .map_err(|_| "invalid message cipher key")?, - mac_key: pb - .mac_key - .as_deref() - .ok_or("missing mac key")? - .try_into() - .map_err(|_| "invalid message MAC key")?, - iv: pb - .iv - .as_deref() - .ok_or("missing iv")? - .try_into() - .map_err(|_| "invalid message IV")?, - counter: pb.index.unwrap_or(0), - })) + // Validate the protobuf has required fields + if pb.cipher_key.as_ref().is_some_and(|b| b.len() == 32) + && pb.mac_key.as_ref().is_some_and(|b| b.len() == 32) + && pb.iv.as_ref().is_some_and(|b| b.len() == 16) + { + // Keep as Serialized for zero-cost round-trip + Ok(Self::Serialized(pb)) + } else { + Err("invalid message key format") + } + } + + /// Get the counter/index without fully parsing the keys. + #[inline] + pub fn counter(&self) -> u32 { + match self { + Self::Keys(k) => k.counter(), + Self::Seed((_, counter)) => *counter, + Self::Serialized(pb) => pb.index.unwrap_or(0), + } } } @@ -108,7 +165,7 @@ impl MessageKeys { } } -#[derive(Clone, Debug)] +#[derive(Clone, Copy, Debug)] pub struct ChainKey { key: [u8; 32], index: u32, @@ -146,6 +203,27 @@ impl ChainKey { ) } + /// Compute both message keys and next chain key in one call, reusing HMAC key setup. + #[inline] + pub fn step_with_message_keys(&self) -> (MessageKeyGenerator, Self) { + let mut hmac = Hmac::::new_from_slice(&self.key) + .expect("HMAC-SHA256 should accept any size key"); + + hmac.update(&Self::MESSAGE_KEY_SEED); + let message_key_seed: [u8; 32] = hmac.finalize_reset().into_bytes().into(); + + hmac.update(&Self::CHAIN_KEY_SEED); + let next_key: [u8; 32] = hmac.finalize().into_bytes().into(); + + let message_keys = MessageKeyGenerator::new_from_seed(&message_key_seed, self.index); + let next_chain = Self { + key: next_key, + index: self.index + 1, + }; + + (message_keys, next_chain) + } + fn calculate_base_material(&self, seed: [u8; 1]) -> [u8; 32] { crypto::hmac_sha256(&self.key, &seed) } @@ -303,7 +381,7 @@ mod tests { /// Test MessageKeyGenerator from seed #[test] fn test_message_key_generator_from_seed() { - let seed = vec![0xBBu8; 32]; + let seed = [0xBBu8; 32]; let counter = 42; let generator = MessageKeyGenerator::new_from_seed(&seed, counter); @@ -339,4 +417,72 @@ mod tests { assert_eq!(keys1.counter(), counter); assert_eq!(keys2.counter(), counter); } + + /// Test that step_with_message_keys produces the same results as + /// calling message_keys() and next_chain_key() separately + #[test] + fn test_step_with_message_keys_equivalence() { + let initial_key = [0x77u8; 32]; + let chain_key = ChainKey::new(initial_key, 5); + + // Get results using separate calls + let message_keys_separate = chain_key.message_keys().generate_keys(); + let next_chain_separate = chain_key.next_chain_key(); + + // Get results using optimized combined call + let (message_keys_gen_combined, next_chain_combined) = chain_key.step_with_message_keys(); + let message_keys_combined = message_keys_gen_combined.generate_keys(); + + // Verify message keys are identical + assert_eq!( + message_keys_separate.cipher_key(), + message_keys_combined.cipher_key() + ); + assert_eq!( + message_keys_separate.mac_key(), + message_keys_combined.mac_key() + ); + assert_eq!(message_keys_separate.iv(), message_keys_combined.iv()); + assert_eq!( + message_keys_separate.counter(), + message_keys_combined.counter() + ); + + // Verify next chain key is identical + assert_eq!(next_chain_separate.key(), next_chain_combined.key()); + assert_eq!(next_chain_separate.index(), next_chain_combined.index()); + } + + /// Test step_with_message_keys over multiple iterations + #[test] + fn test_step_with_message_keys_chain() { + let initial_key = [0x88u8; 32]; + let mut chain_separate = ChainKey::new(initial_key, 0); + let mut chain_combined = ChainKey::new(initial_key, 0); + + // Step both chains 10 times and verify they stay in sync + for i in 0..10 { + let msg_keys_sep = chain_separate.message_keys().generate_keys(); + chain_separate = chain_separate.next_chain_key(); + + let (msg_keys_gen_comb, next_chain) = chain_combined.step_with_message_keys(); + let msg_keys_comb = msg_keys_gen_comb.generate_keys(); + chain_combined = next_chain; + + // Verify message keys match + assert_eq!( + msg_keys_sep.cipher_key(), + msg_keys_comb.cipher_key(), + "cipher_key mismatch at iteration {i}" + ); + + // Verify chain keys match + assert_eq!( + chain_separate.key(), + chain_combined.key(), + "chain key mismatch at iteration {i}" + ); + assert_eq!(chain_separate.index(), chain_combined.index()); + } + } } diff --git a/wacore/libsignal/src/protocol/sender_keys.rs b/wacore/libsignal/src/protocol/sender_keys.rs index ea593d2bf..4e03c7914 100644 --- a/wacore/libsignal/src/protocol/sender_keys.rs +++ b/wacore/libsignal/src/protocol/sender_keys.rs @@ -8,6 +8,9 @@ use std::collections::VecDeque; use itertools::Itertools; use prost::Message; +use hmac::{Hmac, Mac}; +use sha2::Sha256; + use crate::protocol::crypto::hmac_sha256; use crate::protocol::stores::{ SenderKeyRecordStructure, SenderKeyStateStructure, sender_key_state_structure, @@ -49,8 +52,9 @@ impl SenderMessageKey { } pub(crate) fn from_protobuf(smk: sender_key_state_structure::SenderMessageKey) -> Self { - let seed_vec = smk.seed.unwrap_or_default(); - let seed: [u8; 32] = seed_vec + let seed_bytes = smk.seed.unwrap_or_default(); + let seed: [u8; 32] = seed_bytes + .as_ref() .try_into() .expect("SenderMessageKey seed must be exactly 32 bytes"); Self::new(smk.iteration.unwrap_or_default(), seed) @@ -69,14 +73,15 @@ impl SenderMessageKey { } pub(crate) fn as_protobuf(&self) -> sender_key_state_structure::SenderMessageKey { + use prost::bytes::Bytes; sender_key_state_structure::SenderMessageKey { iteration: Some(self.iteration), - seed: Some(self.seed.to_vec()), + seed: Some(Bytes::copy_from_slice(&self.seed)), } } } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Copy)] pub struct SenderChainKey { iteration: u32, chain_key: [u8; 32], @@ -97,7 +102,7 @@ impl SenderChainKey { self.iteration } - pub fn seed(&self) -> &[u8] { + pub fn seed(&self) -> &[u8; 32] { &self.chain_key } @@ -119,6 +124,34 @@ impl SenderChainKey { SenderMessageKey::new(self.iteration, self.get_derivative(Self::MESSAGE_KEY_SEED)) } + /// Compute both sender message key and next chain key in one call, reusing HMAC key setup. + #[inline] + pub fn step_with_message_key(&self) -> Result<(SenderMessageKey, Self), SignalProtocolError> { + let new_iteration = self.iteration.checked_add(1).ok_or_else(|| { + SignalProtocolError::InvalidState( + "sender_chain_key_step", + "Sender chain is too long".into(), + ) + })?; + + let mut hmac = Hmac::::new_from_slice(&self.chain_key) + .expect("HMAC-SHA256 should accept any size key"); + + hmac.update(&[Self::MESSAGE_KEY_SEED]); + let message_key_seed: [u8; 32] = hmac.finalize_reset().into_bytes().into(); + + hmac.update(&[Self::CHAIN_KEY_SEED]); + let next_chain_key: [u8; 32] = hmac.finalize().into_bytes().into(); + + let message_key = SenderMessageKey::new(self.iteration, message_key_seed); + let next_chain = Self { + iteration: new_iteration, + chain_key: next_chain_key, + }; + + Ok((message_key, next_chain)) + } + #[inline] fn get_derivative(&self, label: u8) -> [u8; 32] { let label = [label]; @@ -126,9 +159,10 @@ impl SenderChainKey { } pub(crate) fn as_protobuf(&self) -> sender_key_state_structure::SenderChainKey { + use prost::bytes::Bytes; sender_key_state_structure::SenderChainKey { iteration: Some(self.iteration), - seed: Some(self.chain_key.to_vec()), + seed: Some(Bytes::copy_from_slice(&self.chain_key)), } } } @@ -147,13 +181,15 @@ impl SenderKeyState { signature_key: PublicKey, signature_private_key: Option, ) -> SenderKeyState { + use prost::bytes::Bytes; let chain_key_arr: [u8; 32] = chain_key.try_into().expect("chain_key must be 32 bytes"); let state = SenderKeyStateStructure { sender_key_id: Some(chain_id), sender_chain_key: Some(SenderChainKey::new(iteration, chain_key_arr).as_protobuf()), sender_signing_key: Some(sender_key_state_structure::SenderSigningKey { - public: Some(signature_key.serialize().to_vec()), - private: signature_private_key.map(|k| k.serialize().to_vec()), + public: Some(Bytes::copy_from_slice(&signature_key.serialize())), + private: signature_private_key + .map(|k| Bytes::copy_from_slice(k.serialize().as_ref())), }), sender_message_keys: vec![], }; @@ -225,8 +261,12 @@ impl SenderKeyState { self.state .sender_message_keys .push(sender_message_key.as_protobuf()); - while self.state.sender_message_keys.len() > consts::MAX_MESSAGE_KEYS { - self.state.sender_message_keys.remove(0); + // Remove oldest keys if we exceed capacity. + // Using drain() is O(n) once vs remove(0) in a loop which is O(n) per removal. + let len = self.state.sender_message_keys.len(); + if len > consts::MAX_MESSAGE_KEYS { + let excess = len - consts::MAX_MESSAGE_KEYS; + self.state.sender_message_keys.drain(..excess); } } @@ -529,7 +569,7 @@ mod tests { .next() .expect("sender chain key iteration should succeed"); - state.set_sender_chain_key(next_sck.clone()); + state.set_sender_chain_key(next_sck); let updated_sck = state .sender_chain_key() @@ -724,4 +764,67 @@ mod tests { assert_eq!(state.chain_id(), 12345); assert!(state.sender_chain_key().is_some()); } + + /// Test that step_with_message_key produces the same results as + /// calling sender_message_key() and next() separately + #[test] + fn test_step_with_message_key_equivalence() { + let chain = [0x99u8; 32]; + let sck = SenderChainKey::new(5, chain); + + // Get results using separate calls + let msg_key_separate = sck.sender_message_key(); + let next_chain_separate = sck.next().expect("next should succeed"); + + // Get results using optimized combined call + let (msg_key_combined, next_chain_combined) = sck + .step_with_message_key() + .expect("step_with_message_key should succeed"); + + // Verify message keys are identical + assert_eq!(msg_key_separate.iteration(), msg_key_combined.iteration()); + assert_eq!(msg_key_separate.iv(), msg_key_combined.iv()); + assert_eq!(msg_key_separate.cipher_key(), msg_key_combined.cipher_key()); + + // Verify next chain key is identical + assert_eq!(next_chain_separate.seed(), next_chain_combined.seed()); + assert_eq!( + next_chain_separate.iteration(), + next_chain_combined.iteration() + ); + } + + /// Test step_with_message_key over multiple iterations + #[test] + fn test_step_with_message_key_chain() { + let initial_chain = [0xBBu8; 32]; + let mut chain_separate = SenderChainKey::new(0, initial_chain); + let mut chain_combined = SenderChainKey::new(0, initial_chain); + + // Step both chains 10 times and verify they stay in sync + for i in 0..10 { + let msg_key_sep = chain_separate.sender_message_key(); + chain_separate = chain_separate.next().expect("next should succeed"); + + let (msg_key_comb, next_chain) = chain_combined + .step_with_message_key() + .expect("step_with_message_key should succeed"); + chain_combined = next_chain; + + // Verify message keys match + assert_eq!( + msg_key_sep.cipher_key(), + msg_key_comb.cipher_key(), + "cipher_key mismatch at iteration {i}" + ); + + // Verify chain keys match + assert_eq!( + chain_separate.seed(), + chain_combined.seed(), + "chain key mismatch at iteration {i}" + ); + assert_eq!(chain_separate.iteration(), chain_combined.iteration()); + } + } } diff --git a/wacore/libsignal/src/protocol/session_cipher.rs b/wacore/libsignal/src/protocol/session_cipher.rs index 64a6bcf82..a34b8ccd4 100644 --- a/wacore/libsignal/src/protocol/session_cipher.rs +++ b/wacore/libsignal/src/protocol/session_cipher.rs @@ -78,7 +78,8 @@ pub async fn message_encrypt( let chain_key = session_state.get_sender_chain_key()?; - let message_keys = chain_key.message_keys().generate_keys(); + let (message_keys_gen, next_chain_key) = chain_key.step_with_message_keys(); + let message_keys = message_keys_gen.generate_keys(); let sender_ephemeral = session_state.sender_ratchet_key()?; let previous_counter = session_state.previous_counter(); @@ -153,7 +154,7 @@ pub async fn message_encrypt( )?) }; - session_state.set_sender_chain_key(&chain_key.next_chain_key()); + session_state.set_sender_chain_key(&next_chain_key); // XXX why is this check after everything else?!! if !identity_store @@ -817,14 +818,15 @@ fn get_or_create_message_key( } } - let mut chain_key = chain_key.clone(); + let mut chain_key = *chain_key; while chain_key.index() < counter { - let message_keys = chain_key.message_keys(); + let (message_keys, next_chain) = chain_key.step_with_message_keys(); state.set_message_keys(their_ephemeral, message_keys)?; - chain_key = chain_key.next_chain_key(); + chain_key = next_chain; } - state.set_receiver_chain_key(their_ephemeral, &chain_key.next_chain_key())?; - Ok(chain_key.message_keys()) + let (result_message_keys, next_chain) = chain_key.step_with_message_keys(); + state.set_receiver_chain_key(their_ephemeral, &next_chain)?; + Ok(result_message_keys) } diff --git a/wacore/libsignal/src/protocol/state/session.rs b/wacore/libsignal/src/protocol/state/session.rs index 5a9c5a644..65f0c0135 100644 --- a/wacore/libsignal/src/protocol/state/session.rs +++ b/wacore/libsignal/src/protocol/state/session.rs @@ -84,8 +84,8 @@ impl SessionState { Self { session: SessionStructure { session_version: Some(version as u32), - local_identity_public: Some(our_identity.public_key().serialize().into_vec()), - remote_identity_public: Some(their_identity.serialize().into_vec()), + local_identity_public: Some(our_identity.public_key().serialize().to_vec()), + remote_identity_public: Some(their_identity.serialize().to_vec()), root_key: Some(root_key.key().to_vec()), previous_counter: Some(0), sender_chain: None, @@ -93,7 +93,7 @@ impl SessionState { pending_pre_key: None, remote_registration_id: Some(0), local_registration_id: Some(0), - alice_base_key: Some(alice_base_key.serialize().into_vec()), + alice_base_key: Some(alice_base_key.serialize().to_vec()), needs_refresh: None, pending_key_exchange: None, }, @@ -283,9 +283,10 @@ impl SessionState { } pub fn add_receiver_chain(&mut self, sender: &PublicKey, chain_key: &ChainKey) { + use prost::bytes::Bytes; let chain_key = session_structure::chain::ChainKey { index: Some(chain_key.index()), - key: Some(chain_key.key().to_vec()), + key: Some(Bytes::copy_from_slice(chain_key.key())), }; let chain = session_structure::Chain { @@ -297,14 +298,18 @@ impl SessionState { self.session.receiver_chains.push(chain); - if self.session.receiver_chains.len() > consts::MAX_RECEIVER_CHAINS { + // Remove oldest chains if we exceed capacity (MAX_RECEIVER_CHAINS = 5). + // Using drain() for consistency, though with only 5 elements the difference is negligible. + let len = self.session.receiver_chains.len(); + if len > consts::MAX_RECEIVER_CHAINS { log::info!( "Trimming excessive receiver_chain for session with base key {}, chain count: {}", self.sender_ratchet_key_for_logging() .unwrap_or_else(|e| format!("", e.0)), - self.session.receiver_chains.len() + len ); - self.session.receiver_chains.remove(0); + let excess = len - consts::MAX_RECEIVER_CHAINS; + self.session.receiver_chains.drain(..excess); } } @@ -314,9 +319,10 @@ impl SessionState { } pub fn set_sender_chain(&mut self, sender: &KeyPair, next_chain_key: &ChainKey) { + use prost::bytes::Bytes; let chain_key = session_structure::chain::ChainKey { index: Some(next_chain_key.index()), - key: Some(next_chain_key.key().to_vec()), + key: Some(Bytes::copy_from_slice(next_chain_key.key())), }; let new_chain = session_structure::Chain { @@ -365,9 +371,10 @@ impl SessionState { } pub fn set_sender_chain_key(&mut self, next_chain_key: &ChainKey) { + use prost::bytes::Bytes; let chain_key = session_structure::chain::ChainKey { index: Some(next_chain_key.index()), - key: Some(next_chain_key.key().to_vec()), + key: Some(Bytes::copy_from_slice(next_chain_key.key())), }; // Is it actually valid to call this function with sender_chain == None? @@ -450,10 +457,11 @@ impl SessionState { .get_receiver_chain_index(sender)? .expect("called set_receiver_chain_key for a non-existent chain"); + use prost::bytes::Bytes; self.session.receiver_chains[chain_idx].chain_key = Some(session_structure::chain::ChainKey { index: Some(chain_key.index()), - key: Some(chain_key.key().to_vec()), + key: Some(Bytes::copy_from_slice(chain_key.key())), }); Ok(()) diff --git a/wacore/libsignal/src/protocol/state/signed_prekey.rs b/wacore/libsignal/src/protocol/state/signed_prekey.rs index a8d6fd55b..25d309c39 100644 --- a/wacore/libsignal/src/protocol/state/signed_prekey.rs +++ b/wacore/libsignal/src/protocol/state/signed_prekey.rs @@ -169,7 +169,7 @@ impl KeySerde for PublicKey { impl KeySerde for PrivateKey { fn serialize(&self) -> Vec { - self.serialize() + self.serialize().to_vec() } fn deserialize>(bytes: T) -> Result { diff --git a/wacore/libsignal/src/store/record_helpers.rs b/wacore/libsignal/src/store/record_helpers.rs index b98c8fb94..17d229ea8 100644 --- a/wacore/libsignal/src/store/record_helpers.rs +++ b/wacore/libsignal/src/store/record_helpers.rs @@ -9,7 +9,7 @@ pub fn new_pre_key_record(id: u32, key_pair: &KeyPair) -> wa::PreKeyRecordStruct wa::PreKeyRecordStructure { id: Some(id), public_key: Some(key_pair.public_key.public_key_bytes().to_vec()), - private_key: Some(key_pair.private_key.serialize()), + private_key: Some(key_pair.private_key.serialize().to_vec()), } } @@ -22,7 +22,7 @@ pub fn new_signed_pre_key_record( wa::SignedPreKeyRecordStructure { id: Some(id), public_key: Some(key_pair.public_key.public_key_bytes().to_vec()), - private_key: Some(key_pair.private_key.serialize()), + private_key: Some(key_pair.private_key.serialize().to_vec()), signature: Some(signature.to_vec()), timestamp: Some( timestamp @@ -62,7 +62,7 @@ pub fn prekey_record_to_structure( Ok(wa::PreKeyRecordStructure { id: Some(record.id()?.into()), public_key: Some(record.key_pair()?.public_key.public_key_bytes()[1..].to_vec()), - private_key: Some(record.key_pair()?.private_key.serialize()), + private_key: Some(record.key_pair()?.private_key.serialize().to_vec()), }) } diff --git a/wacore/src/handshake/state.rs b/wacore/src/handshake/state.rs index 4f4555c97..df8dc730e 100644 --- a/wacore/src/handshake/state.rs +++ b/wacore/src/handshake/state.rs @@ -57,10 +57,7 @@ impl HandshakeState { self.noise.authenticate(&server_ephemeral)?; self.noise - .mix_shared_secret( - &self.ephemeral_kp.private_key.serialize(), - &server_ephemeral, - ) + .mix_shared_secret(self.ephemeral_kp.private_key.serialize(), &server_ephemeral) .map_err(|e| HandshakeError::Crypto(e.to_string()))?; let static_decrypted = self @@ -74,7 +71,7 @@ impl HandshakeState { self.noise .mix_shared_secret( - &self.ephemeral_kp.private_key.serialize(), + self.ephemeral_kp.private_key.serialize(), &static_decrypted_arr, ) .map_err(|e| HandshakeError::Crypto(e.to_string()))?; @@ -94,7 +91,7 @@ impl HandshakeState { .map_err(|e| HandshakeError::Crypto(e.to_string()))?; self.noise - .mix_shared_secret(&self.static_kp.private_key.serialize(), &server_ephemeral) + .mix_shared_secret(self.static_kp.private_key.serialize(), &server_ephemeral) .map_err(|e| HandshakeError::Crypto(e.to_string()))?; let encrypted_payload = self diff --git a/wacore/src/send.rs b/wacore/src/send.rs index 0e1ac8fd5..4cd30a8dd 100644 --- a/wacore/src/send.rs +++ b/wacore/src/send.rs @@ -857,7 +857,7 @@ pub async fn create_sender_key_distribution_message_for_group( message_version, state.chain_id(), chain_key.iteration(), - chain_key.seed().to_vec(), + *chain_key.seed(), state .signing_key_public() .map_err(|e| anyhow!("Missing pub key: {:?}", e))?, diff --git a/wacore/src/store/device.rs b/wacore/src/store/device.rs index dad3c69e6..fca711352 100644 --- a/wacore/src/store/device.rs +++ b/wacore/src/store/device.rs @@ -20,7 +20,8 @@ pub mod key_pair_serde { let bytes: Vec = key_pair .private_key .serialize() - .into_iter() + .iter() + .copied() .chain(key_pair.public_key.public_key_bytes().iter().copied()) .collect(); serializer.serialize_bytes(&bytes) diff --git a/waproto/build.rs b/waproto/build.rs index badd2d8af..3c2bddc4d 100644 --- a/waproto/build.rs +++ b/waproto/build.rs @@ -32,6 +32,56 @@ fn main() -> std::io::Result<()> { let mut config = prost_build::Config::new(); config.type_attribute(".", "#[derive(serde::Serialize, serde::Deserialize)]"); + // Use bytes::Bytes instead of Vec for frequently-serialized cryptographic structures. + // This enables O(1) cloning (reference-counted) instead of O(n) copying. + // See: https://docs.rs/prost-build/latest/prost_build/struct.Config.html#method.bytes + config.bytes([ + // Session chain keys (called on every message encrypt/decrypt) + ".whatsapp.SessionStructure.Chain.ChainKey", + ".whatsapp.SessionStructure.Chain.MessageKey", + // Sender key structures (group messaging hot path) + ".whatsapp.SenderKeyStateStructure.SenderChainKey", + ".whatsapp.SenderKeyStateStructure.SenderMessageKey", + ".whatsapp.SenderKeyStateStructure.SenderSigningKey", + ]); + + // Skip serde for Bytes fields since bytes::Bytes doesn't implement Serialize/Deserialize + // without the serde feature which prost doesn't expose. These nested types aren't JSON + // serialized anyway - they're stored as protobuf blobs. + // We use skip + default so serde doesn't try to deserialize these fields. + config.field_attribute( + ".whatsapp.SessionStructure.Chain.ChainKey.key", + "#[serde(skip, default)]", + ); + config.field_attribute( + ".whatsapp.SessionStructure.Chain.MessageKey.cipherKey", + "#[serde(skip, default)]", + ); + config.field_attribute( + ".whatsapp.SessionStructure.Chain.MessageKey.macKey", + "#[serde(skip, default)]", + ); + config.field_attribute( + ".whatsapp.SessionStructure.Chain.MessageKey.iv", + "#[serde(skip, default)]", + ); + config.field_attribute( + ".whatsapp.SenderKeyStateStructure.SenderChainKey.seed", + "#[serde(skip, default)]", + ); + config.field_attribute( + ".whatsapp.SenderKeyStateStructure.SenderMessageKey.seed", + "#[serde(skip, default)]", + ); + config.field_attribute( + ".whatsapp.SenderKeyStateStructure.SenderSigningKey.public", + "#[serde(skip, default)]", + ); + config.field_attribute( + ".whatsapp.SenderKeyStateStructure.SenderSigningKey.private", + "#[serde(skip, default)]", + ); + // Configure prost to output the file to the `src/` directory, // so it can be version-controlled. config.out_dir("src/"); diff --git a/waproto/src/whatsapp.rs b/waproto/src/whatsapp.rs index 125a69520..1ce2f11c2 100644 --- a/waproto/src/whatsapp.rs +++ b/waproto/src/whatsapp.rs @@ -10942,24 +10942,28 @@ pub mod sender_key_state_structure { pub struct SenderChainKey { #[prost(uint32, optional, tag="1")] pub iteration: ::core::option::Option, - #[prost(bytes="vec", optional, tag="2")] - pub seed: ::core::option::Option<::prost::alloc::vec::Vec>, + #[prost(bytes="bytes", optional, tag="2")] + #[serde(skip, default)] + pub seed: ::core::option::Option<::prost::bytes::Bytes>, } #[derive(serde::Serialize, serde::Deserialize)] #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct SenderMessageKey { #[prost(uint32, optional, tag="1")] pub iteration: ::core::option::Option, - #[prost(bytes="vec", optional, tag="2")] - pub seed: ::core::option::Option<::prost::alloc::vec::Vec>, + #[prost(bytes="bytes", optional, tag="2")] + #[serde(skip, default)] + pub seed: ::core::option::Option<::prost::bytes::Bytes>, } #[derive(serde::Serialize, serde::Deserialize)] #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct SenderSigningKey { - #[prost(bytes="vec", optional, tag="1")] - pub public: ::core::option::Option<::prost::alloc::vec::Vec>, - #[prost(bytes="vec", optional, tag="2")] - pub private: ::core::option::Option<::prost::alloc::vec::Vec>, + #[prost(bytes="bytes", optional, tag="1")] + #[serde(skip, default)] + pub public: ::core::option::Option<::prost::bytes::Bytes>, + #[prost(bytes="bytes", optional, tag="2")] + #[serde(skip, default)] + pub private: ::core::option::Option<::prost::bytes::Bytes>, } } #[derive(serde::Serialize, serde::Deserialize)] @@ -11019,20 +11023,24 @@ pub mod session_structure { pub struct ChainKey { #[prost(uint32, optional, tag="1")] pub index: ::core::option::Option, - #[prost(bytes="vec", optional, tag="2")] - pub key: ::core::option::Option<::prost::alloc::vec::Vec>, + #[prost(bytes="bytes", optional, tag="2")] + #[serde(skip, default)] + pub key: ::core::option::Option<::prost::bytes::Bytes>, } #[derive(serde::Serialize, serde::Deserialize)] #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct MessageKey { #[prost(uint32, optional, tag="1")] pub index: ::core::option::Option, - #[prost(bytes="vec", optional, tag="2")] - pub cipher_key: ::core::option::Option<::prost::alloc::vec::Vec>, - #[prost(bytes="vec", optional, tag="3")] - pub mac_key: ::core::option::Option<::prost::alloc::vec::Vec>, - #[prost(bytes="vec", optional, tag="4")] - pub iv: ::core::option::Option<::prost::alloc::vec::Vec>, + #[prost(bytes="bytes", optional, tag="2")] + #[serde(skip, default)] + pub cipher_key: ::core::option::Option<::prost::bytes::Bytes>, + #[prost(bytes="bytes", optional, tag="3")] + #[serde(skip, default)] + pub mac_key: ::core::option::Option<::prost::bytes::Bytes>, + #[prost(bytes="bytes", optional, tag="4")] + #[serde(skip, default)] + pub iv: ::core::option::Option<::prost::bytes::Bytes>, } } #[derive(serde::Serialize, serde::Deserialize)]