From cd69fa5260360756689919da84fa9ff6f7c1b6eb Mon Sep 17 00:00:00 2001
From: Benjamin Saunders <ben.e.saunders@gmail.com>
Date: Sat, 24 Feb 2024 15:29:29 -0800
Subject: [PATCH] Don't double-count buffer consumption in close length checks

PacketBuilder::max_size previously subtracted out the start index and
header size of the packet, and therefore described the admissible size
of the packet's frames. However, most of our logic operates in terms
of absolute buffer positions instead. This was confusing, and led to
erroneous double-counting of space use in close packets.
---
 quinn-proto/src/connection/mod.rs            | 16 ++++++----------
 quinn-proto/src/connection/packet_builder.rs |  6 +++++-
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs
index 79a47ac7f..57940a46f 100644
--- a/quinn-proto/src/connection/mod.rs
+++ b/quinn-proto/src/connection/mod.rs
@@ -759,17 +759,18 @@ impl Connection {
                     "ACKs should leave space for ConnectionClose"
                 );
                 if buf.len() + frame::ConnectionClose::SIZE_BOUND < builder.max_size {
+                    let max_frame_size = builder.max_size - buf.len();
                     match self.state {
                         State::Closed(state::Closed { ref reason }) => {
                             if space_id == SpaceId::Data || reason.is_transport_layer() {
-                                reason.encode(buf, builder.max_size)
+                                reason.encode(buf, max_frame_size)
                             } else {
                                 frame::ConnectionClose {
                                     error_code: TransportErrorCode::APPLICATION_ERROR,
                                     frame_type: None,
                                     reason: Bytes::new(),
                                 }
-                                .encode(buf, builder.max_size)
+                                .encode(buf, max_frame_size)
                             }
                         }
                         State::Draining => frame::ConnectionClose {
@@ -777,7 +778,7 @@ impl Connection {
                             frame_type: None,
                             reason: Bytes::new(),
                         }
-                        .encode(buf, builder.max_size),
+                        .encode(buf, max_frame_size),
                         _ => unreachable!(
                             "tried to make a close packet when the connection wasn't closed"
                         ),
@@ -829,13 +830,8 @@ impl Connection {
                 }
             }
 
-            let sent = self.populate_packet(
-                now,
-                space_id,
-                buf,
-                buf_capacity - builder.tag_len,
-                builder.exact_number,
-            );
+            let sent =
+                self.populate_packet(now, space_id, buf, builder.max_size, builder.exact_number);
 
             // ACK-only packets should only be sent when explicitly allowed. If we write them due
             // to any other reason, there is a bug which leads to one component announcing write
diff --git a/quinn-proto/src/connection/packet_builder.rs b/quinn-proto/src/connection/packet_builder.rs
index c53d09a2d..30a62387f 100644
--- a/quinn-proto/src/connection/packet_builder.rs
+++ b/quinn-proto/src/connection/packet_builder.rs
@@ -18,7 +18,11 @@ pub(super) struct PacketBuilder {
     pub(super) ack_eliciting: bool,
     pub(super) exact_number: u64,
     pub(super) short_header: bool,
+    /// Smallest absolute position in the associated buffer that must be occupied by this packet's
+    /// frames
     pub(super) min_size: usize,
+    /// Largest absolute position in the associated buffer that may be occupied by this packet's
+    /// frames
     pub(super) max_size: usize,
     pub(super) tag_len: usize,
     pub(super) span: tracing::Span,
@@ -147,7 +151,7 @@ impl PacketBuilder {
             buffer.len() + (sample_size + 4).saturating_sub(number.len() + tag_len),
             partial_encode.start + conn.rem_cids.active().len() + 6,
         );
-        let max_size = buffer_capacity - partial_encode.start - partial_encode.header_len - tag_len;
+        let max_size = buffer_capacity - tag_len;
 
         Some(Self {
             datagram_start,