Support storing AuthContext 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 different contexts.
aws · Mar 6, 2025 · b4cd37f · b4cd37f
1 parent a9e7673
commit b4cd37f
Show file tree

Hide file tree

Showing 6 changed files with 84 additions and 5 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::AuthContext;
 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,12 @@ impl Map {
 
         client_id
     }
+
+    #[allow(clippy::type_complexity)]
+    pub fn register_make_auth_context(
+        &self,
+        cb: Box<dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> AuthContext + Send + Sync>,
+    ) {
+        self.store.register_make_auth_context(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 AuthContext = 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,
+    auth_context: AuthContext,
 }
 
 impl SizeOf for Entry {
@@ -58,6 +62,7 @@ impl SizeOf for Entry {
             receiver,
             parameters,
             accessed,
+            auth_context,
         } = self;
         creation_time.size()
             + rehandshake_delta_secs.size()
@@ -68,6 +73,13 @@ impl SizeOf for Entry {
             + receiver.size()
             + parameters.size()
             + accessed.size()
+            + auth_context.size()
+    }
+}
+
+impl SizeOf for AuthContext {
+    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,
+        auth_context: AuthContext,
     ) -> Self {
         // clamp max datagram size to a well-known value
         parameters
@@ -99,6 +112,7 @@ impl Entry {
             receiver,
             parameters,
             accessed: AtomicU8::new(0),
+            auth_context,
         };
         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 auth_context(&self) -> &AuthContext {
+        &self.auth_context
+    }
 }
 
 impl receiver::Error {

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::AuthContext, 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>>,
+    auth_context: Option<AuthContext>,
     map: Map,
 }
 
@@ -41,6 +42,7 @@ impl HandshakingPath {
             endpoint_type,
             secret: None,
             entry: None,
+            auth_context: 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.auth_context = Some(self.map.store.auth_context(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.auth_context.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::AuthContext, stateless_reset, Entry, Store};
 use crate::{
     credentials::{Credentials, Id},
     crypto,
@@ -211,6 +211,15 @@ where
     clock: C,
 
     subscriber: S,
+
+    #[allow(clippy::type_complexity)]
+    mk_auth_context: RwLock<
+        Option<
+            Box<dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> AuthContext + Send + Sync>,
+        >,
+    >,
+
+    dummy_auth_context: AuthContext,
 }
 
 // Share control sockets -- we only send on these so it doesn't really matter if there's only one
@@ -266,6 +275,8 @@ where
             clock,
             subscriber,
             request_handshake: RwLock::new(None),
+            mk_auth_context: RwLock::new(None),
+            dummy_auth_context: Arc::new(()),
         };
 
         // Growing to double our maximum inserted entries should ensure that we never grow again, see:
@@ -645,6 +656,18 @@ where
         self.register_request_handshake(cb);
     }
 
+    #[allow(clippy::type_complexity)]
+    fn register_make_auth_context(
+        &self,
+        cb: Box<dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> AuthContext + Send + Sync>,
+    ) {
+        // FIXME: Maybe panic if already initialized?
+        *self
+            .mk_auth_context
+            .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 +868,18 @@ where
     fn test_stop_cleaner(&self) {
         self.cleaner.stop();
     }
+
+    fn auth_context(&self, session: &dyn s2n_quic_core::crypto::tls::TlsSession) -> AuthContext {
+        if let Some(ctxt) = &*self
+            .mk_auth_context
+            .read()
+            .unwrap_or_else(|e| e.into_inner())
+        {
+            (ctxt)(session)
+        } else {
+            self.dummy_auth_context.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::AuthContext, Entry};
 use crate::{
     credentials::{Credentials, Id},
     packet::{secret_control as control, Packet, WireVersion},
@@ -101,4 +101,12 @@ pub trait Store: 'static + Send + Sync {
 
         Some(state.clone())
     }
+
+    #[allow(clippy::type_complexity)]
+    fn register_make_auth_context(
+        &self,
+        cb: Box<dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> AuthContext + Send + Sync>,
+    );
+
+    fn auth_context(&self, session: &dyn s2n_quic_core::crypto::tls::TlsSession) -> AuthContext;
 }