perf: optimize libsignal reducing allocations

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

@@ -192,7 +192,7 @@ macro_rules! impl_store_wrapper {
 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 private_key = PrivateKey::deserialize(private_key_bytes.serialize())?;
         let ikp = IdentityKeyPair::try_from(private_key)?;
         Ok(ikp)
     }
@@ -339,7 +339,7 @@ impl SignedPreKeyStore for Device {
                 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())
+            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);
 

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

@@ -37,6 +37,7 @@ impl EncryptionBuffer {
 
 thread_local! {
     static ENCRYPTION_BUFFER: RefCell<EncryptionBuffer> = RefCell::new(EncryptionBuffer::new());
+    static DECRYPTION_BUFFER: RefCell<EncryptionBuffer> = RefCell::new(EncryptionBuffer::new());
 }
 
 pub async fn group_encrypt<R: Rng + CryptoRng>(
@@ -69,7 +70,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();
@@ -98,7 +99,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 +142,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 +204,37 @@ 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",
+                    ));
+                }
             }
         }
-    }
+        let result = std::mem::take(buf);
+        buf.reserve(EncryptionBuffer::INITIAL_CAPACITY);
+        Ok::<Vec<u8>, SignalProtocolError>(result)
+    })?;
 
     sender_key_store
         .store_sender_key(sender_key_name, &record)
@@ -313,7 +322,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,