perf: optimize libsignal reducing allocations

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"] }

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,

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<IdentityKeyPair> {
-        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<u32> {
@@ -333,19 +330,9 @@ impl SignedPreKeyStore for Device {
         signed_prekey_id: u32,
     ) -> Result<Option<SignedPreKeyRecordStructure>, 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(),
             );

storages/sqlite-storage/src/sqlite_store.rs

@@ -148,7 +148,7 @@ impl SqliteStore {
 
     fn serialize_keypair(&self, key_pair: &KeyPair) -> Result<Vec<u8>> {
         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
             };

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<u8> {
+    pub fn serialize(&self) -> &[u8; 32] {
         match &self.key {
-            PrivateKeyData::DjbPrivateKey(v) => v.to_vec(),
+            PrivateKeyData::DjbPrivateKey(v) => v,
         }
     }
 
@@ -254,28 +252,28 @@ impl PrivateKey {
         &self,
         message: &[u8],
         csprng: &mut R,
-    ) -> Result<Box<[u8]>, CurveError> {
+    ) -> Result<[u8; 64], CurveError> {
         self.calculate_signature_for_multipart_message(&[message], csprng)
     }
 
     pub fn calculate_signature_for_multipart_message<R: CryptoRng + Rng>(
         &self,
         message: &[&[u8]],
         csprng: &mut R,
-    ) -> Result<Box<[u8]>, 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<Box<[u8]>, 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<Box<[u8]>, CurveError> {
+    ) -> Result<[u8; 64], CurveError> {
         self.private_key.calculate_signature(message, csprng)
     }
 
-    pub fn calculate_agreement(&self, their_key: &PublicKey) -> Result<Box<[u8]>, CurveError> {
+    pub fn calculate_agreement(&self, their_key: &PublicKey) -> Result<[u8; 32], CurveError> {
         self.private_key.calculate_agreement(their_key)
     }
 }

wacore/libsignal/src/protocol/group_cipher.rs

@@ -17,26 +17,37 @@ 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<u8>,
 }
 
-impl EncryptionBuffer {
+impl CryptoBuffer {
     const INITIAL_CAPACITY: usize = 1024;
 
     fn new() -> Self {
         Self {
             buffer: Vec::with_capacity(Self::INITIAL_CAPACITY),
         }
     }
+
+    /// Clears the buffer and returns a mutable reference for writing.
     fn get_buffer(&mut self) -> &mut Vec<u8> {
         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<u8> {
+        std::mem::replace(&mut self.buffer, Vec::with_capacity(Self::INITIAL_CAPACITY))
+    }
 }
 
 thread_local! {
-    static ENCRYPTION_BUFFER: RefCell<EncryptionBuffer> = RefCell::new(EncryptionBuffer::new());
+    static ENCRYPTION_BUFFER: RefCell<CryptoBuffer> = RefCell::new(CryptoBuffer::new());
+    static DECRYPTION_BUFFER: RefCell<CryptoBuffer> = RefCell::new(CryptoBuffer::new());
 }
 
 pub async fn group_encrypt<R: Rng + CryptoRng>(
@@ -69,7 +80,7 @@ pub async fn group_encrypt<R: Rng + CryptoRng>(
         .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<R: Rng + CryptoRng>(
                 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::<Vec<u8>, SignalProtocolError>(result)
+        Ok::<Vec<u8>, SignalProtocolError>(buf_wrapper.take_buffer())
     })?;
 
     let signing_key = sender_key_state
@@ -98,7 +106,7 @@ pub async fn group_encrypt<R: Rng + CryptoRng>(
         &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<SenderMe
     let mut sender_chain_key = sender_chain_key;
 
     while sender_chain_key.iteration() < iteration {
-        state.add_sender_message_key(&sender_chain_key.sender_message_key());
-        sender_chain_key = sender_chain_key.next()?;
+        let (message_key, next_chain) = sender_chain_key.step_with_message_key()?;
+        state.add_sender_message_key(&message_key);
+        sender_chain_key = next_chain;
     }
 
-    state.set_sender_chain_key(sender_chain_key.next()?);
-    Ok(sender_chain_key.sender_message_key())
+    let (result_message_key, next_chain) = sender_chain_key.step_with_message_key()?;
+    state.set_sender_chain_key(next_chain);
+    Ok(result_message_key)
 }
 
 pub async fn group_decrypt(
@@ -201,31 +211,35 @@ pub async fn group_decrypt(
 
     let sender_key = get_sender_key(sender_key_state, skm.iteration())?;
 
-    let mut plaintext = Vec::new();
-    if let Err(e) = aes_256_cbc_decrypt_into(
-        skm.ciphertext(),
-        sender_key.cipher_key(),
-        sender_key.iv(),
-        &mut plaintext,
-    ) {
-        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",
-                ));
+    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::<Vec<u8>, 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<R: Rng + CryptoRng>(
         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)?,

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<Box<[u8]>> {
+    ) -> Result<[u8; 64]> {
         Ok(self.private_key.calculate_signature_for_multipart_message(
             &[
                 ALTERNATE_IDENTITY_SIGNATURE_PREFIX_1,