From b4cd37f42b277e3f8b87e80d4d7ab362703109fe Mon Sep 17 00:00:00 2001 From: Mark Rousskov Date: Thu, 6 Mar 2025 14:49:58 +0000 Subject: [PATCH] Support storing AuthContext in Entry This stores arbitrary metadata (Arc) in each Entry. By default, we share a single `Arc<()>` across all entries, but applications can register a callback to produce different contexts. --- dc/s2n-quic-dc/src/path/secret/map.rs | 13 ++++++- dc/s2n-quic-dc/src/path/secret/map/entry.rs | 21 ++++++++++- .../src/path/secret/map/handshake.rs | 7 +++- dc/s2n-quic-dc/src/path/secret/map/state.rs | 37 ++++++++++++++++++- .../src/path/secret/map/state/tests.rs | 1 + dc/s2n-quic-dc/src/path/secret/map/store.rs | 10 ++++- 6 files changed, 84 insertions(+), 5 deletions(-) diff --git a/dc/s2n-quic-dc/src/path/secret/map.rs b/dc/s2n-quic-dc/src/path/secret/map.rs index 55f782f3c..94a2899d3 100644 --- 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 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 index cb5ab70eb..722b9ce48 100644 --- 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; + #[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 index f0abf285f..f1876cf14 100644 --- 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, entry: Option>, + auth_context: Option, 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, 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 index 7acfec2b6..7dfe96baa 100644 --- 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 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 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> { 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 Drop for State 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 index aadc62001..e5f7ce70d 100644 --- 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 index c6243eedb..3fef8f08a 100644 --- 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 AuthContext + Send + Sync>, + ); + + fn auth_context(&self, session: &dyn s2n_quic_core::crypto::tls::TlsSession) -> AuthContext; }