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,7 +15,7 @@ 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,

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,
+            },
+        )))
     }
 }
 

wacore/noise/src/state.rs

@@ -2,8 +2,8 @@ use crate::error::{NoiseError, Result};
 use hkdf::Hkdf;
 use sha2::{Digest, Sha256};
 use wacore_libsignal::crypto::{
-    GcmInPlaceBuffer, aes_256_gcm_decrypt, aes_256_gcm_decrypt_in_place, aes_256_gcm_encrypt,
-    aes_256_gcm_encrypt_in_place,
+    Aes256GcmDecryption, Aes256GcmEncryption, Aes256GcmKey, CryptoProviderError, GcmInPlaceBuffer,
+    aes_256_gcm_decrypt, aes_256_gcm_encrypt,
 };
 
 /// Buffer kinds accepted by [`NoiseCipher::decrypt_in_place_with_counter`].
@@ -25,22 +25,31 @@ const TAG_LEN: usize = 16;
 /// A cipher wrapper that encapsulates AES-256-GCM encryption/decryption
 /// with counter-based IV generation.
 pub struct NoiseCipher {
-    key: [u8; 32],
+    /// Pre-keyed GCM state: the transport key is fixed for the connection,
+    /// so the AES key schedule and the GHASH subkey are derived once here
+    /// instead of on every frame.
+    key: Aes256GcmKey,
 }
 
 impl NoiseCipher {
     /// Creates a new cipher from a 32-byte key.
     pub fn new(key: &[u8; 32]) -> Result<Self> {
-        Ok(Self { key: *key })
+        Ok(Self {
+            key: Aes256GcmKey::new(key)
+                .map_err(|_| NoiseError::Encrypt(CryptoProviderError::BadInput))?,
+        })
     }
 
     /// Encrypts plaintext using the specified counter for IV generation.
     /// Returns the ciphertext with appended authentication tag (16 bytes).
     pub fn encrypt_with_counter(&self, counter: u32, plaintext: &[u8]) -> Result<Vec<u8>> {
         let iv = generate_iv(counter);
         let mut out = Vec::with_capacity(plaintext.len() + TAG_LEN);
-        aes_256_gcm_encrypt(&self.key, &iv, b"", plaintext, &mut out)
-            .map_err(NoiseError::Encrypt)?;
+        out.extend_from_slice(plaintext);
+        let mut enc = Aes256GcmEncryption::new_with_key(&self.key, &iv, b"")
+            .map_err(|_| NoiseError::Encrypt(CryptoProviderError::BadInput))?;
+        enc.encrypt(&mut out);
+        out.extend_from_slice(&enc.compute_tag());
         Ok(out)
     }
 
@@ -54,7 +63,14 @@ impl NoiseCipher {
         buffer: &mut B,
     ) -> Result<()> {
         let iv = generate_iv(counter);
-        aes_256_gcm_encrypt_in_place(&self.key, &iv, b"", buffer).map_err(NoiseError::Encrypt)
+        let plaintext_len = buffer.len();
+        let mut enc = Aes256GcmEncryption::new_with_key(&self.key, &iv, b"")
+            .map_err(|_| NoiseError::Encrypt(CryptoProviderError::BadInput))?;
+        enc.encrypt(buffer.as_mut_slice());
+        let tag = enc.compute_tag();
+        buffer.resize(plaintext_len + TAG_LEN, 0);
+        buffer.as_mut_slice()[plaintext_len..].copy_from_slice(&tag);
+        Ok(())
     }
 
     /// Decrypts ciphertext (with 16-byte tag appended) in-place within the
@@ -67,7 +83,21 @@ impl NoiseCipher {
         buffer: &mut B,
     ) -> Result<()> {
         let iv = generate_iv(counter);
-        aes_256_gcm_decrypt_in_place(&self.key, &iv, b"", buffer).map_err(NoiseError::Decrypt)
+        let total = buffer.len();
+        if total < TAG_LEN {
+            return Err(NoiseError::Decrypt(CryptoProviderError::BadInput));
+        }
+        let pt_len = total - TAG_LEN;
+        let mut tag = [0u8; TAG_LEN];
+        tag.copy_from_slice(&buffer.as_slice()[pt_len..]);
+
+        let mut dec = Aes256GcmDecryption::new_with_key(&self.key, &iv, b"")
+            .map_err(|_| NoiseError::Decrypt(CryptoProviderError::BadInput))?;
+        dec.decrypt(&mut buffer.as_mut_slice()[..pt_len]);
+        dec.verify_tag(&tag)
+            .map_err(|_| NoiseError::Decrypt(CryptoProviderError::AuthFailed))?;
+        buffer.truncate(pt_len);
+        Ok(())
     }
 }