Skip to content

Commit c64faf9

Browse files
camshaftcamshaft
authored
fix(s2n-quic-dc): route secret control packets for UDP streams (#2747)
1 parent 9fe81de commit c64faf9

File tree

26 files changed

+856
-271
lines changed

26 files changed

+856
-271
lines changed

‎dc/s2n-quic-dc/src/packet/secret_control.rs‎

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,11 @@ const UNKNOWN_PATH_SECRET: u8 = 0b0110_0000;
1818
const STALE_KEY: u8 = 0b0110_0001;
1919
const REPLAY_DETECTED: u8 = 0b0110_0010;
2020

21+
/// Indicates if the packet has a queue_id field
22+
///
23+
/// This is used to route packets to the appropriate sender.
24+
const HAS_QUEUE_ID: u8 = 0b0000_0100;
25+
2126
pub const MAX_PACKET_SIZE: usize = 64;
2227
pub const TAG_LEN: usize = 16;
2328

@@ -35,6 +40,19 @@ macro_rules! impl_tag {
3540

3641
impl Tag {
3742
pub const VALUE: u8 = $tag;
43+
pub const VALUE_WITH_QUEUE_ID: u8 = $tag | HAS_QUEUE_ID;
44+
45+
pub const fn has_queue_id(&self) -> bool {
46+
(self.0 & HAS_QUEUE_ID) != 0
47+
}
48+
49+
pub fn with_queue_id(self, has_queue_id: bool) -> Self {
50+
if has_queue_id {
51+
Self(self.0 | HAS_QUEUE_ID)
52+
} else {
53+
self
54+
}
55+
}
3856
}
3957

4058
impl From<Tag> for u8 {
@@ -48,7 +66,7 @@ macro_rules! impl_tag {
4866
impl<'a> Tag {
4967
fn decode(buffer: Buffer) -> Result<Self> {
5068
let (tag, buffer) = buffer.decode()?;
51-
decoder_invariant!(tag == $tag, "invalid tag");
69+
decoder_invariant!([$tag, $tag | HAS_QUEUE_ID].contains(&tag), "invalid tag");
5270
Ok((Self(tag), buffer))
5371
}
5472
}
@@ -135,8 +153,9 @@ impl<'a> Packet<'a> {
135153
#[inline]
136154
pub fn decode(buffer: DecoderBufferMut<'a>) -> Rm<'a, Self> {
137155
let tag = buffer.peek_byte(0)?;
156+
let base_tag = tag & !HAS_QUEUE_ID;
138157

139-
Ok(match tag {
158+
Ok(match base_tag {
140159
UNKNOWN_PATH_SECRET => {
141160
let (packet, buffer) = unknown_path_secret::Packet::decode(buffer)?;
142161
(Self::UnknownPathSecret(packet), buffer)
@@ -161,6 +180,15 @@ impl<'a> Packet<'a> {
161180
Self::ReplayDetected(p) => p.credential_id(),
162181
}
163182
}
183+
184+
#[inline]
185+
pub fn queue_id(&self) -> Option<VarInt> {
186+
match self {
187+
Self::UnknownPathSecret(p) => p.queue_id(),
188+
Self::StaleKey(p) => p.queue_id(),
189+
Self::ReplayDetected(p) => p.queue_id(),
190+
}
191+
}
164192
}
165193

166194
macro_rules! impl_convert {

‎dc/s2n-quic-dc/src/packet/secret_control/decoder.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,11 @@ macro_rules! impl_packet {
4747

4848
Some(value)
4949
}
50+
51+
#[inline]
52+
pub fn queue_id(&self) -> Option<s2n_quic_core::varint::VarInt> {
53+
self.value.queue_id
54+
}
5055
}
5156
};
5257
}

‎dc/s2n-quic-dc/src/packet/secret_control/replay_detected.rs‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,9 @@ impl_packet!(ReplayDetected);
99
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1010
#[cfg_attr(test, derive(bolero_generator::TypeGenerator))]
1111
pub struct ReplayDetected {
12-
pub wire_version: WireVersion,
1312
pub credential_id: credentials::Id,
13+
pub wire_version: WireVersion,
14+
pub queue_id: Option<VarInt>,
1415
pub rejected_key_id: VarInt,
1516
}
1617

@@ -20,9 +21,12 @@ impl ReplayDetected {
2021
where
2122
C: seal::control::Secret,
2223
{
23-
encoder.encode(&Tag::default());
24+
encoder.encode(&Tag::default().with_queue_id(self.queue_id.is_some()));
2425
encoder.encode(&self.credential_id);
2526
encoder.encode(&self.wire_version);
27+
if let Some(queue_id) = self.queue_id {
28+
encoder.encode(&queue_id);
29+
}
2630
encoder.encode(&self.rejected_key_id);
2731

2832
encoder::finish(encoder, crypto)
@@ -38,13 +42,19 @@ impl<'a> DecoderValue<'a> for ReplayDetected {
3842
#[inline]
3943
fn decode(buffer: DecoderBuffer<'a>) -> R<'a, Self> {
4044
let (tag, buffer) = buffer.decode::<Tag>()?;
41-
decoder_invariant!(tag == Tag::default(), "invalid tag");
4245
let (credential_id, buffer) = buffer.decode()?;
4346
let (wire_version, buffer) = buffer.decode()?;
47+
let (queue_id, buffer) = if tag.has_queue_id() {
48+
let (queue_id, buffer) = buffer.decode()?;
49+
(Some(queue_id), buffer)
50+
} else {
51+
(None, buffer)
52+
};
4453
let (rejected_key_id, buffer) = buffer.decode()?;
4554
let value = Self {
46-
wire_version,
4755
credential_id,
56+
wire_version,
57+
queue_id,
4858
rejected_key_id,
4959
};
5060
Ok((value, buffer))

‎dc/s2n-quic-dc/src/packet/secret_control/stale_key.rs‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,9 @@ impl_packet!(StaleKey);
99
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1010
#[cfg_attr(test, derive(bolero_generator::TypeGenerator))]
1111
pub struct StaleKey {
12-
pub wire_version: WireVersion,
1312
pub credential_id: credentials::Id,
13+
pub wire_version: WireVersion,
14+
pub queue_id: Option<VarInt>,
1415
pub min_key_id: VarInt,
1516
}
1617

@@ -20,9 +21,12 @@ impl StaleKey {
2021
where
2122
C: seal::control::Secret,
2223
{
23-
encoder.encode(&Tag::default());
24+
encoder.encode(&Tag::default().with_queue_id(self.queue_id.is_some()));
2425
encoder.encode(&self.credential_id);
2526
encoder.encode(&self.wire_version);
27+
if let Some(queue_id) = self.queue_id {
28+
encoder.encode(&queue_id);
29+
}
2630
encoder.encode(&self.min_key_id);
2731

2832
encoder::finish(encoder, crypto)
@@ -38,13 +42,19 @@ impl<'a> DecoderValue<'a> for StaleKey {
3842
#[inline]
3943
fn decode(buffer: DecoderBuffer<'a>) -> R<'a, Self> {
4044
let (tag, buffer) = buffer.decode::<Tag>()?;
41-
decoder_invariant!(tag == Tag::default(), "invalid tag");
4245
let (credential_id, buffer) = buffer.decode()?;
4346
let (wire_version, buffer) = buffer.decode()?;
47+
let (queue_id, buffer) = if tag.has_queue_id() {
48+
let (queue_id, buffer) = buffer.decode()?;
49+
(Some(queue_id), buffer)
50+
} else {
51+
(None, buffer)
52+
};
4453
let (min_key_id, buffer) = buffer.decode()?;
4554
let value = Self {
4655
wire_version,
4756
credential_id,
57+
queue_id,
4858
min_key_id,
4959
};
5060
Ok((value, buffer))

‎dc/s2n-quic-dc/src/packet/secret_control/unknown_path_secret.rs‎

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ impl<'a> Packet<'a> {
2121
value: UnknownPathSecret {
2222
wire_version: WireVersion::ZERO,
2323
credential_id: id,
24+
queue_id: None,
2425
},
2526
crypto_tag: &stateless_reset[..],
2627
}
@@ -48,25 +49,37 @@ impl<'a> Packet<'a> {
4849
aws_lc_rs::constant_time::verify_slices_are_equal(self.crypto_tag, stateless_reset).ok()?;
4950
Some(&self.value)
5051
}
52+
53+
#[inline]
54+
pub fn queue_id(&self) -> Option<VarInt> {
55+
self.value.queue_id
56+
}
5157
}
5258

5359
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
5460
#[cfg_attr(test, derive(bolero_generator::TypeGenerator))]
5561
pub struct UnknownPathSecret {
56-
pub wire_version: WireVersion,
5762
pub credential_id: credentials::Id,
63+
pub wire_version: WireVersion,
64+
pub queue_id: Option<VarInt>,
5865
}
5966

6067
impl UnknownPathSecret {
61-
pub const PACKET_SIZE: usize =
62-
size_of::<Tag>() + size_of::<u8>() + size_of::<credentials::Id>() + TAG_LEN;
68+
pub const MAX_PACKET_SIZE: usize = size_of::<Tag>()
69+
+ size_of::<u8>()
70+
+ size_of::<credentials::Id>()
71+
+ size_of::<VarInt>()
72+
+ TAG_LEN;
6373

6474
#[inline]
6575
pub fn encode(&self, mut encoder: EncoderBuffer, stateless_reset_tag: &[u8; TAG_LEN]) -> usize {
6676
let before = encoder.len();
67-
encoder.encode(&Tag::default());
77+
encoder.encode(&Tag::default().with_queue_id(self.queue_id.is_some()));
6878
encoder.encode(&&self.credential_id[..]);
6979
encoder.encode(&self.wire_version);
80+
if let Some(queue_id) = self.queue_id {
81+
encoder.encode(&queue_id);
82+
}
7083
encoder.encode(&&stateless_reset_tag[..]);
7184
let after = encoder.len();
7285
after - before
@@ -77,12 +90,19 @@ impl<'a> DecoderValue<'a> for UnknownPathSecret {
7790
#[inline]
7891
fn decode(buffer: DecoderBuffer<'a>) -> R<'a, Self> {
7992
let (tag, buffer) = buffer.decode::<Tag>()?;
80-
decoder_invariant!(tag == Tag::default(), "invalid tag");
8193
let (credential_id, buffer) = buffer.decode()?;
8294
let (wire_version, buffer) = buffer.decode()?;
95+
let (queue_id, buffer) = if tag.has_queue_id() {
96+
let (queue_id, buffer) = buffer.decode()?;
97+
(Some(queue_id), buffer)
98+
} else {
99+
(None, buffer)
100+
};
101+
83102
let value = Self {
84103
wire_version,
85104
credential_id,
105+
queue_id,
86106
};
87107
Ok((value, buffer))
88108
}
@@ -102,7 +122,7 @@ mod tests {
102122
bolero::check!()
103123
.with_type::<(UnknownPathSecret, [u8; TAG_LEN])>()
104124
.for_each(|(value, stateless_reset)| {
105-
let mut buffer = [0u8; UnknownPathSecret::PACKET_SIZE];
125+
let mut buffer = [0u8; UnknownPathSecret::MAX_PACKET_SIZE];
106126
let len = {
107127
let encoder = s2n_codec::EncoderBuffer::new(&mut buffer);
108128
value.encode(encoder, stateless_reset)

‎dc/s2n-quic-dc/src/packet/tag.rs‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -66,22 +66,25 @@ decoder_value!(
6666
let (tag, buffer) = buffer.decode()?;
6767
Ok((Self::Control(tag), buffer))
6868
}
69-
super::secret_control::stale_key::Tag::VALUE => {
69+
super::secret_control::stale_key::Tag::VALUE
70+
| super::secret_control::stale_key::Tag::VALUE_WITH_QUEUE_ID => {
7071
let (tag, buffer) = buffer.decode()?;
7172
Ok((Self::StaleKey(tag), buffer))
7273
}
73-
super::secret_control::replay_detected::Tag::VALUE => {
74+
super::secret_control::replay_detected::Tag::VALUE
75+
| super::secret_control::replay_detected::Tag::VALUE_WITH_QUEUE_ID => {
7476
let (tag, buffer) = buffer.decode()?;
7577
Ok((Self::ReplayDetected(tag), buffer))
7678
}
77-
super::secret_control::unknown_path_secret::Tag::VALUE => {
79+
super::secret_control::unknown_path_secret::Tag::VALUE
80+
| super::secret_control::unknown_path_secret::Tag::VALUE_WITH_QUEUE_ID => {
7881
let (tag, buffer) = buffer.decode()?;
7982
Ok((Self::UnknownPathSecret(tag), buffer))
8083
}
8184
// reserve this range for other packet types
82-
0b0110_0011..=0b0111_1111 => Err(s2n_codec::DecoderError::InvariantViolation(
83-
"unexpected packet tag",
84-
)),
85+
0b0110_0011 | 0b0110_0111 | 0b0110_1000..=0b0111_1111 => Err(
86+
s2n_codec::DecoderError::InvariantViolation("unexpected packet tag"),
87+
),
8588
0b1000_0000..=0b1111_1111 => Err(s2n_codec::DecoderError::InvariantViolation(
8689
"only short packets are accepted",
8790
)),

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

Lines changed: 46 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use crate::{
99
stream::TransportFeatures,
1010
};
1111
use core::fmt;
12-
use s2n_quic_core::{dc, time};
12+
use s2n_quic_core::{dc, time, varint::VarInt};
1313
use std::{net::SocketAddr, sync::Arc};
1414

1515
mod cleaner;
@@ -143,23 +143,29 @@ impl Map {
143143
pub fn open_once(
144144
&self,
145145
credentials: &Credentials,
146+
queue_id: Option<VarInt>,
146147
control_out: &mut Vec<u8>,
147148
) -> Option<open::Once> {
148-
let entry = self.store.pre_authentication(credentials, control_out)?;
149-
let opener = entry.uni_opener(self.clone(), credentials);
149+
let entry = self
150+
.store
151+
.pre_authentication(credentials, queue_id, control_out)?;
152+
let opener = entry.uni_opener(self.clone(), credentials, queue_id);
150153
Some(opener)
151154
}
152155

153156
pub fn pair_for_credentials(
154157
&self,
155158
credentials: &Credentials,
159+
queue_id: Option<VarInt>,
156160
features: &TransportFeatures,
157161
control_out: &mut Vec<u8>,
158162
) -> Option<(entry::Bidirectional, dc::ApplicationParams)> {
159-
let entry = self.store.pre_authentication(credentials, control_out)?;
163+
let entry = self
164+
.store
165+
.pre_authentication(credentials, queue_id, control_out)?;
160166

161167
let params = entry.parameters();
162-
let keys = entry.bidi_remote(self.clone(), credentials, features);
168+
let keys = entry.bidi_remote(self.clone(), credentials, queue_id, features);
163169

164170
Some((keys, params))
165171
}
@@ -173,7 +179,41 @@ impl Map {
173179
}
174180

175181
pub fn handle_control_packet(&self, packet: &control::Packet, peer: &SocketAddr) {
176-
self.store.handle_control_packet(packet, peer)
182+
match packet {
183+
control::Packet::StaleKey(packet) => {
184+
let _ = self.handle_stale_key_packet(packet, peer);
185+
}
186+
control::Packet::ReplayDetected(packet) => {
187+
let _ = self.handle_replay_detected_packet(packet, peer);
188+
}
189+
control::Packet::UnknownPathSecret(packet) => {
190+
let _ = self.handle_unknown_path_secret_packet(packet, peer);
191+
}
192+
}
193+
}
194+
195+
pub fn handle_stale_key_packet<'a>(
196+
&self,
197+
packet: &'a control::stale_key::Packet,
198+
peer: &SocketAddr,
199+
) -> Option<&'a control::StaleKey> {
200+
self.store.handle_stale_key_packet(packet, peer)
201+
}
202+
203+
pub fn handle_replay_detected_packet<'a>(
204+
&self,
205+
packet: &'a control::replay_detected::Packet,
206+
peer: &SocketAddr,
207+
) -> Option<&'a control::ReplayDetected> {
208+
self.store.handle_replay_detected_packet(packet, peer)
209+
}
210+
211+
pub fn handle_unknown_path_secret_packet<'a>(
212+
&self,
213+
packet: &'a control::unknown_path_secret::Packet,
214+
peer: &SocketAddr,
215+
) -> Option<&'a control::UnknownPathSecret> {
216+
self.store.handle_unknown_path_secret_packet(packet, peer)
177217
}
178218

179219
#[doc(hidden)]

0 commit comments

Comments
 (0)