Skip to content

Commit 55cd809

Browse files
maddeleinecamshaft
maddeleine
authored and
camshaft
committed
feat(s2n-quic): implement updatable connection limits (#2508)
1 parent f21fb99 commit 55cd809

File tree

12 files changed

+182
-9
lines changed

12 files changed

+182
-9
lines changed

‎quic/s2n-quic-core/src/connection/limits.rs

+64-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
22
// SPDX-License-Identifier: Apache-2.0
3-
3+
#[cfg(feature = "alloc")]
4+
use crate::application::ServerName;
45
use crate::{
56
ack,
67
event::{api::SocketAddress, IntoEvent},
@@ -12,6 +13,8 @@ use crate::{
1213
MaxDatagramFrameSize, MaxIdleTimeout, MigrationSupport, TransportParameters,
1314
},
1415
};
16+
#[cfg(feature = "alloc")]
17+
use bytes::Bytes;
1518
use core::time::Duration;
1619
use s2n_codec::decoder_invariant;
1720

@@ -53,6 +56,30 @@ impl<'a> ConnectionInfo<'a> {
5356
}
5457
}
5558

59+
#[non_exhaustive]
60+
#[derive(Debug)]
61+
#[cfg(feature = "alloc")]
62+
pub struct HandshakeInfo<'a> {
63+
pub remote_address: SocketAddress<'a>,
64+
pub server_name: Option<&'a ServerName>,
65+
pub application_protocol: &'a Bytes,
66+
}
67+
68+
#[cfg(feature = "alloc")]
69+
impl<'a> HandshakeInfo<'a> {
70+
pub fn new(
71+
remote_address: &'a inet::SocketAddress,
72+
server_name: Option<&'a ServerName>,
73+
application_protocol: &'a Bytes,
74+
) -> HandshakeInfo<'a> {
75+
Self {
76+
remote_address: remote_address.into_event(),
77+
server_name,
78+
application_protocol,
79+
}
80+
}
81+
}
82+
5683
#[derive(Clone, Copy, Debug)]
5784
pub struct Limits {
5885
pub(crate) max_idle_timeout: MaxIdleTimeout,
@@ -397,16 +424,41 @@ impl Limits {
397424
}
398425
}
399426

427+
#[must_use]
428+
#[derive(Debug)]
429+
pub struct UpdatableLimits<'a>(&'a mut Limits);
430+
431+
impl<'a> UpdatableLimits<'a> {
432+
pub fn new(limits: &'a mut Limits) -> UpdatableLimits<'a> {
433+
UpdatableLimits(limits)
434+
}
435+
436+
pub fn with_stream_batch_size(&mut self, size: u8) {
437+
self.0.stream_batch_size = size;
438+
}
439+
}
440+
400441
/// Creates limits for a given connection
401442
pub trait Limiter: 'static + Send {
402443
fn on_connection(&mut self, info: &ConnectionInfo) -> Limits;
444+
445+
/// Provides another opportunity to change connection limits with information
446+
/// from the handshake
447+
#[inline]
448+
#[cfg(feature = "alloc")]
449+
fn on_post_handshake(&mut self, info: &HandshakeInfo, limits: &mut UpdatableLimits) {
450+
let _ = info;
451+
let _ = limits;
452+
}
403453
}
404454

405455
/// Implement Limiter for a Limits struct
406456
impl Limiter for Limits {
407457
fn on_connection(&mut self, _into: &ConnectionInfo) -> Limits {
408458
*self
409459
}
460+
#[cfg(feature = "alloc")]
461+
fn on_post_handshake(&mut self, _info: &HandshakeInfo, _limits: &mut UpdatableLimits) {}
410462
}
411463

412464
#[cfg(test)]
@@ -429,4 +481,15 @@ mod tests {
429481
assert!(limits.with_bidirectional_remote_data_window(data).is_ok());
430482
assert!(limits.with_unidirectional_data_window(data).is_ok());
431483
}
484+
485+
// Limits can be updated through the UpdatableLimits wrapper
486+
#[test]
487+
fn updatable_limits() {
488+
let mut limits = Limits::default();
489+
assert_eq!(limits.stream_batch_size, 1);
490+
let mut updatable_limits = UpdatableLimits::new(&mut limits);
491+
let new_size = 10;
492+
updatable_limits.with_stream_batch_size(new_size);
493+
assert_eq!(limits.stream_batch_size, new_size);
494+
}
432495
}

‎quic/s2n-quic-transport/src/connection/connection_container/tests.rs

+5
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@ impl connection::Trait for TestConnection {
134134
_subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
135135
_datagram: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
136136
_dc_endpoint: &mut <Self::Config as endpoint::Config>::DcEndpoint,
137+
_conn_limits_endpoint: &mut <Self::Config as endpoint::Config>::ConnectionLimits,
137138
) -> Result<(), connection::Error> {
138139
Ok(())
139140
}
@@ -148,6 +149,7 @@ impl connection::Trait for TestConnection {
148149
_packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
149150
_datagram_endpoint: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
150151
_dc_endpoint: &mut <Self::Config as endpoint::Config>::DcEndpoint,
152+
_conn_limits_endpoint: &mut <Self::Config as endpoint::Config>::ConnectionLimits,
151153
) -> Result<(), ProcessingError> {
152154
Ok(())
153155
}
@@ -163,6 +165,7 @@ impl connection::Trait for TestConnection {
163165
_packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
164166
_datagram_endpoint: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
165167
_dc_endpoint: &mut <Self::Config as endpoint::Config>::DcEndpoint,
168+
_conn_limits_endpoint: &mut <Self::Config as endpoint::Config>::ConnectionLimits,
166169
) -> Result<(), ProcessingError> {
167170
Ok(())
168171
}
@@ -178,6 +181,7 @@ impl connection::Trait for TestConnection {
178181
_packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
179182
_datagram_endpoint: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
180183
_dc_endpoint: &mut <Self::Config as endpoint::Config>::DcEndpoint,
184+
_connection_limits_endpoint: &mut <Self::Config as endpoint::Config>::ConnectionLimits,
181185
) -> Result<(), ProcessingError> {
182186
Ok(())
183187
}
@@ -193,6 +197,7 @@ impl connection::Trait for TestConnection {
193197
_packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
194198
_datagram_endpoint: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
195199
_dc_endpoint: &mut <Self::Config as endpoint::Config>::DcEndpoint,
200+
_limits_endpoint: &mut <Self::Config as endpoint::Config>::ConnectionLimits,
196201
) -> Result<(), ProcessingError> {
197202
Ok(())
198203
}

‎quic/s2n-quic-transport/src/connection/connection_impl.rs

+13-1
Original file line numberDiff line numberDiff line change
@@ -272,6 +272,7 @@ impl<Config: endpoint::Config> ConnectionImpl<Config> {
272272
subscriber: &mut Config::EventSubscriber,
273273
datagram: &mut Config::DatagramEndpoint,
274274
dc: &mut Config::DcEndpoint,
275+
limits: &mut Config::ConnectionLimits,
275276
) -> Result<(), connection::Error> {
276277
let mut publisher = self.event_context.publisher(timestamp, subscriber);
277278
let space_manager = &mut self.space_manager;
@@ -285,6 +286,7 @@ impl<Config: endpoint::Config> ConnectionImpl<Config> {
285286
&mut publisher,
286287
datagram,
287288
dc,
289+
limits,
288290
) {
289291
Poll::Ready(Ok(())) => {}
290292
// use `from` instead of `into` so the location is correctly captured
@@ -658,6 +660,7 @@ impl<Config: endpoint::Config> connection::Trait for ConnectionImpl<Config> {
658660
parameters.event_subscriber,
659661
parameters.datagram_endpoint,
660662
parameters.dc_endpoint,
663+
parameters.limits_endpoint,
661664
) {
662665
connection.with_event_publisher(
663666
parameters.timestamp,
@@ -1128,12 +1131,13 @@ impl<Config: endpoint::Config> connection::Trait for ConnectionImpl<Config> {
11281131
subscriber: &mut Config::EventSubscriber,
11291132
datagram: &mut Config::DatagramEndpoint,
11301133
dc: &mut Config::DcEndpoint,
1134+
conn_limits: &mut Config::ConnectionLimits,
11311135
) -> Result<(), connection::Error> {
11321136
// reset the queued state first so that new wakeup request are not missed
11331137
self.wakeup_handle.wakeup_handled();
11341138

11351139
// check if crypto progress can be made
1136-
self.update_crypto_state(timestamp, subscriber, datagram, dc)?;
1140+
self.update_crypto_state(timestamp, subscriber, datagram, dc, conn_limits)?;
11371141

11381142
// return an error if the application set one
11391143
self.error?;
@@ -1220,6 +1224,7 @@ impl<Config: endpoint::Config> connection::Trait for ConnectionImpl<Config> {
12201224
packet_interceptor: &mut Config::PacketInterceptor,
12211225
datagram_endpoint: &mut Config::DatagramEndpoint,
12221226
dc_endpoint: &mut Config::DcEndpoint,
1227+
connection_limits_endpoint: &mut Config::ConnectionLimits,
12231228
) -> Result<(), ProcessingError> {
12241229
//= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
12251230
//= type=TODO
@@ -1261,6 +1266,7 @@ impl<Config: endpoint::Config> connection::Trait for ConnectionImpl<Config> {
12611266
packet_interceptor,
12621267
datagram_endpoint,
12631268
dc_endpoint,
1269+
connection_limits_endpoint,
12641270
)?;
12651271
}
12661272

@@ -1278,6 +1284,7 @@ impl<Config: endpoint::Config> connection::Trait for ConnectionImpl<Config> {
12781284
packet_interceptor: &mut Config::PacketInterceptor,
12791285
datagram_endpoint: &mut Config::DatagramEndpoint,
12801286
dc_endpoint: &mut Config::DcEndpoint,
1287+
connection_limits_endpoint: &mut Config::ConnectionLimits,
12811288
) -> Result<(), ProcessingError> {
12821289
if let Some((space, handshake_status)) = self.space_manager.initial_mut() {
12831290
let mut publisher = self.event_context.publisher(datagram.timestamp, subscriber);
@@ -1339,6 +1346,7 @@ impl<Config: endpoint::Config> connection::Trait for ConnectionImpl<Config> {
13391346
subscriber,
13401347
datagram_endpoint,
13411348
dc_endpoint,
1349+
connection_limits_endpoint,
13421350
)?;
13431351

13441352
// notify the connection a packet was processed
@@ -1359,6 +1367,7 @@ impl<Config: endpoint::Config> connection::Trait for ConnectionImpl<Config> {
13591367
packet_interceptor: &mut Config::PacketInterceptor,
13601368
datagram_endpoint: &mut Config::DatagramEndpoint,
13611369
dc_endpoint: &mut Config::DcEndpoint,
1370+
connection_limits_endpoint: &mut Config::ConnectionLimits,
13621371
) -> Result<(), ProcessingError> {
13631372
let mut publisher = self.event_context.publisher(datagram.timestamp, subscriber);
13641373

@@ -1444,6 +1453,7 @@ impl<Config: endpoint::Config> connection::Trait for ConnectionImpl<Config> {
14441453
subscriber,
14451454
datagram_endpoint,
14461455
dc_endpoint,
1456+
connection_limits_endpoint,
14471457
)?;
14481458

14491459
// notify the connection a packet was processed
@@ -1464,6 +1474,7 @@ impl<Config: endpoint::Config> connection::Trait for ConnectionImpl<Config> {
14641474
packet_interceptor: &mut Config::PacketInterceptor,
14651475
datagram_endpoint: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
14661476
dc_endpoint: &mut Config::DcEndpoint,
1477+
limits_endpoint: &mut Config::ConnectionLimits,
14671478
) -> Result<(), ProcessingError> {
14681479
let mut publisher = self.event_context.publisher(datagram.timestamp, subscriber);
14691480

@@ -1561,6 +1572,7 @@ impl<Config: endpoint::Config> connection::Trait for ConnectionImpl<Config> {
15611572
&mut publisher,
15621573
datagram_endpoint,
15631574
dc_endpoint,
1575+
limits_endpoint,
15641576
)?;
15651577
}
15661578
// notify the connection a packet was processed

‎quic/s2n-quic-transport/src/connection/connection_trait.rs

+11
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,7 @@ pub trait ConnectionTrait: 'static + Send + Sized {
111111
subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
112112
datagram: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
113113
dc_endpoint: &mut <Self::Config as endpoint::Config>::DcEndpoint,
114+
conn_limits: &mut <Self::Config as endpoint::Config>::ConnectionLimits,
114115
) -> Result<(), connection::Error>;
115116

116117
// Packet handling
@@ -126,6 +127,7 @@ pub trait ConnectionTrait: 'static + Send + Sized {
126127
packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
127128
datagram_endpoint: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
128129
dc_endpoint: &mut <Self::Config as endpoint::Config>::DcEndpoint,
130+
connection_limits_endpoint: &mut <Self::Config as endpoint::Config>::ConnectionLimits,
129131
) -> Result<(), ProcessingError>;
130132

131133
/// Is called when an unprotected initial packet had been received
@@ -139,6 +141,7 @@ pub trait ConnectionTrait: 'static + Send + Sized {
139141
packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
140142
datagram_endpoint: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
141143
dc_endpoint: &mut <Self::Config as endpoint::Config>::DcEndpoint,
144+
connection_limits_endpoint: &mut <Self::Config as endpoint::Config>::ConnectionLimits,
142145
) -> Result<(), ProcessingError>;
143146

144147
/// Is called when a handshake packet had been received
@@ -152,6 +155,7 @@ pub trait ConnectionTrait: 'static + Send + Sized {
152155
packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
153156
datagram_endpoint: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
154157
dc_endpoint: &mut <Self::Config as endpoint::Config>::DcEndpoint,
158+
connection_limits_endpoint: &mut <Self::Config as endpoint::Config>::ConnectionLimits,
155159
) -> Result<(), ProcessingError>;
156160

157161
/// Is called when a short packet had been received
@@ -165,6 +169,7 @@ pub trait ConnectionTrait: 'static + Send + Sized {
165169
packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
166170
datagram_endpoint: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
167171
dc_endpoint: &mut <Self::Config as endpoint::Config>::DcEndpoint,
172+
limits_endpoint: &mut <Self::Config as endpoint::Config>::ConnectionLimits,
168173
) -> Result<(), ProcessingError>;
169174

170175
/// Is called when a version negotiation packet had been received
@@ -225,6 +230,7 @@ pub trait ConnectionTrait: 'static + Send + Sized {
225230
packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
226231
datagram_endpoint: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
227232
dc_endpoint: &mut <Self::Config as endpoint::Config>::DcEndpoint,
233+
connection_limits_endpoint: &mut <Self::Config as endpoint::Config>::ConnectionLimits,
228234
check_for_stateless_reset: &mut bool,
229235
) -> Result<(), connection::Error> {
230236
macro_rules! emit_drop_reason {
@@ -286,6 +292,7 @@ pub trait ConnectionTrait: 'static + Send + Sized {
286292
packet_interceptor,
287293
datagram_endpoint,
288294
dc_endpoint,
295+
connection_limits_endpoint,
289296
),
290297
ProtectedPacket::VersionNegotiation(packet) => self.handle_version_negotiation_packet(
291298
datagram,
@@ -303,6 +310,7 @@ pub trait ConnectionTrait: 'static + Send + Sized {
303310
packet_interceptor,
304311
datagram_endpoint,
305312
dc_endpoint,
313+
connection_limits_endpoint,
306314
),
307315
ProtectedPacket::ZeroRtt(packet) => self.handle_zero_rtt_packet(
308316
datagram,
@@ -320,6 +328,7 @@ pub trait ConnectionTrait: 'static + Send + Sized {
320328
packet_interceptor,
321329
datagram_endpoint,
322330
dc_endpoint,
331+
connection_limits_endpoint,
323332
),
324333
ProtectedPacket::Retry(packet) => {
325334
self.handle_retry_packet(datagram, path_id, packet, subscriber, packet_interceptor)
@@ -378,6 +387,7 @@ pub trait ConnectionTrait: 'static + Send + Sized {
378387
packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
379388
datagram_endpoint: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
380389
dc_endpoint: &mut <Self::Config as endpoint::Config>::DcEndpoint,
390+
connection_limits_endpoint: &mut <Self::Config as endpoint::Config>::ConnectionLimits,
381391
check_for_stateless_reset: &mut bool,
382392
) -> Result<(), connection::Error> {
383393
macro_rules! emit_drop_reason {
@@ -432,6 +442,7 @@ pub trait ConnectionTrait: 'static + Send + Sized {
432442
packet_interceptor,
433443
datagram_endpoint,
434444
dc_endpoint,
445+
connection_limits_endpoint,
435446
check_for_stateless_reset,
436447
);
437448

‎quic/s2n-quic-transport/src/connection/mod.rs

+2
Original file line numberDiff line numberDiff line change
@@ -87,4 +87,6 @@ pub struct Parameters<'a, Cfg: endpoint::Config> {
8787
pub dc_endpoint: &'a mut Cfg::DcEndpoint,
8888
/// The event subscriber for the endpoint
8989
pub event_subscriber: &'a mut Cfg::EventSubscriber,
90+
/// The connection limits provider
91+
pub limits_endpoint: &'a mut Cfg::ConnectionLimits,
9092
}

‎quic/s2n-quic-transport/src/endpoint/initial.rs

+3
Original file line numberDiff line numberDiff line change
@@ -315,6 +315,7 @@ impl<Config: endpoint::Config> endpoint::Endpoint<Config> {
315315
datagram_endpoint: endpoint_context.datagram,
316316
dc_endpoint: endpoint_context.dc,
317317
open_registry: None,
318+
limits_endpoint: endpoint_context.connection_limits,
318319
};
319320

320321
let mut connection = <Config as endpoint::Config>::Connection::new(connection_parameters)?;
@@ -367,6 +368,7 @@ impl<Config: endpoint::Config> endpoint::Endpoint<Config> {
367368
endpoint_context.packet_interceptor,
368369
endpoint_context.datagram,
369370
endpoint_context.dc,
371+
endpoint_context.connection_limits,
370372
)
371373
.map_err(|err| {
372374
use connection::ProcessingError;
@@ -391,6 +393,7 @@ impl<Config: endpoint::Config> endpoint::Endpoint<Config> {
391393
endpoint_context.packet_interceptor,
392394
endpoint_context.datagram,
393395
endpoint_context.dc,
396+
endpoint_context.connection_limits,
394397
&mut false,
395398
)?;
396399

0 commit comments

Comments
 (0)