Support storing ApplicationData in Entry

This stores arbitrary metadata (Arc<dyn Any>) in each Entry. By default, we share a single `Arc<()>` across all entries, but applications can register a callback to produce arbitrary other values.
aws · Mar 7, 2025 · 465fd1b · 465fd1b
1 parent a9e7673
commit 465fd1b
Show file tree

Hide file tree

Showing 7 changed files with 101 additions and 6 deletions.
diff --git a/dc/s2n-quic-dc/src/path/secret/map.rs b/dc/s2n-quic-dc/src/path/secret/map.rs
@@ -8,6 +8,7 @@ use crate::{
     path::secret::{open, seal, stateless_reset},
     stream::TransportFeatures,
 };
+pub use entry::ApplicationData;
 use s2n_quic_core::{dc, time};
 use std::{net::SocketAddr, sync::Arc};
 
@@ -26,7 +27,7 @@ pub mod testing;
 #[cfg(test)]
 mod event_tests;
 
-use entry::Entry;
+pub use entry::Entry;
 use store::Store;
 
 pub use entry::{ApplicationPair, Bidirectional, ControlPair};
@@ -202,6 +203,7 @@ impl Map {
                 receiver::State::new(),
                 dc::testing::TEST_APPLICATION_PARAMS,
                 dc::testing::TEST_REHANDSHAKE_PERIOD,
+                Arc::new(()),
             );
             let entry = Arc::new(entry);
             provider.store.test_insert(entry);
@@ -255,6 +257,7 @@ impl Map {
                 super::receiver::State::new(),
                 dc::testing::TEST_APPLICATION_PARAMS,
                 dc::testing::TEST_REHANDSHAKE_PERIOD,
+                Arc::new(()),
             );
             let entry = Arc::new(entry);
             map.store.test_insert(entry);
@@ -269,4 +272,14 @@ impl Map {
 
         client_id
     }
+
+    #[allow(clippy::type_complexity)]
+    pub fn register_make_application_data(
+        &self,
+        cb: Box<
+            dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> ApplicationData + Send + Sync,
+        >,
+    ) {
+        self.store.register_make_application_data(cb);
+    }
 }
diff --git a/dc/s2n-quic-dc/src/path/secret/map/entry.rs b/dc/s2n-quic-dc/src/path/secret/map/entry.rs
@@ -20,6 +20,7 @@ use rand::Rng as _;
 use s2n_codec::EncoderBuffer;
 use s2n_quic_core::{dc, varint::VarInt};
 use std::{
+    any::Any,
     net::SocketAddr,
     sync::{
         atomic::{AtomicU32, AtomicU8, Ordering},
@@ -31,8 +32,10 @@ use std::{
 #[cfg(test)]
 mod tests;
 
+pub type ApplicationData = Arc<dyn Any + Send + Sync>;
+
 #[derive(Debug)]
-pub(super) struct Entry {
+pub struct Entry {
     creation_time: Instant,
     rehandshake_delta_secs: AtomicU32,
     peer: SocketAddr,
@@ -44,6 +47,7 @@ pub(super) struct Entry {
     // we store this as a u8 to allow the cleaner to separately "take" accessed for id and addr
     // maps while not having two writes and wasting an extra byte of space.
     accessed: AtomicU8,
+    application_data: ApplicationData,
 }
 
 impl SizeOf for Entry {
@@ -58,6 +62,7 @@ impl SizeOf for Entry {
             receiver,
             parameters,
             accessed,
+            application_data,
         } = self;
         creation_time.size()
             + rehandshake_delta_secs.size()
@@ -68,6 +73,13 @@ impl SizeOf for Entry {
             + receiver.size()
             + parameters.size()
             + accessed.size()
+            + application_data.size()
+    }
+}
+
+impl SizeOf for ApplicationData {
+    fn size(&self) -> usize {
+        std::mem::size_of_val(self)
     }
 }
 
@@ -82,6 +94,7 @@ impl Entry {
         receiver: receiver::State,
         parameters: dc::ApplicationParams,
         rehandshake_time: Duration,
+        application_data: ApplicationData,
     ) -> Self {
         // clamp max datagram size to a well-known value
         parameters
@@ -99,6 +112,7 @@ impl Entry {
             receiver,
             parameters,
             accessed: AtomicU8::new(0),
+            application_data,
         };
         entry.rehandshake_time_reschedule(rehandshake_time);
         entry
@@ -123,6 +137,7 @@ impl Entry {
             receiver,
             dc::testing::TEST_APPLICATION_PARAMS,
             dc::testing::TEST_REHANDSHAKE_PERIOD,
+            Arc::new(()),
         ))
     }
 
@@ -287,6 +302,10 @@ impl Entry {
     pub fn control_sealer(&self) -> crate::crypto::awslc::seal::control::Secret {
         self.secret.control_sealer()
     }
+
+    pub fn application_data(&self) -> &ApplicationData {
+        &self.application_data
+    }
 }
 
 impl receiver::Error {

diff --git a/dc/s2n-quic-dc/src/path/secret/map/entry/tests.rs b/dc/s2n-quic-dc/src/path/secret/map/entry/tests.rs
@@ -15,7 +15,7 @@ fn entry_size() {
     if should_check {
         assert_eq!(
             Entry::fake((std::net::Ipv4Addr::LOCALHOST, 0).into(), None).size(),
-            295
+            311
         );
     }
 }
diff --git a/dc/s2n-quic-dc/src/path/secret/map/handshake.rs b/dc/s2n-quic-dc/src/path/secret/map/handshake.rs
@@ -1,7 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use super::{Entry, Map};
+use super::{entry::ApplicationData, Entry, Map};
 use crate::{
     packet::secret_control as control,
     path::secret::{receiver, schedule, sender},
@@ -24,6 +24,7 @@ pub struct HandshakingPath {
     endpoint_type: s2n_quic_core::endpoint::Type,
     secret: Option<schedule::Secret>,
     entry: Option<Arc<Entry>>,
+    application_data: Option<ApplicationData>,
     map: Map,
 }
 
@@ -41,6 +42,7 @@ impl HandshakingPath {
             endpoint_type,
             secret: None,
             entry: None,
+            application_data: None,
             map,
         }
     }
@@ -82,6 +84,8 @@ impl dc::Path for HandshakingPath {
         &mut self,
         session: &impl s2n_quic_core::crypto::tls::TlsSession,
     ) -> Result<Vec<s2n_quic_core::stateless_reset::Token>, s2n_quic_core::transport::Error> {
+        self.application_data = Some(self.map.application_data.application_data(session));
+
         let mut material = Zeroizing::new([0; TLS_EXPORTER_LENGTH]);
         session
             .tls_exporter(
@@ -134,6 +138,7 @@ impl dc::Path for HandshakingPath {
             receiver,
             self.parameters.clone(),
             self.map.store.rehandshake_period(),
+            self.application_data.take().unwrap(),
         );
         let entry = Arc::new(entry);
         self.entry = Some(entry.clone());

diff --git a/dc/s2n-quic-dc/src/path/secret/map/state.rs b/dc/s2n-quic-dc/src/path/secret/map/state.rs
@@ -1,7 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use super::{cleaner::Cleaner, stateless_reset, Entry, Store};
+use super::{cleaner::Cleaner, entry::ApplicationData, stateless_reset, Entry, Store};
 use crate::{
     credentials::{Credentials, Id},
     crypto,
@@ -211,6 +211,19 @@ where
     clock: C,
 
     subscriber: S,
+
+    #[allow(clippy::type_complexity)]
+    mk_application_data: RwLock<
+        Option<
+            Box<
+                dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> ApplicationData
+                    + Send
+                    + Sync,
+            >,
+        >,
+    >,
+
+    dummy_application_data: ApplicationData,
 }
 
 // Share control sockets -- we only send on these so it doesn't really matter if there's only one
@@ -266,6 +279,8 @@ where
             clock,
             subscriber,
             request_handshake: RwLock::new(None),
+            mk_application_data: RwLock::new(None),
+            dummy_application_data: Arc::new(()),
         };
 
         // Growing to double our maximum inserted entries should ensure that we never grow again, see:
@@ -645,6 +660,20 @@ where
         self.register_request_handshake(cb);
     }
 
+    #[allow(clippy::type_complexity)]
+    fn register_make_application_data(
+        &self,
+        cb: Box<
+            dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> ApplicationData + Send + Sync,
+        >,
+    ) {
+        // FIXME: Maybe panic if already initialized?
+        *self
+            .mk_application_data
+            .write()
+            .unwrap_or_else(|e| e.into_inner()) = Some(cb);
+    }
+
     fn get_by_addr_untracked(&self, peer: &SocketAddr) -> Option<Arc<Entry>> {
         self.peers.get(*peer)
     }
@@ -845,6 +874,21 @@ where
     fn test_stop_cleaner(&self) {
         self.cleaner.stop();
     }
+
+    fn application_data(
+        &self,
+        session: &dyn s2n_quic_core::crypto::tls::TlsSession,
+    ) -> ApplicationData {
+        if let Some(ctxt) = &*self
+            .mk_application_data
+            .read()
+            .unwrap_or_else(|e| e.into_inner())
+        {
+            (ctxt)(session)
+        } else {
+            self.dummy_application_data.clone()
+        }
+    }
 }
 
 impl<C, S> Drop for State<C, S>

diff --git a/dc/s2n-quic-dc/src/path/secret/map/state/tests.rs b/dc/s2n-quic-dc/src/path/secret/map/state/tests.rs
@@ -139,6 +139,7 @@ impl Model {
                     receiver::State::new(),
                     dc::testing::TEST_APPLICATION_PARAMS,
                     dc::testing::TEST_REHANDSHAKE_PERIOD,
+                    Arc::new(()),
                 )));
 
                 self.invariants.insert(Invariant::ContainsIp(ip));

diff --git a/dc/s2n-quic-dc/src/path/secret/map/store.rs b/dc/s2n-quic-dc/src/path/secret/map/store.rs
@@ -1,7 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use super::Entry;
+use super::{entry::ApplicationData, Entry};
 use crate::{
     credentials::{Credentials, Id},
     packet::{secret_control as control, Packet, WireVersion},
@@ -101,4 +101,17 @@ pub trait Store: 'static + Send + Sync {
 
         Some(state.clone())
     }
+
+    #[allow(clippy::type_complexity)]
+    fn register_make_application_data(
+        &self,
+        cb: Box<
+            dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> ApplicationData + Send + Sync,
+        >,
+    );
+
+    fn application_data(
+        &self,
+        session: &dyn s2n_quic_core::crypto::tls::TlsSession,
+    ) -> ApplicationData;
 }