Use a stack allocation for header protection (#1978)

The use of `Vec` here is unnecessary. We can use a fixed sized array instead. The performance gain here is negligible, but the code becomes cleaner, so that's a win.
mozilla · Jul 12, 2024 · 59dc0ab · 59dc0ab
1 parent f1c04d2
commit 59dc0ab
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 24 deletions.
diff --git a/neqo-crypto/src/hp.rs b/neqo-crypto/src/hp.rs
@@ -57,7 +57,7 @@ impl Debug for HpKey {
 }
 
 impl HpKey {
-    const SAMPLE_SIZE: usize = 16;
+    pub const SAMPLE_SIZE: usize = 16;
 
     /// QUIC-specific API for extracting a header-protection key.
     ///
@@ -124,13 +124,6 @@ impl HpKey {
         Ok(res)
     }
 
-    /// Get the sample size, which is also the output size.
-    #[must_use]
-    #[allow(clippy::unused_self)] // To maintain an API contract.
-    pub const fn sample_size(&self) -> usize {
-        Self::SAMPLE_SIZE
-    }
-
     const fn block_size(&self) -> usize {
         match self {
             Self::Aes(_) => 16,
@@ -148,8 +141,9 @@ impl HpKey {
     /// # Panics
     ///
     /// When the mechanism for our key is not supported.
-    pub fn mask(&self, sample: &[u8]) -> Res<Vec<u8>> {
-        let mut output = vec![0_u8; self.block_size()];
+    /// Or when the sample provided is not at least `self.sample_size()` bytes.
+    pub fn mask(&self, sample: &[u8]) -> Res<[u8; Self::SAMPLE_SIZE]> {
+        let mut output = [0; Self::SAMPLE_SIZE];
 
         match self {
             Self::Aes(context) => {
@@ -164,7 +158,7 @@ impl HpKey {
                         c_int::try_from(Self::SAMPLE_SIZE).unwrap(),
                     )
                 })?;
-                assert_eq!(usize::try_from(output_len).unwrap(), output.len());
+                debug_assert_eq!(usize::try_from(output_len).unwrap(), output.len());
                 Ok(output)
             }
 
@@ -185,11 +179,11 @@ impl HpKey {
                         output[..].as_mut_ptr(),
                         &mut output_len,
                         c_uint::try_from(output.len())?,
-                        output[..].as_ptr(),
-                        c_uint::try_from(output.len())?,
+                        [0; Self::SAMPLE_SIZE].as_ptr(),
+                        c_uint::try_from(Self::SAMPLE_SIZE)?,
                     )
                 })?;
-                assert_eq!(usize::try_from(output_len).unwrap(), output.len());
+                debug_assert_eq!(usize::try_from(output_len).unwrap(), output.len());
                 Ok(output)
             }
         }

diff --git a/neqo-crypto/tests/hp.rs b/neqo-crypto/tests/hp.rs
@@ -61,10 +61,7 @@ fn aes256() {
 fn chacha20_ctr() {
     const EXPECTED: &[u8] = &[
         0x34, 0x11, 0xb3, 0x53, 0x02, 0x0b, 0x16, 0xda, 0x0a, 0x85, 0x5a, 0x52, 0x0d, 0x06, 0x07,
-        0x1f, 0x4a, 0xb1, 0xaf, 0xf7, 0x83, 0xa8, 0xf0, 0x29, 0xc3, 0x19, 0xef, 0x57, 0x48, 0xe7,
-        0x8e, 0x3e, 0x11, 0x91, 0xe1, 0xd5, 0x92, 0x8f, 0x61, 0x6d, 0x3f, 0x3d, 0x76, 0xb5, 0x29,
-        0xf1, 0x62, 0x2f, 0x1e, 0xad, 0xdd, 0x23, 0x59, 0x45, 0xac, 0xd2, 0x19, 0x8a, 0xb4, 0x1f,
-        0x2f, 0x52, 0x46, 0x89,
+        0x1f,
     ];
 
     hp_test(TLS_CHACHA20_POLY1305_SHA256, EXPECTED);

diff --git a/neqo-transport/src/crypto.rs b/neqo-transport/src/crypto.rs
@@ -631,9 +631,9 @@ impl CryptoDxState {
         )
     }
 
-    pub fn compute_mask(&self, sample: &[u8]) -> Res<Vec<u8>> {
+    pub fn compute_mask(&self, sample: &[u8]) -> Res<[u8; HpKey::SAMPLE_SIZE]> {
         let mask = self.hpkey.mask(sample)?;
-        qtrace!([self], "HP sample={} mask={}", hex(sample), hex(&mask));
+        qtrace!([self], "HP sample={} mask={}", hex(sample), hex(mask));
         Ok(mask)
     }
 
@@ -710,9 +710,7 @@ impl CryptoDxState {
     /// This is the difference between the size of the header protection sample
     /// and the AEAD expansion.
     pub const fn extra_padding(&self) -> usize {
-        self.hpkey
-            .sample_size()
-            .saturating_sub(self.aead.expansion())
+        HpKey::SAMPLE_SIZE.saturating_sub(self.aead.expansion())
     }
 }
 

diff --git a/test-fixture/src/header_protection.rs b/test-fixture/src/header_protection.rs
@@ -135,7 +135,7 @@ pub fn apply_header_protection(hp: &HpKey, packet: &mut [u8], pn_bytes: Range<us
     qtrace!(
         "sample={} mask={}",
         hex_with_len(&packet[sample_start..sample_end]),
-        hex_with_len(&mask)
+        hex_with_len(mask)
     );
     packet[0] ^= mask[0] & 0xf;
     for i in 0..(pn_bytes.end - pn_bytes.start) {