Send ECN bits and use stride instead of custom split protocol

Author: matheus23

iroh-relay/src/client/conn.rs

@@ -19,7 +19,7 @@ use crate::client::streams::{MaybeTlsStream, ProxyStream};
 use crate::{
     protos::{
         handshake,
-        relay::{ClientToRelayMsg, Error as ProtoError, RelayToClientMsg},
+        relay::{ClientToRelayMsg, Error as ProtoError, RelayToClientMsg, MAX_PAYLOAD_SIZE},
         streams::WsBytesFramed,
     },
     MAX_PACKET_SIZE,
@@ -143,9 +143,9 @@ impl Sink<ClientToRelayMsg> for Conn {
     }
 
     fn start_send(mut self: Pin<&mut Self>, frame: ClientToRelayMsg) -> Result<(), Self::Error> {
-        if let ClientToRelayMsg::SendPacket { .. } = &frame {
-            let size = frame.encoded_len();
-            snafu::ensure!(size <= MAX_PACKET_SIZE, ExceedsMaxPacketSizeSnafu { size });
+        if let ClientToRelayMsg::Datagrams { datagrams, .. } = &frame {
+            let size = datagrams.contents.len();
+            snafu::ensure!(size <= MAX_PAYLOAD_SIZE, ExceedsMaxPacketSizeSnafu { size });
         }
 
         Pin::new(&mut self.conn)

iroh-relay/src/protos/common.rs

@@ -24,9 +24,9 @@ pub enum FrameType {
     /// The server frame type for authentication denial
     ServerDeniesAuth = 3,
     /// 32B dest pub key + ECN byte + segment size u16 + datagrams contents
-    SendPacket = 4,
+    ClientToRelayDatagrams = 4,
     /// 32B src pub key + ECN byte + segment size u16 + datagrams contents
-    RecvPacket = 6,
+    RelayToClientDatagrams = 6,
     /// Sent from server to client to signal that a previous sender is no longer connected.
     ///
     /// That is, if A sent to B, and then if A disconnects, the server sends `FrameType::PeerGone`

iroh-relay/src/protos/relay.rs

@@ -21,6 +21,12 @@ use crate::KeyCache;
 /// including its on-wire framing overhead)
 pub const MAX_PACKET_SIZE: usize = 64 * 1024;
 
+/// Maximum size a datagram payload is allowed to be.
+///
+/// This is [`MAX_PACKET_SIZE`] minus the length of an encoded public key minus 3 bytes,
+/// one for ECN, and two for the segment size.
+pub const MAX_PAYLOAD_SIZE: usize = MAX_PACKET_SIZE - NodeId::LENGTH - 3;
+
 /// The maximum frame size.
 ///
 /// This is also the minimum burst size that a rate-limiter has to accept.
@@ -74,12 +80,11 @@ pub enum Error {
 #[derive(derive_more::Debug, Clone, PartialEq, Eq)]
 pub enum RelayToClientMsg {
     /// Represents datagrams sent from relays (originally sent to them by another client).
-    ReceivedPacket {
+    Datagrams {
         /// The [`NodeId`] of the original sender.
-        src_key: NodeId,
-        /// The received packet bytes.
-        #[debug(skip)]
-        content: Bytes,
+        remote_node_id: NodeId,
+        /// The datagrams and related metadata.
+        datagrams: Datagrams,
     },
     /// Indicates that the client identified by the underlying public key had previously sent you a
     /// packet but has now disconnected from the relay.
@@ -123,19 +128,85 @@ pub enum ClientToRelayMsg {
     /// with the payload sent previously in the ping.
     Pong([u8; 8]),
     /// Request from the client to relay datagrams to given remote node.
-    SendPacket {
+    Datagrams {
         /// The remote node to relay to.
-        dst_key: NodeId,
+        dst_node_id: NodeId,
         /// The datagrams and related metadata to relay.
-        packet: Bytes,
+        datagrams: Datagrams,
     },
 }
 
+/// One or multiple datagrams being transferred via the relay.
+///
+/// This type is modeled after [`quinn_proto::Transmit`]
+/// (or even more similarly `quinn_udp::Transmit`, but we don't depend on that library here).
+#[derive(derive_more::Debug, Clone, PartialEq, Eq)]
+pub struct Datagrams {
+    /// Explicit congestion notification bits
+    pub ecn: Option<quinn_proto::EcnCodepoint>,
+    /// The segment size if this transmission contains multiple datagrams.
+    /// This is `None` if the transmit only contains a single datagram
+    pub segment_size: Option<u16>,
+    /// The contents of the datagram(s)
+    #[debug(skip)]
+    pub contents: Bytes,
+}
+
+impl<T: AsRef<[u8]>> From<T> for Datagrams {
+    fn from(bytes: T) -> Self {
+        Self {
+            ecn: None,
+            segment_size: None,
+            contents: Bytes::copy_from_slice(bytes.as_ref()),
+        }
+    }
+}
+
+impl Datagrams {
+    fn write_to<O: BufMut>(&self, mut dst: O) -> O {
+        let ecn = self.ecn.map_or(0, |ecn| ecn as u8);
+        let segment_size = self.segment_size.unwrap_or_default();
+        dst.put_u8(ecn);
+        dst.put_u16(segment_size);
+        dst.put(self.contents.as_ref());
+        dst
+    }
+
+    fn encoded_len(&self) -> usize {
+        1 // ECN byte
+        + 2 // segment size
+        + self.contents.len()
+    }
+
+    fn from_bytes(bytes: Bytes) -> Result<Self, Error> {
+        // 1 bytes ECN, 2 bytes segment size
+        snafu::ensure!(bytes.len() > 3, InvalidFrameSnafu);
+
+        let ecn_byte = bytes[0];
+        let ecn = quinn_proto::EcnCodepoint::from_bits(ecn_byte);
+
+        let segment_size = u16::from_be_bytes(bytes[1..3].try_into().expect("length checked"));
+        let segment_size = if segment_size == 0 {
+            None
+        } else {
+            Some(segment_size)
+        };
+
+        let contents = bytes.slice(3..);
+
+        Ok(Self {
+            ecn,
+            segment_size,
+            contents,
+        })
+    }
+}
+
 impl RelayToClientMsg {
     /// Returns this frame's corresponding frame type.
     pub fn typ(&self) -> FrameType {
         match self {
-            Self::ReceivedPacket { .. } => FrameType::RecvPacket,
+            Self::Datagrams { .. } => FrameType::RelayToClientDatagrams,
             Self::NodeGone { .. } => FrameType::NodeGone,
             Self::Ping { .. } => FrameType::Ping,
             Self::Pong { .. } => FrameType::Pong,
@@ -156,12 +227,12 @@ impl RelayToClientMsg {
     pub(crate) fn write_to<O: BufMut>(&self, mut dst: O) -> O {
         dst = self.typ().write_to(dst);
         match self {
-            Self::ReceivedPacket {
-                src_key: remote_node_id,
-                content,
+            Self::Datagrams {
+                remote_node_id,
+                datagrams,
             } => {
                 dst.put(remote_node_id.as_ref());
-                dst.put(content.as_ref());
+                dst = datagrams.write_to(dst);
             }
             Self::NodeGone(node_id) => {
                 dst.put(node_id.as_ref());
@@ -189,9 +260,9 @@ impl RelayToClientMsg {
     #[cfg(feature = "server")]
     pub(crate) fn encoded_len(&self) -> usize {
         let payload_len = match self {
-            Self::ReceivedPacket { content, .. } => {
+            Self::Datagrams { datagrams, .. } => {
                 32 // nodeid
-                + content.len()
+                + datagrams.encoded_len()
             }
             Self::NodeGone(_) => 32,
             Self::Ping(_) | Self::Pong(_) => 8,
@@ -218,14 +289,17 @@ impl RelayToClientMsg {
         );
 
         let res = match frame_type {
-            FrameType::RecvPacket => {
+            FrameType::RelayToClientDatagrams => {
                 snafu::ensure!(content.len() >= NodeId::LENGTH, InvalidFrameSnafu);
 
-                let src_key = cache
+                let remote_node_id = cache
                     .key_from_slice(&content[..NodeId::LENGTH])
                     .context(InvalidPublicKeySnafu)?;
-                let content = content.slice(NodeId::LENGTH..);
-                Self::ReceivedPacket { src_key, content }
+                let datagrams = Datagrams::from_bytes(content.slice(NodeId::LENGTH..))?;
+                Self::Datagrams {
+                    remote_node_id,
+                    datagrams,
+                }
             }
             FrameType::NodeGone => {
                 snafu::ensure!(content.len() == NodeId::LENGTH, InvalidFrameSnafu);
@@ -282,7 +356,7 @@ impl RelayToClientMsg {
 impl ClientToRelayMsg {
     pub(crate) fn typ(&self) -> FrameType {
         match self {
-            Self::SendPacket { .. } => FrameType::SendPacket,
+            Self::Datagrams { .. } => FrameType::ClientToRelayDatagrams,
             Self::Ping { .. } => FrameType::Ping,
             Self::Pong { .. } => FrameType::Pong,
         }
@@ -298,9 +372,12 @@ impl ClientToRelayMsg {
     pub(crate) fn write_to<O: BufMut>(&self, mut dst: O) -> O {
         dst = self.typ().write_to(dst);
         match self {
-            Self::SendPacket { dst_key, packet } => {
-                dst.put(dst_key.as_ref());
-                dst.put(packet.as_ref());
+            Self::Datagrams {
+                dst_node_id,
+                datagrams,
+            } => {
+                dst.put(dst_node_id.as_ref());
+                dst = datagrams.write_to(dst);
             }
             Self::Ping(data) => {
                 dst.put(&data[..]);
@@ -315,9 +392,9 @@ impl ClientToRelayMsg {
     pub(crate) fn encoded_len(&self) -> usize {
         let payload_len = match self {
             Self::Ping(_) | Self::Pong(_) => 8,
-            Self::SendPacket { packet, .. } => {
+            Self::Datagrams { datagrams, .. } => {
                 32 // node id
-                + packet.len()
+                + datagrams.encoded_len()
             }
         };
         1 // frame type (all frame types currently encode as 1 byte varint)
@@ -338,12 +415,15 @@ impl ClientToRelayMsg {
         );
 
         let res = match frame_type {
-            FrameType::SendPacket => {
-                let dst_key = cache
+            FrameType::ClientToRelayDatagrams => {
+                let dst_node_id = cache
                     .key_from_slice(&content[..NodeId::LENGTH])
                     .context(InvalidPublicKeySnafu)?;
-                let packet = content.slice(NodeId::LENGTH..);
-                Self::SendPacket { dst_key, packet }
+                let datagrams = Datagrams::from_bytes(content.slice(NodeId::LENGTH..))?;
+                Self::Datagrams {
+                    dst_node_id,
+                    datagrams,
+                }
             }
             FrameType::Ping => {
                 snafu::ensure!(content.len() == 8, InvalidFrameSnafu);
@@ -419,9 +499,13 @@ mod tests {
                 "0a 2a 2a 2a 2a 2a 2a 2a 2a",
             ),
             (
-                RelayToClientMsg::ReceivedPacket {
-                    src_key: client_key.public(),
-                    content: "Hello World!".into(),
+                RelayToClientMsg::Datagrams {
+                    remote_node_id: client_key.public(),
+                    datagrams: Datagrams {
+                        ecn: Some(quinn::EcnCodepoint::Ce),
+                        segment_size: Some(6),
+                        contents: "Hello World!".into(),
+                    },
                 }
                 .write_to(Vec::new()),
                 "06 19 7f 6b 23 e1 6c 85 32 c6 ab c8 38 fa cd 5e
@@ -455,9 +539,13 @@ mod tests {
                 "0a 2a 2a 2a 2a 2a 2a 2a 2a",
             ),
             (
-                ClientToRelayMsg::SendPacket {
-                    dst_key: client_key.public(),
-                    packet: "Goodbye!".into(),
+                ClientToRelayMsg::Datagrams {
+                    dst_node_id: client_key.public(),
+                    datagrams: Datagrams {
+                        ecn: Some(quinn::EcnCodepoint::Ce),
+                        segment_size: Some(6),
+                        contents: "Hello World!".into(),
+                    },
                 }
                 .write_to(Vec::new()),
                 "04 19 7f 6b 23 e1 6c 85 32 c6 ab c8 38 fa cd 5e
@@ -485,22 +573,42 @@ mod proptests {
         secret_key().prop_map(|key| key.public())
     }
 
-    /// Generates random data, up to the maximum packet size minus the given number of bytes
-    fn data(consumed: usize) -> impl Strategy<Value = Bytes> {
-        let len = MAX_PACKET_SIZE - consumed;
-        prop::collection::vec(any::<u8>(), 0..len).prop_map(Bytes::from)
+    fn ecn() -> impl Strategy<Value = Option<quinn_proto::EcnCodepoint>> {
+        (0..=3).prop_map(|n| match n {
+            1 => Some(quinn_proto::EcnCodepoint::Ce),
+            2 => Some(quinn_proto::EcnCodepoint::Ect0),
+            3 => Some(quinn_proto::EcnCodepoint::Ect1),
+            _ => None,
+        })
+    }
+
+    fn datagrams() -> impl Strategy<Value = Datagrams> {
+        (
+            ecn(),
+            prop::option::of(MAX_PAYLOAD_SIZE / 20..MAX_PAYLOAD_SIZE),
+            prop::collection::vec(any::<u8>(), 0..MAX_PAYLOAD_SIZE),
+        )
+            .prop_map(|(ecn, segment_size, data)| Datagrams {
+                ecn,
+                segment_size: segment_size.map(|ss| std::cmp::min(data.len(), ss) as u16),
+                contents: Bytes::from(data),
+            })
     }
 
     /// Generates a random valid frame
     fn server_client_frame() -> impl Strategy<Value = RelayToClientMsg> {
-        let recv_packet = (key(), data(32))
-            .prop_map(|(src_key, content)| RelayToClientMsg::ReceivedPacket { src_key, content });
-        let node_gone = key().prop_map(|node_id| RelayToClientMsg::NodeGone(node_id));
+        let recv_packet = (key(), datagrams()).prop_map(|(remote_node_id, datagrams)| {
+            RelayToClientMsg::Datagrams {
+                remote_node_id,
+                datagrams,
+            }
+        });
+        let node_gone = key().prop_map(RelayToClientMsg::NodeGone);
         let ping = prop::array::uniform8(any::<u8>()).prop_map(RelayToClientMsg::Ping);
         let pong = prop::array::uniform8(any::<u8>()).prop_map(RelayToClientMsg::Pong);
         let health = ".{0,65536}"
-            .prop_filter("exceeds max payload size", |s| {
-                s.len() < 65536 // a single unicode character can match a regex "." but take up multiple bytes
+            .prop_filter("exceeds MAX_PAYLOAD_SIZE", |s| {
+                s.len() < MAX_PAYLOAD_SIZE // a single unicode character can match a regex "." but take up multiple bytes
             })
             .prop_map(|problem| RelayToClientMsg::Health { problem });
         let restarting = (any::<u32>(), any::<u32>()).prop_map(|(reconnect_in, try_for)| {
@@ -513,8 +621,11 @@ mod proptests {
     }
 
     fn client_server_frame() -> impl Strategy<Value = ClientToRelayMsg> {
-        let send_packet = (key(), data(32))
-            .prop_map(|(dst_key, packet)| ClientToRelayMsg::SendPacket { dst_key, packet });
+        let send_packet =
+            (key(), datagrams()).prop_map(|(dst_node_id, datagrams)| ClientToRelayMsg::Datagrams {
+                dst_node_id,
+                datagrams,
+            });
         let ping = prop::array::uniform8(any::<u8>()).prop_map(ClientToRelayMsg::Ping);
         let pong = prop::array::uniform8(any::<u8>()).prop_map(ClientToRelayMsg::Pong);
         prop_oneof![send_packet, ping, pong]
@@ -548,5 +659,12 @@ mod proptests {
             let actual_encoded_len = frame.to_bytes().len();
             prop_assert_eq!(claimed_encoded_len, actual_encoded_len);
         }
+
+        #[test]
+        fn datagrams_encoded_len(datagrams in datagrams()) {
+            let claimed_encoded_len = datagrams.encoded_len();
+            let actual_encoded_len = datagrams.write_to(Vec::new()).len();
+            prop_assert_eq!(claimed_encoded_len, actual_encoded_len);
+        }
     }
 }