Skip to content

Commit 8670e83

Browse files
Mark-SimulacrumMark-Simulacrum
authored
feat(s2n-quic-dc): Support storing ApplicationData in Entry (#2515)
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.
1 parent c82daf7 commit 8670e83

File tree

7 files changed

+101
-6
lines changed

7 files changed

+101
-6
lines changed

‎dc/s2n-quic-dc/src/path/secret/map.rs‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use crate::{
88
path::secret::{open, seal, stateless_reset},
99
stream::TransportFeatures,
1010
};
11+
pub use entry::ApplicationData;
1112
use s2n_quic_core::{dc, time};
1213
use std::{net::SocketAddr, sync::Arc};
1314

@@ -26,7 +27,7 @@ pub mod testing;
2627
#[cfg(test)]
2728
mod event_tests;
2829

29-
use entry::Entry;
30+
pub use entry::Entry;
3031
use store::Store;
3132

3233
pub use entry::{ApplicationPair, Bidirectional, ControlPair};
@@ -202,6 +203,7 @@ impl Map {
202203
receiver::State::new(),
203204
dc::testing::TEST_APPLICATION_PARAMS,
204205
dc::testing::TEST_REHANDSHAKE_PERIOD,
206+
Arc::new(()),
205207
);
206208
let entry = Arc::new(entry);
207209
provider.store.test_insert(entry);
@@ -255,6 +257,7 @@ impl Map {
255257
super::receiver::State::new(),
256258
dc::testing::TEST_APPLICATION_PARAMS,
257259
dc::testing::TEST_REHANDSHAKE_PERIOD,
260+
Arc::new(()),
258261
);
259262
let entry = Arc::new(entry);
260263
map.store.test_insert(entry);
@@ -269,4 +272,14 @@ impl Map {
269272

270273
client_id
271274
}
275+
276+
#[allow(clippy::type_complexity)]
277+
pub fn register_make_application_data(
278+
&self,
279+
cb: Box<
280+
dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> ApplicationData + Send + Sync,
281+
>,
282+
) {
283+
self.store.register_make_application_data(cb);
284+
}
272285
}

‎dc/s2n-quic-dc/src/path/secret/map/entry.rs‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ use rand::Rng as _;
2020
use s2n_codec::EncoderBuffer;
2121
use s2n_quic_core::{dc, varint::VarInt};
2222
use std::{
23+
any::Any,
2324
net::SocketAddr,
2425
sync::{
2526
atomic::{AtomicU32, AtomicU8, Ordering},
@@ -31,8 +32,10 @@ use std::{
3132
#[cfg(test)]
3233
mod tests;
3334

35+
pub type ApplicationData = Arc<dyn Any + Send + Sync>;
36+
3437
#[derive(Debug)]
35-
pub(super) struct Entry {
38+
pub struct Entry {
3639
creation_time: Instant,
3740
rehandshake_delta_secs: AtomicU32,
3841
peer: SocketAddr,
@@ -44,6 +47,7 @@ pub(super) struct Entry {
4447
// we store this as a u8 to allow the cleaner to separately "take" accessed for id and addr
4548
// maps while not having two writes and wasting an extra byte of space.
4649
accessed: AtomicU8,
50+
application_data: ApplicationData,
4751
}
4852

4953
impl SizeOf for Entry {
@@ -58,6 +62,7 @@ impl SizeOf for Entry {
5862
receiver,
5963
parameters,
6064
accessed,
65+
application_data,
6166
} = self;
6267
creation_time.size()
6368
+ rehandshake_delta_secs.size()
@@ -68,6 +73,13 @@ impl SizeOf for Entry {
6873
+ receiver.size()
6974
+ parameters.size()
7075
+ accessed.size()
76+
+ application_data.size()
77+
}
78+
}
79+
80+
impl SizeOf for ApplicationData {
81+
fn size(&self) -> usize {
82+
std::mem::size_of_val(self)
7183
}
7284
}
7385

@@ -82,6 +94,7 @@ impl Entry {
8294
receiver: receiver::State,
8395
parameters: dc::ApplicationParams,
8496
rehandshake_time: Duration,
97+
application_data: ApplicationData,
8598
) -> Self {
8699
// clamp max datagram size to a well-known value
87100
parameters
@@ -99,6 +112,7 @@ impl Entry {
99112
receiver,
100113
parameters,
101114
accessed: AtomicU8::new(0),
115+
application_data,
102116
};
103117
entry.rehandshake_time_reschedule(rehandshake_time);
104118
entry
@@ -123,6 +137,7 @@ impl Entry {
123137
receiver,
124138
dc::testing::TEST_APPLICATION_PARAMS,
125139
dc::testing::TEST_REHANDSHAKE_PERIOD,
140+
Arc::new(()),
126141
))
127142
}
128143

@@ -287,6 +302,10 @@ impl Entry {
287302
pub fn control_sealer(&self) -> crate::crypto::awslc::seal::control::Secret {
288303
self.secret.control_sealer()
289304
}
305+
306+
pub fn application_data(&self) -> &ApplicationData {
307+
&self.application_data
308+
}
290309
}
291310

292311
impl receiver::Error {

‎dc/s2n-quic-dc/src/path/secret/map/entry/tests.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ fn entry_size() {
1515
if should_check {
1616
assert_eq!(
1717
Entry::fake((std::net::Ipv4Addr::LOCALHOST, 0).into(), None).size(),
18-
295
18+
311
1919
);
2020
}
2121
}

‎dc/s2n-quic-dc/src/path/secret/map/handshake.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
22
// SPDX-License-Identifier: Apache-2.0
33

4-
use super::{Entry, Map};
4+
use super::{entry::ApplicationData, Entry, Map};
55
use crate::{
66
packet::secret_control as control,
77
path::secret::{receiver, schedule, sender},
@@ -24,6 +24,7 @@ pub struct HandshakingPath {
2424
endpoint_type: s2n_quic_core::endpoint::Type,
2525
secret: Option<schedule::Secret>,
2626
entry: Option<Arc<Entry>>,
27+
application_data: Option<ApplicationData>,
2728
map: Map,
2829
}
2930

@@ -41,6 +42,7 @@ impl HandshakingPath {
4142
endpoint_type,
4243
secret: None,
4344
entry: None,
45+
application_data: None,
4446
map,
4547
}
4648
}
@@ -82,6 +84,8 @@ impl dc::Path for HandshakingPath {
8284
&mut self,
8385
session: &impl s2n_quic_core::crypto::tls::TlsSession,
8486
) -> Result<Vec<s2n_quic_core::stateless_reset::Token>, s2n_quic_core::transport::Error> {
87+
self.application_data = Some(self.map.store.application_data(session));
88+
8589
let mut material = Zeroizing::new([0; TLS_EXPORTER_LENGTH]);
8690
session
8791
.tls_exporter(
@@ -134,6 +138,7 @@ impl dc::Path for HandshakingPath {
134138
receiver,
135139
self.parameters.clone(),
136140
self.map.store.rehandshake_period(),
141+
self.application_data.take().unwrap(),
137142
);
138143
let entry = Arc::new(entry);
139144
self.entry = Some(entry.clone());

‎dc/s2n-quic-dc/src/path/secret/map/state.rs‎

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
22
// SPDX-License-Identifier: Apache-2.0
33

4-
use super::{cleaner::Cleaner, stateless_reset, Entry, Store};
4+
use super::{cleaner::Cleaner, entry::ApplicationData, stateless_reset, Entry, Store};
55
use crate::{
66
credentials::{Credentials, Id},
77
crypto,
@@ -211,6 +211,19 @@ where
211211
clock: C,
212212

213213
subscriber: S,
214+
215+
#[allow(clippy::type_complexity)]
216+
mk_application_data: RwLock<
217+
Option<
218+
Box<
219+
dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> ApplicationData
220+
+ Send
221+
+ Sync,
222+
>,
223+
>,
224+
>,
225+
226+
dummy_application_data: ApplicationData,
214227
}
215228

216229
// Share control sockets -- we only send on these so it doesn't really matter if there's only one
@@ -266,6 +279,8 @@ where
266279
clock,
267280
subscriber,
268281
request_handshake: RwLock::new(None),
282+
mk_application_data: RwLock::new(None),
283+
dummy_application_data: Arc::new(()),
269284
};
270285

271286
// Growing to double our maximum inserted entries should ensure that we never grow again, see:
@@ -645,6 +660,20 @@ where
645660
self.register_request_handshake(cb);
646661
}
647662

663+
#[allow(clippy::type_complexity)]
664+
fn register_make_application_data(
665+
&self,
666+
cb: Box<
667+
dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> ApplicationData + Send + Sync,
668+
>,
669+
) {
670+
// FIXME: Maybe panic if already initialized?
671+
*self
672+
.mk_application_data
673+
.write()
674+
.unwrap_or_else(|e| e.into_inner()) = Some(cb);
675+
}
676+
648677
fn get_by_addr_untracked(&self, peer: &SocketAddr) -> Option<Arc<Entry>> {
649678
self.peers.get(*peer)
650679
}
@@ -845,6 +874,21 @@ where
845874
fn test_stop_cleaner(&self) {
846875
self.cleaner.stop();
847876
}
877+
878+
fn application_data(
879+
&self,
880+
session: &dyn s2n_quic_core::crypto::tls::TlsSession,
881+
) -> ApplicationData {
882+
if let Some(ctxt) = &*self
883+
.mk_application_data
884+
.read()
885+
.unwrap_or_else(|e| e.into_inner())
886+
{
887+
(ctxt)(session)
888+
} else {
889+
self.dummy_application_data.clone()
890+
}
891+
}
848892
}
849893

850894
impl<C, S> Drop for State<C, S>

‎dc/s2n-quic-dc/src/path/secret/map/state/tests.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,7 @@ impl Model {
139139
receiver::State::new(),
140140
dc::testing::TEST_APPLICATION_PARAMS,
141141
dc::testing::TEST_REHANDSHAKE_PERIOD,
142+
Arc::new(()),
142143
)));
143144

144145
self.invariants.insert(Invariant::ContainsIp(ip));

‎dc/s2n-quic-dc/src/path/secret/map/store.rs‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
22
// SPDX-License-Identifier: Apache-2.0
33

4-
use super::Entry;
4+
use super::{entry::ApplicationData, Entry};
55
use crate::{
66
credentials::{Credentials, Id},
77
packet::{secret_control as control, Packet, WireVersion},
@@ -101,4 +101,17 @@ pub trait Store: 'static + Send + Sync {
101101

102102
Some(state.clone())
103103
}
104+
105+
#[allow(clippy::type_complexity)]
106+
fn register_make_application_data(
107+
&self,
108+
cb: Box<
109+
dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> ApplicationData + Send + Sync,
110+
>,
111+
);
112+
113+
fn application_data(
114+
&self,
115+
session: &dyn s2n_quic_core::crypto::tls::TlsSession,
116+
) -> ApplicationData;
104117
}

0 commit comments

Comments
 (0)