feat: reliable stream reset

noq-proto/src/config/mod.rs

@@ -36,13 +36,15 @@ pub use transport::{AckFrequencyConfig, IdleTimeout, MtuDiscoveryConfig, Transpo
 /// Global configuration for the endpoint, affecting all connections
 ///
 /// Default values should be suitable for most internet applications.
-#[derive(Clone)]
+#[derive(Clone, derive_more::Debug)]
 pub struct EndpointConfig {
+    #[debug("HmacKey")]
     pub(crate) reset_key: Arc<dyn HmacKey>,
     pub(crate) max_udp_payload_size: VarInt,
     /// CID generator factory
     ///
     /// Create a cid generator for local cid in Endpoint struct
+    #[debug("ConnectionIdGenerator")]
     pub(crate) connection_id_generator_factory:
         Arc<dyn Fn() -> Box<dyn ConnectionIdGenerator> + Send + Sync>,
     pub(crate) supported_versions: Vec<u32>,
@@ -51,6 +53,7 @@ pub struct EndpointConfig {
     pub(crate) min_reset_interval: Duration,
     /// Optional seed to be used internally for random number generation
     pub(crate) rng_seed: Option<[u8; 32]>,
+    pub(crate) reset_stream_at: bool,
 }
 
 impl EndpointConfig {
@@ -66,6 +69,7 @@ impl EndpointConfig {
             grease_quic_bit: true,
             min_reset_interval: Duration::from_millis(20),
             rng_seed: None,
+            reset_stream_at: true,
         }
     }
 
@@ -162,18 +166,13 @@ impl EndpointConfig {
         self.rng_seed = seed;
         self
     }
-}
 
-impl fmt::Debug for EndpointConfig {
-    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt.debug_struct("EndpointConfig")
-            // reset_key not debug
-            .field("max_udp_payload_size", &self.max_udp_payload_size)
-            // cid_generator_factory not debug
-            .field("supported_versions", &self.supported_versions)
-            .field("grease_quic_bit", &self.grease_quic_bit)
-            .field("rng_seed", &self.rng_seed)
-            .finish_non_exhaustive()
+    /// Enables the QUIC Stream Resets with Partial Delivery draft-07 extentsion.
+    ///
+    /// <https://datatracker.ietf.org/doc/html/draft-ietf-quic-reliable-stream-reset>
+    pub fn reliable_stream_reset(&mut self, value: bool) -> &mut Self {
+        self.reset_stream_at = value;
+        self
     }
 }
 

noq-proto/src/connection/assembler.rs

@@ -58,6 +58,20 @@ impl Assembler {
 
     /// Get the the next chunk
     pub(super) fn read(&mut self, max_length: usize, ordered: bool) -> Option<Chunk> {
+        self.read_capped(max_length, ordered, u64::MAX)
+    }
+
+    /// Like [`Self::read`], but never returns data at or beyond `offset_limit`.
+    ///
+    /// Used to deliver only the reliable prefix of a stream subject to a RESET_STREAM_AT. The limit
+    /// is applied by stream offset rather than by the read cursor, because in unordered mode the
+    /// cursor is a running total of bytes handed out, not the contiguous prefix offset.
+    pub(super) fn read_capped(
+        &mut self,
+        max_length: usize,
+        ordered: bool,
+        offset_limit: u64,
+    ) -> Option<Chunk> {
         loop {
             let mut chunk = self.data.peek_mut()?;
 
@@ -82,6 +96,13 @@ impl Assembler {
                 }
             }
 
+            // Never hand out data at or beyond the offset limit, and never let a single chunk
+            // straddle it.
+            if chunk.offset >= offset_limit {
+                return None;
+            }
+            let max_length = max_length.min((offset_limit - chunk.offset) as usize);
+
             return Some(if max_length < chunk.bytes.len() {
                 self.bytes_read += max_length as u64;
                 let offset = chunk.offset;

noq-proto/src/connection/mod.rs

@@ -98,8 +98,8 @@ pub use streams::StreamsState;
 #[cfg(not(fuzzing))]
 use streams::StreamsState;
 pub use streams::{
-    Chunks, ClosedStream, FinishError, ReadError, ReadableError, RecvStream, SendStream,
-    ShouldTransmit, StreamEvent, Streams, WriteError,
+    Chunks, ClosedStream, FinishError, ReadError, ReadableError, RecvStream, ResetStreamAtError,
+    SendStream, ShouldTransmit, StreamEvent, Streams, WriteError,
 };
 
 mod timer;
@@ -3152,6 +3152,9 @@ impl Connection {
             for (id, _) in retransmits.reset_stream.iter() {
                 self.streams.reset_acked(*id);
             }
+            for (id, reliable_size) in retransmits.reset_stream_at.iter() {
+                self.streams.reset_at_acked(*id, *reliable_size);
+            }
         }
 
         for frame in info.stream_frames {
@@ -5099,6 +5102,18 @@ impl Connection {
                         self.spaces[SpaceId::Data].pending.max_data = true;
                     }
                 }
+                Frame::ResetStreamAt(frame) => {
+                    // We only advertise (and thus permit) RESET_STREAM_AT when the extension is
+                    // enabled; receiving one otherwise is a protocol violation.
+                    if !self.endpoint_config.reset_stream_at {
+                        return Err(TransportError::PROTOCOL_VIOLATION(
+                            "RESET_STREAM_AT frame received without negotiating the extension",
+                        ));
+                    }
+                    if self.streams.received_reset_at(frame)?.should_transmit() {
+                        self.spaces[SpaceId::Data].pending.max_data = true;
+                    }
+                }
                 Frame::DataBlocked(DataBlocked(offset)) => {
                     debug!(offset, "peer claims to be blocked at connection level");
                 }
@@ -7725,6 +7740,10 @@ impl SentFrames {
             StreamsBlocked(streams_blocked) => {
                 self.retransmits_mut().streams_blocked[streams_blocked.dir as usize] = true
             }
+            ResetStreamAt(frame) => self
+                .retransmits_mut()
+                .reset_stream_at
+                .push((frame.id, frame.reliable_size)),
         }
     }
 }

noq-proto/src/connection/qlog.rs

@@ -989,6 +989,17 @@ impl ToQlog for frame::StreamMetaEncoder {
     }
 }
 
+#[cfg(feature = "qlog")]
+impl ToQlog for frame::ResetStreamAt {
+    fn to_qlog(&self) -> QuicFrame {
+        // TODO: Teach qlog about this frame type.
+        QuicFrame::Unknown {
+            frame_type_bytes: Some(self.get_type().to_u64()),
+            raw: None,
+        }
+    }
+}
+
 #[cfg(feature = "qlog")]
 impl Frame {
     /// Converts a [`crate::Frame`] into a [`QuicFrame`].
@@ -1077,6 +1088,7 @@ impl Frame {
             Self::AddAddress(f) => f.to_qlog(),
             Self::ReachOut(f) => f.to_qlog(),
             Self::RemoveAddress(f) => f.to_qlog(),
+            Self::ResetStreamAt(f) => f.to_qlog(),
         }
     }
 }

noq-proto/src/connection/send_buffer.rs

@@ -94,6 +94,47 @@ impl SendBufferData {
         }
     }
 
+    /// Discard data from the end of the buffer, retaining only the bytes before `offset`.
+    ///
+    /// `offset` is an absolute stream offset. Values at or beyond [`Self::range`]`.end` are a
+    /// no-op (nothing to drop); values at or below [`Self::range`]`.start` clear the buffer. The
+    /// front offset ([`Self::range`]`.start`) is never moved, so already-acknowledged-and-popped
+    /// data is unaffected.
+    fn truncate(&mut self, offset: u64) {
+        // Number of bytes to retain from the front of the buffer.
+        let new_len = offset.saturating_sub(self.offset).min(self.len as u64) as usize;
+        if new_len >= self.len {
+            return;
+        }
+
+        // Walk the stored segments (and finally `last_segment`), keeping the first `new_len`
+        // bytes and dropping everything after them.
+        let mut kept = 0usize;
+        for segment in self.segments.iter_mut() {
+            if kept >= new_len {
+                segment.clear();
+            } else if kept + segment.len() > new_len {
+                segment.truncate(new_len - kept);
+                kept = new_len;
+            } else {
+                kept += segment.len();
+            }
+        }
+        // Any remainder lives in `last_segment`. If the segments already cover `new_len`, the
+        // whole `last_segment` is beyond the cut and is dropped.
+        self.last_segment.truncate(new_len.saturating_sub(kept));
+
+        self.len = new_len;
+
+        // remove empty segments
+        self.segments.retain(|s| !s.is_empty());
+
+        // shrink segments if we have a lot of unused capacity
+        if self.segments.len() * 4 < self.segments.capacity() {
+            self.segments.shrink_to_fit();
+        }
+    }
+
     /// Discard data from the front of the buffer
     ///
     /// Calling this with n > len() is allowed and will simply clear the buffer.
@@ -232,6 +273,24 @@ impl SendBuffer {
         self.retransmits.remove(0..self.fully_acked_offset());
     }
 
+    /// Discard buffered data at or beyond `offset`, abandoning it.
+    ///
+    /// Used to implement RESET_STREAM_AT: data past the reliable size is no longer (re)transmitted.
+    /// The new end offset is clamped to the fully-acknowledged offset, so already-sent-and-acked
+    /// data is retained. Bookkeeping for the discarded tail (unsent cursor, queued retransmits, and
+    /// acknowledged ranges) is dropped to keep the buffer consistent.
+    pub(super) fn truncate(&mut self, offset: u64) {
+        let old_end = self.offset();
+        self.data.truncate(offset);
+        let new_end = self.offset();
+        debug_assert!(new_end <= old_end);
+
+        // We no longer have the discarded data, so we must not try to send or track it.
+        self.unsent = self.unsent.min(new_end);
+        self.retransmits.remove(new_end..old_end);
+        self.acks.remove(new_end..old_end);
+    }
+
     /// Compute the next range to transmit on this stream and update state to account for that
     /// transmission.
     ///
@@ -324,7 +383,7 @@ impl SendBuffer {
     }
 
     /// Offset up to which all data has been acknowledged
-    fn fully_acked_offset(&self) -> u64 {
+    pub(super) fn fully_acked_offset(&self) -> u64 {
         self.data.range().start
     }
 
@@ -561,6 +620,94 @@ mod tests {
         assert!(buf.acks.is_empty());
     }
 
+    #[test]
+    fn truncate_basic() {
+        let mut buf = SendBuffer::new();
+        const MSG: &[u8] = b"Hello, world!"; // 13 bytes, coalesced into last_segment
+        buf.write(MSG);
+        assert_eq!(buf.offset(), 13);
+
+        buf.truncate(5);
+        assert_eq!(buf.offset(), 5);
+        assert_eq!(aggregate_unacked(&buf), &MSG[..5]);
+
+        // Truncating at or past the end is a no-op.
+        buf.truncate(100);
+        assert_eq!(buf.offset(), 5);
+        assert_eq!(aggregate_unacked(&buf), &MSG[..5]);
+
+        // poll_transmit never yields data beyond the truncation point.
+        assert_eq!(buf.poll_transmit(64), (0..5, true));
+        assert_eq!(buf.poll_transmit(64), (5..5, true));
+        assert!(!buf.has_unsent_data());
+    }
+
+    #[test]
+    fn truncate_multiple_segments() {
+        let mut buf = SendBuffer::new();
+        // Segments larger than MAX_COMBINE are stored as standalone segments.
+        buf.write(Bytes::from(vec![1u8; 2000]));
+        buf.write(Bytes::from(vec![2u8; 2000]));
+        assert_eq!(buf.offset(), 4000);
+
+        // Cut inside the second segment.
+        buf.truncate(3000);
+        assert_eq!(buf.offset(), 3000);
+        let data = aggregate_unacked(&buf);
+        assert_eq!(data.len(), 3000);
+        assert!(data[..2000].iter().all(|&x| x == 1));
+        assert!(data[2000..].iter().all(|&x| x == 2));
+
+        // Cut exactly at the boundary of the first segment.
+        buf.truncate(2000);
+        assert_eq!(buf.offset(), 2000);
+        let data = aggregate_unacked(&buf);
+        assert_eq!(data.len(), 2000);
+        assert!(data.iter().all(|&x| x == 1));
+    }
+
+    #[test]
+    fn truncate_after_ack_and_below_front() {
+        let mut buf = SendBuffer::new();
+        const MSG: &[u8] = b"abcdefghij"; // 10 bytes
+        buf.write(MSG);
+        assert_eq!(buf.poll_transmit(64), (0..10, true));
+        buf.ack(0..4); // drop the first 4 bytes from the front
+        assert_eq!(buf.fully_acked_offset(), 4);
+
+        // Truncate within the retained range.
+        buf.truncate(7);
+        assert_eq!(buf.offset(), 7);
+        assert_eq!(aggregate_unacked(&buf), b"efg");
+
+        // Truncating below the (advanced) front clears the buffer but keeps the front offset.
+        buf.truncate(2);
+        assert_eq!(buf.offset(), 4);
+        assert_eq!(buf.fully_acked_offset(), 4);
+        assert!(buf.is_fully_acked());
+        assert!(aggregate_unacked(&buf).is_empty());
+    }
+
+    #[test]
+    fn truncate_drops_retransmits_and_unacked() {
+        let mut buf = SendBuffer::new();
+        const MSG: &[u8] = b"Hello, world with extra data!"; // 29 bytes
+        buf.write(MSG);
+        assert_eq!(buf.poll_transmit(64), (0..29, true));
+        // Mark a tail range lost, queuing it for retransmission.
+        buf.retransmit(20..29);
+        assert_eq!(buf.unacked(), 29);
+        assert!(buf.has_unsent_data());
+
+        // Truncating away the lost tail must drop the queued retransmit and update accounting,
+        // otherwise a later poll_transmit would request data we no longer have and panic.
+        buf.truncate(15);
+        assert_eq!(buf.offset(), 15);
+        assert_eq!(buf.unacked(), 15);
+        assert!(!buf.has_unsent_data());
+        assert_eq!(buf.poll_transmit(64), (15..15, true));
+    }
+
     fn aggregate_unacked(buf: &SendBuffer) -> Vec<u8> {
         buf.data.to_vec()
     }
@@ -600,6 +747,8 @@ mod proptests {
         Retransmit(Range<u64>),
         // poll_transmit with the given max len
         PollTransmit(#[strategy(16usize..1024)] usize),
+        // truncate the buffer at a random offset within the sent range
+        Truncate(u64),
     }
 
     /// Map a range into a target range
@@ -654,6 +803,21 @@ mod proptests {
                     trace!("Op::Retransmit({:?})", range);
                     sb.retransmit(range);
                 }
+                Op::Truncate(raw) => {
+                    // Choose a truncation point within the data that has actually been sent.
+                    let target = if max_send_offset == 0 {
+                        0
+                    } else {
+                        raw % (max_send_offset + 1)
+                    };
+                    trace!("Op::Truncate({})", target);
+                    sb.truncate(target);
+                    // `truncate` clamps the new end to the fully-acked offset, so read back the
+                    // authoritative end and mirror it in the reference model.
+                    let new_end = sb.offset();
+                    buf.truncate(new_end as usize);
+                    max_send_offset = max_send_offset.min(new_end);
+                }
                 Op::PollTransmit(max_len) => {
                     trace!("Op::PollTransmit({})", max_len);
                     let (range, _partial) = sb.poll_transmit(max_len);