fix(proto): do not coalesce into a too small datagram tail

noq-proto/src/connection/mod.rs

@@ -1595,8 +1595,9 @@ impl Connection {
                 pad_datagram = PadDatagram::No;
             }
 
-            // If coalescing another packet into the existing datagram, there should
-            // still be enough space for a whole packet.
+            // If a previous space left a datagram whose tail is too small for another
+            // packet, its CONNECTION_CLOSE branch finished that datagram, so anything we
+            // coalesce into still has room for a whole packet.
             if transmit.datagram_start_offset() < transmit.len() {
                 debug_assert!(transmit.datagram_remaining_mut() >= MIN_PACKET_SPACE);
             }
@@ -1710,15 +1711,17 @@ impl Connection {
                 // Send a close frame in every possible space for robustness, per
                 // RFC9000 "Immediate Close during the Handshake". Don't bother trying
                 // to send anything else.
-                // TODO(flub): This breaks during the handshake if we can not coalesce
-                //    packets due to space reasons: the next space would either fail a
-                //    debug_assert checking for enough packet space or produce an invalid
-                //    packet. We need to keep track of per-space pending CONNECTION_CLOSE to
-                //    be able to send these across multiple calls to poll_transmit. Then
-                //    check for coalescing space here because initial packets need to be in
-                //    padded datagrams. And also add space checks for CONNECTION_CLOSE in
-                //    space_can_send so it would stop a GSO batch if the datagram is too
-                //    small for another CONNECTION_CLOSE packet.
+                // TODO(flub): We need to keep track of per-space pending CONNECTION_CLOSE to
+                //    be able to send these across multiple calls to poll_transmit. And also
+                //    add space checks for CONNECTION_CLOSE in space_can_send so it would
+                //    stop a GSO batch if the datagram is too small for another
+                //    CONNECTION_CLOSE packet.
+                // If what is left of this datagram is too small for another packet, finish
+                // it so the next space starts a fresh datagram rather than a packet being
+                // coalesced past its end.
+                if transmit.datagram_remaining_mut() < MIN_PACKET_SPACE {
+                    transmit.finish_datagram();
+                }
                 return PollPathSpaceStatus::WrotePacket {
                     last_packet_number: last_pn,
                     pad_datagram,

noq-proto/src/connection/transmit_buf.rs

@@ -153,6 +153,24 @@ impl<'a> TransmitBuf<'a> {
         self.buf_capacity = self.buf.len();
     }
 
+    /// Finishes the current datagram, so the next packet starts a new datagram
+    ///
+    /// Used when what is left of the current datagram is too small to hold another packet.
+    /// The datagram then ends up shorter than the segment size: if it is the first
+    /// datagram of the batch the segment size is clipped to it, otherwise the batch has to
+    /// end with it, because only the first and the last datagram of a GSO batch may be
+    /// smaller than the segment size.
+    pub(super) fn finish_datagram(&mut self) {
+        debug_assert!(self.num_datagrams > 0);
+        if self.num_datagrams == 1 {
+            self.clip_segment_size();
+        } else {
+            // This datagram is shorter than the segment size, so the batch has to end here.
+            self.buf_capacity = self.buf.len();
+            self.max_datagrams = NonZeroUsize::new(self.num_datagrams).unwrap_or(NonZeroUsize::MIN);
+        }
+    }
+
     /// Returns the GSO segment size
     ///
     /// This is also the maximum size datagrams are allowed to be. The first and last

noq-proto/src/tests/mod.rs

@@ -4937,3 +4937,82 @@ fn regression_initial_coalescing_large_cid() {
     pair.time += Duration::from_secs(5);
     pair.drive_client(); // this used to try to build a packet without enough datagram space
 }
+
+/// A CONNECTION_CLOSE written while more than one packet number space still has keys hands
+/// its datagram to the next space to coalesce into, without checking that what is left of
+/// the datagram can hold another packet.
+///
+/// A client pads its Initial CONNECTION_CLOSE packet to `MIN_INITIAL_SIZE`. With an
+/// `initial_mtu` above `MIN_INITIAL_SIZE` that leaves a tail smaller than a packet, and the
+/// Handshake space used to write into that tail: a failed `MIN_PACKET_SPACE` debug assertion
+/// in the debug profile, a packet written past the end of the datagram in the release
+/// profile.
+///
+/// The tail has to stay below `MIN_PACKET_SPACE`, above that it holds a whole packet and
+/// coalescing into it is fine. `min_mtu` at `MIN_INITIAL_SIZE` and an `initial_mtu` of 1250
+/// leave 50 bytes.
+#[test]
+fn close_during_handshake_does_not_coalesce_into_a_too_small_datagram_tail() {
+    let _guard = subscribe();
+
+    const MIN_MTU: u16 = MIN_INITIAL_SIZE;
+    const MTU: u16 = 1250;
+
+    let mut transport = TransportConfig::default();
+    transport.min_mtu(MIN_MTU);
+    transport.initial_mtu(MTU);
+    transport.mtu_discovery_config(None);
+    let transport = Arc::new(transport);
+
+    let mut server_cfg = server_config();
+    server_cfg.transport = transport.clone();
+    let mut pair = Pair::new(Default::default(), server_cfg);
+
+    let mut client_cfg = client_config();
+    client_cfg.transport = transport;
+
+    // The client's Initial goes out and the server's flight comes back, but the client only
+    // *receives* it: it now has Handshake keys while still holding its Initial keys, because
+    // a client discards those when it sends its first Handshake packet, which it has not
+    // done yet.
+    let client_ch = pair.begin_connect(client_cfg);
+    pair.drive_client();
+    pair.drive_server();
+
+    let now = pair.time;
+    pair.client.drive_incoming(now);
+    let events: Vec<_> = pair
+        .client
+        .conn_events
+        .remove(&client_ch)
+        .unwrap_or_default()
+        .into_iter()
+        .collect();
+    let conn = pair.client.connections.get_mut(&client_ch).unwrap();
+    for event in events {
+        conn.handle_event(event);
+    }
+    conn.close(now, VarInt::from_u32(42), Bytes::from_static(b"bye"));
+
+    let mut buf = Vec::new();
+    let transmit = conn
+        .poll_transmit(now, NonZeroUsize::new(10).expect("known"), &mut buf)
+        .expect("a CONNECTION_CLOSE is sent");
+
+    // The Initial CONNECTION_CLOSE packet fills its datagram to MIN_INITIAL_SIZE, and the
+    // 50 byte tail that leaves in a MTU sized datagram is not handed on: the Handshake and
+    // Data CONNECTION_CLOSE packets go into a datagram of their own.
+    let segment_size = transmit
+        .segment_size
+        .expect("more than one datagram was written");
+    assert_eq!(
+        segment_size,
+        usize::from(MIN_INITIAL_SIZE),
+        "the datagram carrying the Initial CONNECTION_CLOSE ends at MIN_INITIAL_SIZE"
+    );
+    assert!(
+        transmit.size > segment_size,
+        "expected a CONNECTION_CLOSE in a further space, got {} bytes",
+        transmit.size
+    );
+}