feat(s2n-quic-dc): Further shrink path secret entry (#2339)

aws · Oct 1, 2024 · cb5087c · cb5087c
1 parent 2a735a9
commit cb5087c
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 38 deletions.
diff --git a/dc/s2n-quic-dc/src/path/secret/map.rs b/dc/s2n-quic-dc/src/path/secret/map.rs
@@ -297,8 +297,6 @@ impl Map {
         peer: SocketAddr,
     ) -> Option<(seal::Once, Credentials, ApplicationParams)> {
         let state = self.state.peers.get_by_key(&peer)?;
-        state.mark_live(self.state.cleaner.epoch());
-
         let (sealer, credentials) = state.uni_sealer();
         Some((sealer, credentials, state.parameters))
     }
@@ -319,8 +317,6 @@ impl Map {
         features: &TransportFeatures,
     ) -> Option<(Bidirectional, ApplicationParams)> {
         let state = self.state.peers.get_by_key(&peer)?;
-        state.mark_live(self.state.cleaner.epoch());
-
         let keys = state.bidi_local(features);
 
         Some((keys, state.parameters))
@@ -399,7 +395,6 @@ impl Map {
                 let Some(packet) = packet.authenticate(&key) else {
                     return;
                 };
-                state.mark_live(self.state.cleaner.epoch());
                 state.sender.update_for_stale_key(packet.min_key_id);
                 self.state
                     .handled_control_packets
@@ -447,7 +442,6 @@ impl Map {
             packet.encode(encoder, &stateless_reset);
             return None;
         };
-        state.mark_live(self.state.cleaner.epoch());
 
         match state.receiver.pre_authentication(identity) {
             Ok(()) => {}
@@ -465,7 +459,6 @@ impl Map {
     pub(super) fn insert(&self, entry: Arc<Entry>) {
         // On insert clear our interest in a handshake.
         self.state.requested_handshakes.pin().remove(&entry.peer);
-        entry.mark_live(self.state.cleaner.epoch());
         let id = *entry.secret.id();
         let peer = entry.peer;
         if self.state.ids.insert(id, entry.clone()).is_some() {
@@ -613,12 +606,6 @@ pub(super) struct Entry {
     peer: SocketAddr,
     secret: schedule::Secret,
     retired: IsRetired,
-    // Last time the entry was pulled out of the State map.
-    // This is not necessarily the last time the entry was used but it's close enough for our
-    // purposes: if the entry is not being pulled out of the State map, it's hopefully not going to
-    // start getting pulled out shortly. This is used for the LRU mechanism, see the Cleaner impl
-    // for details.
-    used_at: AtomicU64,
     sender: sender::State,
     receiver: receiver::State,
     parameters: ApplicationParams,
@@ -640,7 +627,6 @@ impl SizeOf for Entry {
             peer,
             secret,
             retired,
-            used_at,
             sender,
             receiver,
             parameters,
@@ -650,7 +636,6 @@ impl SizeOf for Entry {
             + peer.size()
             + secret.size()
             + retired.size()
-            + used_at.size()
             + sender.size()
             + receiver.size()
             + parameters.size()
@@ -717,7 +702,6 @@ impl Entry {
             peer,
             secret,
             retired: Default::default(),
-            used_at: AtomicU64::new(0),
             sender,
             receiver,
             parameters,
@@ -728,10 +712,6 @@ impl Entry {
         self.retired.0.store(at_epoch, Ordering::Relaxed);
     }
 
-    fn mark_live(&self, at_epoch: u64) {
-        self.used_at.store(at_epoch, Ordering::Relaxed);
-    }
-
     fn uni_sealer(&self) -> (seal::Once, Credentials) {
         let key_id = self.sender.next_key_id();
         let credentials = Credentials {

diff --git a/dc/s2n-quic-dc/src/path/secret/map/test.rs b/dc/s2n-quic-dc/src/path/secret/map/test.rs
@@ -327,6 +327,6 @@ fn no_memory_growth() {
 fn entry_size() {
     // This gates to running only on specific GHA to reduce false positives.
     if std::env::var("S2N_QUIC_RUN_VERSION_SPECIFIC_TESTS").is_ok() {
-        assert_eq!(fake_entry(0).size(), 350);
+        assert_eq!(fake_entry(0).size(), 270);
     }
 }
diff --git a/dc/s2n-quic-dc/src/path/secret/schedule.rs b/dc/s2n-quic-dc/src/path/secret/schedule.rs
@@ -11,8 +11,10 @@ use aws_lc_rs::{
     hmac,
 };
 use s2n_quic_core::{dc, varint::VarInt};
+use zeroize::Zeroizing;
 
 pub use s2n_quic_core::endpoint;
+
 pub const MAX_KEY_LEN: usize = 32;
 const MAX_HMAC_KEY_LEN: usize = 1024 / 8;
 
@@ -104,32 +106,30 @@ pub type ExportSecret = [u8; 32];
 #[derive(Debug)]
 pub struct Secret {
     id: Id,
-    prk: Prk,
+    export_secret: Zeroizing<ExportSecret>,
     endpoint: endpoint::Type,
     ciphersuite: Ciphersuite,
 }
 
 impl super::map::SizeOf for Id {}
-impl super::map::SizeOf for Prk {
+impl super::map::SizeOf for endpoint::Type {}
+impl super::map::SizeOf for Ciphersuite {}
+impl super::map::SizeOf for Zeroizing<ExportSecret> {
     fn size(&self) -> usize {
-        // FIXME: maybe don't use this type since it has overhead and needs this weird assert to
-        // check the mode?
-        assert!(format!("{:?}", self).contains("mode: Expand"));
+        // Zeroizing uses Drop just for zeroing, but that doesn't add any space.
         std::mem::size_of::<Self>()
     }
 }
-impl super::map::SizeOf for endpoint::Type {}
-impl super::map::SizeOf for Ciphersuite {}
 
 impl super::map::SizeOf for Secret {
     fn size(&self) -> usize {
         let Secret {
             id,
-            prk,
+            export_secret,
             endpoint,
             ciphersuite,
         } = self;
-        id.size() + prk.size() + endpoint.size() + ciphersuite.size()
+        id.size() + export_secret.size() + endpoint.size() + ciphersuite.size()
     }
 }
 
@@ -141,22 +141,27 @@ impl Secret {
         endpoint: endpoint::Type,
         export_secret: &ExportSecret,
     ) -> Self {
-        let prk = Prk::new_less_safe(ciphersuite.hkdf(), export_secret);
-
         let mut v = Self {
             id: Default::default(),
-            prk,
+            export_secret: Zeroizing::new(*export_secret),
             endpoint,
             ciphersuite,
         };
 
         let mut id = Id::default();
-        v.prk.expand_into(&[&[16], b" pid"], &mut *id);
+        v.prk().expand_into(&[&[16], b" pid"], &mut *id);
         v.id = id;
 
         v
     }
 
+    // Note that Prk doesn't allocate when constructed with new_less_safe (or even traverse to C),
+    // but we can store it in far less space (104 -> 32 bytes) if we store just the secret
+    // directly.
+    fn prk(&self) -> Prk {
+        Prk::new_less_safe(self.ciphersuite.hkdf(), &*self.export_secret)
+    }
+
     #[inline]
     pub fn id(&self) -> &Id {
         &self.id
@@ -176,7 +181,7 @@ impl Secret {
         debug_assert!(out_len <= u16::MAX as usize);
 
         let (out, _) = out.split_at_mut(out_len);
-        self.prk.expand_into(
+        self.prk().expand_into(
             &[
                 &(out_len as u16).to_be_bytes(),
                 b" bidi",
@@ -244,7 +249,7 @@ impl Secret {
         debug_assert!(out_len <= u16::MAX as usize);
 
         let (out, _) = out.split_at_mut(out_len);
-        self.prk.expand_into(
+        self.prk().expand_into(
             &[
                 &(out_len as u16).to_be_bytes(),
                 b" bidi",
@@ -293,7 +298,7 @@ impl Secret {
         debug_assert!(out_len <= u16::MAX as usize);
 
         let (out, _) = out.split_at_mut(out_len);
-        self.prk.expand_into(
+        self.prk().expand_into(
             &[
                 &(out_len as u16).to_be_bytes(),
                 b" uni",
@@ -340,7 +345,7 @@ impl Secret {
         debug_assert!(out_len <= u16::MAX as usize);
 
         let (out, _) = out.split_at_mut(out_len);
-        self.prk.expand_into(
+        self.prk().expand_into(
             &[
                 &(out_len as u16).to_be_bytes(),
                 b" ctl",