fix(iroh-relay): Fix proptests, make Datagrams::segment_size be an Option<NonZeroU16> (#3404)

## Description

The proptests were failing in some rare case where
`Datagrams::segment_size` was generated as `Some(0)`. (For some reason
we haven't hit that before?).

I added the failing case hash to the proptest regressions file and fixed
the problem.

I also changed `Datagrams::segment_size` to be an `Option<NonZeroU16>`
instead of `Option<u16>`, which checks this invariant on the type level
now.

## Change checklist
<!-- Remove any that are not relevant. -->
- [x] Self-review.

Author: matheus23

iroh-relay/proptest-regressions/protos/relay.txt

@@ -6,3 +6,4 @@
 # everyone who runs the test benefits from these saved cases.
 cc 9295f5287162dfb180e5826e563c2cea08b477b803ef412ff8351eb5c3eb45ef # shrinks to frame = KeepAlive
 cc 753aabcf8ae2b4e4a52f451d58339aab85a4b61108afdf4b9600f97b3a33bf42 # shrinks to frame = Health { problem: None }
+cc 2b45ae945ff922d4c3dfbad31a5e57c535c0ee2739906272f21ed77f8b862528 # shrinks to frame = Datagrams { remote_node_id: PublicKey(3b6a27bcceb6a42d62a3a8d02a6f0d73653215771de243a63ac048a18b59da29), datagrams: Datagrams { ecn: None, segment_size: Some(0), .. } }

iroh-relay/src/protos/relay.rs

@@ -7,6 +7,8 @@
 //!  * server then sends [`FrameType::RelayToClientDatagram`] or [`FrameType::RelayToClientDatagramBatch`] to recipient
 //!  * server sends [`FrameType::NodeGone`] when the other client disconnects
 
+use std::num::NonZeroU16;
+
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use iroh_base::{NodeId, SignatureError};
 use n0_future::time::Duration;
@@ -138,7 +140,7 @@ pub struct Datagrams {
     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>,
+    pub segment_size: Option<NonZeroU16>,
     /// The contents of the datagram(s)
     #[debug(skip)]
     pub contents: Bytes,
@@ -159,7 +161,7 @@ impl Datagrams {
         let ecn = self.ecn.map_or(0, |ecn| ecn as u8);
         dst.put_u8(ecn);
         if let Some(segment_size) = self.segment_size {
-            dst.put_u16(segment_size);
+            dst.put_u16(segment_size.into());
         }
         dst.put(self.contents.as_ref());
         dst
@@ -185,11 +187,7 @@ impl Datagrams {
 
         let segment_size = if is_batch {
             let segment_size = bytes.get_u16(); // length checked above
-            if segment_size == 0 {
-                None
-            } else {
-                Some(segment_size)
-            }
+            NonZeroU16::new(segment_size)
         } else {
             None
         };
@@ -519,7 +517,7 @@ mod tests {
                     remote_node_id: client_key.public(),
                     datagrams: Datagrams {
                         ecn: Some(quinn::EcnCodepoint::Ce),
-                        segment_size: Some(6),
+                        segment_size: NonZeroU16::new(6),
                         contents: "Hello World!".into(),
                     },
                 }
@@ -589,7 +587,7 @@ mod tests {
                     dst_node_id: client_key.public(),
                     datagrams: Datagrams {
                         ecn: Some(quinn::EcnCodepoint::Ce),
-                        segment_size: Some(6),
+                        segment_size: NonZeroU16::new(6),
                         contents: "Hello World!".into(),
                     },
                 }
@@ -668,7 +666,9 @@ mod proptests {
         )
             .prop_map(|(ecn, segment_size, data)| Datagrams {
                 ecn,
-                segment_size: segment_size.map(|ss| std::cmp::min(data.len(), ss) as u16),
+                segment_size: segment_size
+                    .map(|ss| std::cmp::min(data.len(), ss) as u16)
+                    .and_then(NonZeroU16::new),
                 contents: Bytes::from(data),
             })
     }

iroh/src/magicsock/transports/relay.rs

@@ -1,5 +1,6 @@
 use std::{
     io,
+    num::NonZeroU16,
     task::{Context, Poll},
 };
 
@@ -104,7 +105,7 @@ impl RelayTransport {
             meta_out.stride = dm
                 .datagrams
                 .segment_size
-                .map_or(dm.datagrams.contents.len(), |s| s as usize);
+                .map_or(dm.datagrams.contents.len(), |s| u16::from(s) as usize);
             meta_out.ecn = None;
             meta_out.dst_ip = None; // TODO: insert the relay url for this relay
 
@@ -314,7 +315,10 @@ fn datagrams_from_transmit(transmit: &Transmit<'_>) -> Datagrams {
             quinn_udp::EcnCodepoint::Ect1 => quinn_proto::EcnCodepoint::Ect1,
             quinn_udp::EcnCodepoint::Ce => quinn_proto::EcnCodepoint::Ce,
         }),
-        segment_size: transmit.segment_size.map(|ss| ss as u16),
+        segment_size: transmit
+            .segment_size
+            .map(|ss| ss as u16)
+            .and_then(NonZeroU16::new),
         contents: Bytes::copy_from_slice(transmit.contents),
     }
 }