Skip to content

Commit 184cff8

Browse files
matheus23matheus23
committed
feat(iroh-relay)!: Save 2 bytes for non-GSO datagrams
1 parent acfeda5 commit 184cff8

File tree

4 files changed

+111
-31
lines changed

4 files changed

+111
-31
lines changed

‎iroh-relay/src/protos/common.rs

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,14 @@ pub enum FrameType {
2323
ServerConfirmsAuth = 2,
2424
/// The server frame type for authentication denial
2525
ServerDeniesAuth = 3,
26+
/// 32B dest pub key + ECN bytes + one datagram's content
27+
ClientToRelayDatagram = 4,
2628
/// 32B dest pub key + ECN byte + segment size u16 + datagrams contents
27-
ClientToRelayDatagrams = 4,
29+
ClientToRelayDatagrams = 5,
30+
/// 32B src pub key + ECN bytes + one datagram's content
31+
RelayToClientDatagram = 6,
2832
/// 32B src pub key + ECN byte + segment size u16 + datagrams contents
29-
RelayToClientDatagrams = 6,
33+
RelayToClientDatagrams = 7,
3034
/// Sent from server to client to signal that a previous sender is no longer connected.
3135
///
3236
/// That is, if A sent to B, and then if A disconnects, the server sends `FrameType::PeerGone`

‎iroh-relay/src/protos/relay.rs

Lines changed: 100 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
//! * server then sends [`FrameType::RelayToClientDatagrams`] to recipient
88
//! * server sends [`FrameType::NodeGone`] when the other client disconnects
99
10-
use bytes::{BufMut, Bytes, BytesMut};
10+
use bytes::{Buf, BufMut, Bytes, BytesMut};
1111
use iroh_base::{NodeId, SignatureError};
1212
use n0_future::time::{self, Duration};
1313
use nested_enum_utils::common_fields;
@@ -165,39 +165,46 @@ impl<T: AsRef<[u8]>> From<T> for Datagrams {
165165
impl Datagrams {
166166
fn write_to<O: BufMut>(&self, mut dst: O) -> O {
167167
let ecn = self.ecn.map_or(0, |ecn| ecn as u8);
168-
let segment_size = self.segment_size.unwrap_or_default();
169168
dst.put_u8(ecn);
170-
dst.put_u16(segment_size);
169+
if let Some(segment_size) = self.segment_size {
170+
dst.put_u16(segment_size);
171+
}
171172
dst.put(self.contents.as_ref());
172173
dst
173174
}
174175

175176
fn encoded_len(&self) -> usize {
176177
1 // ECN byte
177-
+ 2 // segment size
178+
+ self.segment_size.map_or(0, |_| 2) // segment size, when None, then a packed representation is assumed
178179
+ self.contents.len()
179180
}
180181

181-
fn from_bytes(bytes: Bytes) -> Result<Self, Error> {
182-
// 1 bytes ECN, 2 bytes segment size
183-
snafu::ensure!(bytes.len() > 3, InvalidFrameSnafu);
182+
fn from_bytes(mut bytes: Bytes, is_batch: bool) -> Result<Self, Error> {
183+
if is_batch {
184+
// 1 bytes ECN, 2 bytes segment size
185+
snafu::ensure!(bytes.len() >= 3, InvalidFrameSnafu);
186+
} else {
187+
snafu::ensure!(bytes.len() >= 1, InvalidFrameSnafu);
188+
}
184189

185-
let ecn_byte = bytes[0];
190+
let ecn_byte = bytes.get_u8();
186191
let ecn = quinn_proto::EcnCodepoint::from_bits(ecn_byte);
187192

188-
let segment_size = u16::from_be_bytes(bytes[1..3].try_into().expect("length checked"));
189-
let segment_size = if segment_size == 0 {
190-
None
193+
let segment_size = if is_batch {
194+
let segment_size = bytes.get_u16(); // length checked above
195+
if segment_size == 0 {
196+
None
197+
} else {
198+
Some(segment_size)
199+
}
191200
} else {
192-
Some(segment_size)
201+
None
193202
};
194203

195-
let contents = bytes.slice(3..);
196-
197204
Ok(Self {
198205
ecn,
199206
segment_size,
200-
contents,
207+
contents: bytes,
201208
})
202209
}
203210
}
@@ -206,7 +213,13 @@ impl RelayToClientMsg {
206213
/// Returns this frame's corresponding frame type.
207214
pub fn typ(&self) -> FrameType {
208215
match self {
209-
Self::Datagrams { .. } => FrameType::RelayToClientDatagrams,
216+
Self::Datagrams { datagrams, .. } => {
217+
if datagrams.segment_size.is_some() {
218+
FrameType::RelayToClientDatagrams
219+
} else {
220+
FrameType::RelayToClientDatagram
221+
}
222+
}
210223
Self::NodeGone { .. } => FrameType::NodeGone,
211224
Self::Ping { .. } => FrameType::Ping,
212225
Self::Pong { .. } => FrameType::Pong,
@@ -289,13 +302,16 @@ impl RelayToClientMsg {
289302
);
290303

291304
let res = match frame_type {
292-
FrameType::RelayToClientDatagrams => {
305+
FrameType::RelayToClientDatagram | FrameType::RelayToClientDatagrams => {
293306
snafu::ensure!(content.len() >= NodeId::LENGTH, InvalidFrameSnafu);
294307

295308
let remote_node_id = cache
296309
.key_from_slice(&content[..NodeId::LENGTH])
297310
.context(InvalidPublicKeySnafu)?;
298-
let datagrams = Datagrams::from_bytes(content.slice(NodeId::LENGTH..))?;
311+
let datagrams = Datagrams::from_bytes(
312+
content.slice(NodeId::LENGTH..),
313+
frame_type == FrameType::RelayToClientDatagrams,
314+
)?;
299315
Self::Datagrams {
300316
remote_node_id,
301317
datagrams,
@@ -356,7 +372,13 @@ impl RelayToClientMsg {
356372
impl ClientToRelayMsg {
357373
pub(crate) fn typ(&self) -> FrameType {
358374
match self {
359-
Self::Datagrams { .. } => FrameType::ClientToRelayDatagrams,
375+
Self::Datagrams { datagrams, .. } => {
376+
if datagrams.segment_size.is_some() {
377+
FrameType::ClientToRelayDatagrams
378+
} else {
379+
FrameType::ClientToRelayDatagram
380+
}
381+
}
360382
Self::Ping { .. } => FrameType::Ping,
361383
Self::Pong { .. } => FrameType::Pong,
362384
}
@@ -415,11 +437,14 @@ impl ClientToRelayMsg {
415437
);
416438

417439
let res = match frame_type {
418-
FrameType::ClientToRelayDatagrams => {
440+
FrameType::ClientToRelayDatagram | FrameType::ClientToRelayDatagrams => {
419441
let dst_node_id = cache
420442
.key_from_slice(&content[..NodeId::LENGTH])
421443
.context(InvalidPublicKeySnafu)?;
422-
let datagrams = Datagrams::from_bytes(content.slice(NodeId::LENGTH..))?;
444+
let datagrams = Datagrams::from_bytes(
445+
content.slice(NodeId::LENGTH..),
446+
frame_type == FrameType::ClientToRelayDatagrams,
447+
)?;
423448
Self::Datagrams {
424449
dst_node_id,
425450
datagrams,
@@ -508,9 +533,39 @@ mod tests {
508533
},
509534
}
510535
.write_to(Vec::new()),
511-
"06 19 7f 6b 23 e1 6c 85 32 c6 ab c8 38 fa cd 5e
512-
a7 89 be 0c 76 b2 92 03 34 03 9b fa 8b 3d 36 8d
513-
61 48 65 6c 6c 6f 20 57 6f 72 6c 64 21",
536+
// frame type
537+
// public key first 16 bytes
538+
// public key second 16 bytes
539+
// ECN byte
540+
// segment size
541+
// hello world contents bytes
542+
"07
543+
19 7f 6b 23 e1 6c 85 32 c6 ab c8 38 fa cd 5e a7
544+
89 be 0c 76 b2 92 03 34 03 9b fa 8b 3d 36 8d 61
545+
03
546+
00 06
547+
48 65 6c 6c 6f 20 57 6f 72 6c 64 21",
548+
),
549+
(
550+
RelayToClientMsg::Datagrams {
551+
remote_node_id: client_key.public(),
552+
datagrams: Datagrams {
553+
ecn: Some(quinn::EcnCodepoint::Ce),
554+
segment_size: None,
555+
contents: "Hello World!".into(),
556+
},
557+
}
558+
.write_to(Vec::new()),
559+
// frame type
560+
// public key first 16 bytes
561+
// public key second 16 bytes
562+
// ECN byte
563+
// hello world contents bytes
564+
"06
565+
19 7f 6b 23 e1 6c 85 32 c6 ab c8 38 fa cd 5e a7
566+
89 be 0c 76 b2 92 03 34 03 9b fa 8b 3d 36 8d 61
567+
03
568+
48 65 6c 6c 6f 20 57 6f 72 6c 64 21",
514569
),
515570
(
516571
RelayToClientMsg::Restarting {
@@ -552,6 +607,27 @@ mod tests {
552607
a7 89 be 0c 76 b2 92 03 34 03 9b fa 8b 3d 36 8d
553608
61 47 6f 6f 64 62 79 65 21",
554609
),
610+
(
611+
ClientToRelayMsg::Datagrams {
612+
dst_node_id: client_key.public(),
613+
datagrams: Datagrams {
614+
ecn: Some(quinn::EcnCodepoint::Ce),
615+
segment_size: None,
616+
contents: "Hello World!".into(),
617+
},
618+
}
619+
.write_to(Vec::new()),
620+
// frame type
621+
// public key first 16 bytes
622+
// public key second 16 bytes
623+
// ECN byte
624+
// hello world contents
625+
"04
626+
19 7f 6b 23 e1 6c 85 32 c6 ab c8 38 fa cd 5e a7
627+
89 be 0c 76 b2 92 03 34 03 9b fa 8b 3d 36 8d 61
628+
03
629+
48 65 6c 6c 6f 20 57 6f 72 6c 64 21",
630+
),
555631
]);
556632

557633
Ok(())

‎iroh-relay/src/server/client.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -640,7 +640,7 @@ mod tests {
640640
data: Datagrams::from(&data[..]),
641641
};
642642
send_queue_s.send(packet.clone()).await.context("send")?;
643-
let frame = recv_frame(FrameType::RelayToClientDatagrams, &mut io_rw)
643+
let frame = recv_frame(FrameType::RelayToClientDatagram, &mut io_rw)
644644
.await
645645
.e()?;
646646
assert_eq!(
@@ -657,7 +657,7 @@ mod tests {
657657
.send(packet.clone())
658658
.await
659659
.context("send")?;
660-
let frame = recv_frame(FrameType::RelayToClientDatagrams, &mut io_rw)
660+
let frame = recv_frame(FrameType::RelayToClientDatagram, &mut io_rw)
661661
.await
662662
.e()?;
663663
assert_eq!(
@@ -731,7 +731,7 @@ mod tests {
731731
let mut stream = RelayedStream::test_limited(io_read, LIMIT / 10, LIMIT)?;
732732

733733
// Prepare a frame to send, assert its size.
734-
let data = Datagrams::from(b"hello world!!1");
734+
let data = Datagrams::from(b"hello world!!!!!");
735735
let target = SecretKey::generate(rand::thread_rng()).public();
736736
let frame = ClientToRelayMsg::Datagrams {
737737
dst_node_id: target,

‎iroh-relay/src/server/clients.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -255,7 +255,7 @@ mod tests {
255255
// send packet
256256
let data = b"hello world!";
257257
clients.send_packet(a_key, Datagrams::from(&data[..]), b_key, &metrics)?;
258-
let frame = recv_frame(FrameType::RelayToClientDatagrams, &mut a_rw).await?;
258+
let frame = recv_frame(FrameType::RelayToClientDatagram, &mut a_rw).await?;
259259
assert_eq!(
260260
frame,
261261
RelayToClientMsg::Datagrams {
@@ -266,7 +266,7 @@ mod tests {
266266

267267
// send disco packet
268268
clients.send_disco_packet(a_key, Datagrams::from(&data[..]), b_key, &metrics)?;
269-
let frame = recv_frame(FrameType::RelayToClientDatagrams, &mut a_rw).await?;
269+
let frame = recv_frame(FrameType::RelayToClientDatagram, &mut a_rw).await?;
270270
assert_eq!(
271271
frame,
272272
RelayToClientMsg::Datagrams {

0 commit comments

Comments
 (0)