From 365ac1df1ef24162a5dd351148abbef70df252e5 Mon Sep 17 00:00:00 2001
From: Mark Rousskov <thismark@amazon.com>
Date: Mon, 23 Sep 2024 16:30:52 +0000
Subject: [PATCH] Remove caching of hmac control key

This shrinks the Entry state by ~700 bytes.
---
 dc/s2n-quic-dc/src/path/secret/map.rs      |  4 ++--
 dc/s2n-quic-dc/src/path/secret/map/test.rs |  1 +
 dc/s2n-quic-dc/src/path/secret/sender.rs   | 21 +++++----------------
 3 files changed, 8 insertions(+), 18 deletions(-)

diff --git a/dc/s2n-quic-dc/src/path/secret/map.rs b/dc/s2n-quic-dc/src/path/secret/map.rs
index 3a2271994..9ad753999 100644
--- a/dc/s2n-quic-dc/src/path/secret/map.rs
+++ b/dc/s2n-quic-dc/src/path/secret/map.rs
@@ -437,7 +437,7 @@ impl Map {
 
         match packet {
             control::Packet::StaleKey(packet) => {
-                let Some(packet) = packet.authenticate(key) else {
+                let Some(packet) = packet.authenticate(&key) else {
                     return;
                 };
                 state.mark_live(self.state.cleaner.epoch());
@@ -447,7 +447,7 @@ impl Map {
                     .fetch_add(1, Ordering::Relaxed);
             }
             control::Packet::ReplayDetected(packet) => {
-                let Some(_packet) = packet.authenticate(key) else {
+                let Some(_packet) = packet.authenticate(&key) else {
                     return;
                 };
                 self.state
diff --git a/dc/s2n-quic-dc/src/path/secret/map/test.rs b/dc/s2n-quic-dc/src/path/secret/map/test.rs
index 81ff02da8..da389a484 100644
--- a/dc/s2n-quic-dc/src/path/secret/map/test.rs
+++ b/dc/s2n-quic-dc/src/path/secret/map/test.rs
@@ -310,6 +310,7 @@ fn check_invariants_no_overflow() {
 }
 
 #[test]
+#[cfg(all(target_pointer_width = "64", target_os = "linux"))]
 fn entry_size() {
     // Don't run this outside of GHA -- it is likely to break on random changes (e.g., new Rust
     // versions) and so is expected to be somewhat flaky. Ideally we'd have something like "Is this
diff --git a/dc/s2n-quic-dc/src/path/secret/sender.rs b/dc/s2n-quic-dc/src/path/secret/sender.rs
index c3e5eddb4..99785c88b 100644
--- a/dc/s2n-quic-dc/src/path/secret/sender.rs
+++ b/dc/s2n-quic-dc/src/path/secret/sender.rs
@@ -3,7 +3,6 @@
 
 use super::schedule;
 use crate::{crypto::awslc::open, packet::secret_control};
-use once_cell::sync::OnceCell;
 use s2n_quic_core::varint::VarInt;
 use std::sync::atomic::{AtomicU64, Ordering};
 
@@ -13,16 +12,6 @@ type StatelessReset = [u8; secret_control::TAG_LEN];
 pub struct State {
     current_id: AtomicU64,
     pub(super) stateless_reset: StatelessReset,
-    control_secret: OnceCell<open::control::Secret>,
-}
-
-impl super::map::SizeOf for OnceCell<open::control::Secret> {
-    fn size(&self) -> usize {
-        // FIXME: OnceCell stores the value inline, but it also has a AtomicPtr to a list of
-        // waiters. That should be ~empty after init finishes. We should probably just initialize
-        // this eagerly when creating the secret rather than adding extra mutable state.
-        std::mem::size_of::<Self>()
-    }
 }
 
 impl super::map::SizeOf for StatelessReset {}
@@ -32,9 +21,8 @@ impl super::map::SizeOf for State {
         let State {
             current_id,
             stateless_reset,
-            control_secret,
         } = self;
-        current_id.size() + stateless_reset.size() + control_secret.size()
+        current_id.size() + stateless_reset.size()
     }
 }
 
@@ -43,7 +31,6 @@ impl State {
         Self {
             current_id: AtomicU64::new(0),
             stateless_reset,
-            control_secret: Default::default(),
         }
     }
 
@@ -69,8 +56,10 @@ impl State {
     }
 
     #[inline]
-    pub fn control_secret(&self, secret: &schedule::Secret) -> &open::control::Secret {
-        self.control_secret.get_or_init(|| secret.control_opener())
+    pub fn control_secret(&self, secret: &schedule::Secret) -> open::control::Secret {
+        // We don't try to cache this, hmac init is cheap (~200-600ns depending on algorithm) and
+        // the space requirement is huge (700+ bytes)
+        secret.control_opener()
     }
 
     /// Update the sender for a received stale key packet.