perf(noise): pre-key the transport AES-GCM once per connection

wacore/libsignal/src/crypto/aes_gcm.rs

@@ -26,18 +26,24 @@ struct GcmGhash {
 
 impl GcmGhash {
     fn new(h: &[u8; TAG_SIZE], ghash_pad: [u8; TAG_SIZE], associated_data: &[u8]) -> Result<Self> {
-        let mut ghash = GHash::new(h.into());
+        Ok(Self::from_keyed(
+            GHash::new(h.into()),
+            ghash_pad,
+            associated_data,
+        ))
+    }
 
+    fn from_keyed(mut ghash: GHash, ghash_pad: [u8; TAG_SIZE], associated_data: &[u8]) -> Self {
         ghash.update_padded(associated_data);
 
-        Ok(Self {
+        Self {
             ghash,
             ghash_pad,
             msg_buf: [0u8; TAG_SIZE],
             msg_buf_offset: 0,
             ad_len: associated_data.len(),
             msg_len: 0,
-        })
+        }
     }
 
     fn update(&mut self, msg: &[u8]) {
@@ -122,6 +128,43 @@ fn setup_gcm(key: &[u8], nonce: &[u8], associated_data: &[u8]) -> Result<(Aes256
     Ok((ctr, ghash))
 }
 
+/// Key-dependent AES-256-GCM state computed once: the AES key schedule and
+/// the keyed GHASH outlive individual seals when the key never changes (the
+/// Noise transport key lives for a whole connection), leaving only the
+/// nonce-dependent setup per call.
+#[derive(Clone)]
+pub struct Aes256GcmKey {
+    cipher: Aes256,
+    keyed_ghash: GHash,
+}
+
+impl Aes256GcmKey {
+    pub fn new(key: &[u8]) -> Result<Self> {
+        let cipher = Aes256::new_from_slice(key).map_err(|_| Error::InvalidKeySize)?;
+        let mut h: Block<Aes256> = [0u8; TAG_SIZE].into();
+        cipher.encrypt_block(&mut h);
+        let h: [u8; TAG_SIZE] = h.into();
+        Ok(Self {
+            cipher,
+            keyed_ghash: GHash::new((&h).into()),
+        })
+    }
+
+    fn setup(&self, nonce: &[u8], associated_data: &[u8]) -> Result<(Aes256Ctr32, GcmGhash)> {
+        if nonce.len() != NONCE_SIZE {
+            return Err(Error::InvalidNonceSize);
+        }
+
+        let mut ctr = Aes256Ctr32::new(self.cipher.clone(), nonce, 1)?;
+
+        let mut ghash_pad = [0u8; 16];
+        ctr.process(&mut ghash_pad);
+
+        let ghash = GcmGhash::from_keyed(self.keyed_ghash.clone(), ghash_pad, associated_data);
+        Ok((ctr, ghash))
+    }
+}
+
 pub struct Aes256GcmEncryption {
     ctr: Aes256Ctr32,
     ghash: GcmGhash,
@@ -136,6 +179,11 @@ impl Aes256GcmEncryption {
         Ok(Self { ctr, ghash })
     }
 
+    pub fn new_with_key(key: &Aes256GcmKey, nonce: &[u8], associated_data: &[u8]) -> Result<Self> {
+        let (ctr, ghash) = key.setup(nonce, associated_data)?;
+        Ok(Self { ctr, ghash })
+    }
+
     pub fn encrypt(&mut self, buf: &mut [u8]) {
         self.ctr.process(buf);
         self.ghash.update(buf);
@@ -160,6 +208,11 @@ impl Aes256GcmDecryption {
         Ok(Self { ctr, ghash })
     }
 
+    pub fn new_with_key(key: &Aes256GcmKey, nonce: &[u8], associated_data: &[u8]) -> Result<Self> {
+        let (ctr, ghash) = key.setup(nonce, associated_data)?;
+        Ok(Self { ctr, ghash })
+    }
+
     pub fn decrypt(&mut self, buf: &mut [u8]) {
         self.ghash.update(buf);
         self.ctr.process(buf);
@@ -508,4 +561,44 @@ mod tests {
         dec.verify_tag(&expected_tag).unwrap();
         assert_eq!(decrypted, plaintext);
     }
+
+    /// The pre-keyed path must stay byte-identical to the per-call setup;
+    /// any drift would corrupt every Noise frame.
+    #[test]
+    fn pre_keyed_matches_per_call_setup() {
+        let key = [0x42u8; 32];
+        let pre_keyed = Aes256GcmKey::new(&key).unwrap();
+
+        for (i, len) in [0usize, 1, 15, 16, 17, 1500, 4096].iter().enumerate() {
+            let mut nonce = [0u8; NONCE_SIZE];
+            nonce[11] = i as u8;
+            let aad: &[u8] = if i % 2 == 0 { b"" } else { b"associated" };
+            let plaintext: Vec<u8> = (0..*len).map(|b| b as u8).collect();
+
+            let mut plain_ct = plaintext.clone();
+            let mut enc = Aes256GcmEncryption::new(&key, &nonce, aad).unwrap();
+            enc.encrypt(&mut plain_ct);
+            let plain_tag = enc.compute_tag();
+
+            let mut pk_ct = plaintext.clone();
+            let mut enc = Aes256GcmEncryption::new_with_key(&pre_keyed, &nonce, aad).unwrap();
+            enc.encrypt(&mut pk_ct);
+            let pk_tag = enc.compute_tag();
+
+            assert_eq!(plain_ct, pk_ct);
+            assert_eq!(plain_tag, pk_tag);
+
+            let mut dec = Aes256GcmDecryption::new_with_key(&pre_keyed, &nonce, aad).unwrap();
+            dec.decrypt(&mut pk_ct);
+            dec.verify_tag(&pk_tag).unwrap();
+            assert_eq!(pk_ct, plaintext);
+
+            let mut bad_tag = pk_tag;
+            bad_tag[0] ^= 1;
+            let mut ct_again = plain_ct.clone();
+            let mut dec = Aes256GcmDecryption::new_with_key(&pre_keyed, &nonce, aad).unwrap();
+            dec.decrypt(&mut ct_again);
+            assert!(dec.verify_tag(&bad_tag).is_err());
+        }
+    }
 }

wacore/libsignal/src/crypto/mod.rs

@@ -15,16 +15,26 @@ pub use aes_cbc::{
     DecryptionError, EncryptionError, aes_256_cbc_decrypt_into, aes_256_cbc_encrypt_into,
 };
 pub use aes_ctr::Aes256Ctr32;
-pub use aes_gcm::{Aes256GcmDecryption, Aes256GcmEncryption};
+pub use aes_gcm::{Aes256GcmDecryption, Aes256GcmEncryption, Aes256GcmKey};
 pub use error::{Error, Result};
 pub use hash::{
     CryptographicHash, CryptographicMac, SHA1_OUTPUT_SIZE, SHA256_OUTPUT_SIZE, SHA512_OUTPUT_SIZE,
 };
 pub use provider::{
-    CryptoProviderError, GcmInPlaceBuffer, RustCryptoProvider, SignalCryptoProvider,
+    CryptoProviderError, GcmInPlaceBuffer, RustCryptoProvider, SignalCryptoProvider, TransportAead,
     set_crypto_provider,
 };
 
+/// Connection-lifetime transport AEAD for one fixed key, from the active
+/// [`SignalCryptoProvider`]. The default RustCrypto path precomputes the
+/// key-dependent state once.
+#[inline]
+pub fn transport_aead(
+    key: &[u8; 32],
+) -> std::result::Result<Box<dyn TransportAead>, CryptoProviderError> {
+    provider::provider().transport_aead(key)
+}
+
 /// AES-256-GCM seal. Appends `ciphertext || tag(16)` to `out`.
 /// Delegates to the active [`SignalCryptoProvider`].
 #[inline]

wacore/libsignal/src/crypto/provider.rs

@@ -13,7 +13,7 @@ use bytes::BytesMut;
 use hmac::{Hmac, KeyInit, Mac};
 use sha2::Sha256;
 
-use crate::crypto::aes_gcm::{Aes256GcmDecryption, Aes256GcmEncryption};
+use crate::crypto::aes_gcm::{Aes256GcmDecryption, Aes256GcmEncryption, Aes256GcmKey};
 
 const GCM_TAG: usize = 16;
 
@@ -83,6 +83,55 @@ pub enum CryptoProviderError {
     BackendFailed,
 }
 
+/// Connection-lifetime AES-256-GCM handle for the Noise transport: one key,
+/// many nonces. Lets a provider precompute key-dependent state once instead
+/// of per frame; the default routes every call through the configured
+/// provider so custom providers keep observing transport crypto.
+pub trait TransportAead: Send + Sync {
+    fn encrypt_in_place(
+        &self,
+        nonce: &[u8; 12],
+        aad: &[u8],
+        buffer: &mut dyn GcmInPlaceBuffer,
+    ) -> Result<(), CryptoProviderError>;
+
+    fn decrypt_in_place(
+        &self,
+        nonce: &[u8; 12],
+        aad: &[u8],
+        buffer: &mut dyn GcmInPlaceBuffer,
+    ) -> Result<(), CryptoProviderError>;
+}
+
+/// Default [`TransportAead`]: per-call dispatch through the provider that
+/// created it, with no precomputed state.
+struct PerCallTransportAead<P: ?Sized + 'static> {
+    provider: &'static P,
+    key: [u8; 32],
+}
+
+impl<P: SignalCryptoProvider + ?Sized> TransportAead for PerCallTransportAead<P> {
+    fn encrypt_in_place(
+        &self,
+        nonce: &[u8; 12],
+        aad: &[u8],
+        buffer: &mut dyn GcmInPlaceBuffer,
+    ) -> Result<(), CryptoProviderError> {
+        self.provider
+            .aes_256_gcm_encrypt_in_place(&self.key, nonce, aad, buffer)
+    }
+
+    fn decrypt_in_place(
+        &self,
+        nonce: &[u8; 12],
+        aad: &[u8],
+        buffer: &mut dyn GcmInPlaceBuffer,
+    ) -> Result<(), CryptoProviderError> {
+        self.provider
+            .aes_256_gcm_decrypt_in_place(&self.key, nonce, aad, buffer)
+    }
+}
+
 /// Pluggable crypto primitives used by libsignal and higher-level callers.
 ///
 /// All methods are one-shot (no streaming state): inputs fully known at call
@@ -171,6 +220,20 @@ pub trait SignalCryptoProvider: Send + Sync + 'static {
         buffer.as_mut_slice().copy_from_slice(&out);
         Ok(())
     }
+
+    /// Connection-lifetime transport AEAD for one fixed key. Default keeps
+    /// per-call dispatch through this same provider; override to precompute
+    /// key-dependent state (key schedule, GHASH subkey) once. The `'static`
+    /// receiver ties the handle to the installed provider's lifetime.
+    fn transport_aead(
+        &'static self,
+        key: &[u8; 32],
+    ) -> Result<Box<dyn TransportAead>, CryptoProviderError> {
+        Ok(Box::new(PerCallTransportAead {
+            provider: self,
+            key: *key,
+        }))
+    }
 }
 
 static CRYPTO_PROVIDER: OnceLock<Box<dyn SignalCryptoProvider>> = OnceLock::new();
@@ -335,6 +398,56 @@ impl SignalCryptoProvider for RustCryptoProvider {
         buffer.truncate(pt_len);
         Ok(())
     }
+
+    fn transport_aead(
+        &'static self,
+        key: &[u8; 32],
+    ) -> Result<Box<dyn TransportAead>, CryptoProviderError> {
+        Ok(Box::new(
+            Aes256GcmKey::new(key).map_err(|_| CryptoProviderError::BadInput)?,
+        ))
+    }
+}
+
+impl TransportAead for Aes256GcmKey {
+    fn encrypt_in_place(
+        &self,
+        nonce: &[u8; 12],
+        aad: &[u8],
+        buffer: &mut dyn GcmInPlaceBuffer,
+    ) -> Result<(), CryptoProviderError> {
+        let plaintext_len = buffer.len();
+        let mut enc = Aes256GcmEncryption::new_with_key(self, nonce, aad)
+            .map_err(|_| CryptoProviderError::BadInput)?;
+        enc.encrypt(buffer.as_mut_slice());
+        let tag = enc.compute_tag();
+        buffer.resize(plaintext_len + GCM_TAG, 0);
+        buffer.as_mut_slice()[plaintext_len..].copy_from_slice(&tag);
+        Ok(())
+    }
+
+    fn decrypt_in_place(
+        &self,
+        nonce: &[u8; 12],
+        aad: &[u8],
+        buffer: &mut dyn GcmInPlaceBuffer,
+    ) -> Result<(), CryptoProviderError> {
+        let total = buffer.len();
+        if total < GCM_TAG {
+            return Err(CryptoProviderError::BadInput);
+        }
+        let pt_len = total - GCM_TAG;
+        let mut tag = [0u8; GCM_TAG];
+        tag.copy_from_slice(&buffer.as_slice()[pt_len..]);
+
+        let mut dec = Aes256GcmDecryption::new_with_key(self, nonce, aad)
+            .map_err(|_| CryptoProviderError::BadInput)?;
+        dec.decrypt(&mut buffer.as_mut_slice()[..pt_len]);
+        dec.verify_tag(&tag)
+            .map_err(|_| CryptoProviderError::AuthFailed)?;
+        buffer.truncate(pt_len);
+        Ok(())
+    }
 }
 
 #[cfg(test)]

wacore/noise/src/handshake.rs

@@ -380,7 +380,7 @@ pub struct IkFallbackInputs {
 pub enum IkServerHelloOutcome {
     /// Server accepted IK; cert payload was decrypted and the cipher keys
     /// are derivable. The cached cert chain on the device stays untouched.
-    Continue(IkHandshakeOutcome),
+    Continue(Box<IkHandshakeOutcome>),
     /// Server rejected IK by replying with `static != null`. Switch to
     /// XX-fallback using the carryover inputs. Boxed since `IkFallbackInputs`
     /// holds two KeyPairs and a payload, dwarfing `Continue`'s two ciphers.
@@ -630,10 +630,12 @@ impl IkHandshakeState {
         let _cert_plaintext = noise.decrypt(&cert_payload)?;
 
         let (write_cipher, read_cipher) = noise.finish()?;
-        Ok(IkServerHelloOutcome::Continue(IkHandshakeOutcome {
-            write_cipher,
-            read_cipher,
-        }))
+        Ok(IkServerHelloOutcome::Continue(Box::new(
+            IkHandshakeOutcome {
+                write_cipher,
+                read_cipher,
+            },
+        )))
     }
 }