fix: Include an ACK with a CONNECTION_CLOSE (#1854)

* fix: Include an ACK with a CONNECTION_CLOSE

Fixes #1161

* Address review comments

* Send ACK before CC if there is space

* Address code review

* Address more review comments

* Update neqo-transport/src/tracking.rs

Co-authored-by: Martin Thomson <mt@lowentropy.net>
Signed-off-by: Lars Eggert <lars@eggert.org>

* Update neqo-transport/src/tracking.rs

Co-authored-by: Martin Thomson <mt@lowentropy.net>
Signed-off-by: Lars Eggert <lars@eggert.org>

* Address code review

---------

Signed-off-by: Lars Eggert <lars@eggert.org>
Co-authored-by: Martin Thomson <mt@lowentropy.net>

neqo-transport/src/connection/mod.rs

@@ -56,7 +56,7 @@ use crate::{
         self, TransportParameter, TransportParameterId, TransportParameters,
         TransportParametersHandler,
     },
-    tracking::{AckTracker, PacketNumberSpace, SentPacket},
+    tracking::{AckTracker, PacketNumberSpace, RecvdPackets, SentPacket},
     version::{Version, WireVersion},
     AppError, ConnectionError, Error, Res, StreamId,
 };
@@ -2189,6 +2189,40 @@ impl Connection {
         (tokens, ack_eliciting, padded)
     }
 
+    fn write_closing_frames(
+        &mut self,
+        close: &ClosingFrame,
+        builder: &mut PacketBuilder,
+        space: PacketNumberSpace,
+        now: Instant,
+        path: &PathRef,
+        tokens: &mut Vec<RecoveryToken>,
+    ) {
+        if builder.remaining() > ClosingFrame::MIN_LENGTH + RecvdPackets::USEFUL_ACK_LEN {
+            // Include an ACK frame with the CONNECTION_CLOSE.
+            let limit = builder.limit();
+            builder.set_limit(limit - ClosingFrame::MIN_LENGTH);
+            self.acks.immediate_ack(now);
+            self.acks.write_frame(
+                space,
+                now,
+                path.borrow().rtt().estimate(),
+                builder,
+                tokens,
+                &mut self.stats.borrow_mut().frame_tx,
+            );
+            builder.set_limit(limit);
+        }
+        // ConnectionError::Application is only allowed at 1RTT.
+        let sanitized = if space == PacketNumberSpace::ApplicationData {
+            None
+        } else {
+            close.sanitize()
+        };
+        sanitized.as_ref().unwrap_or(close).write_frame(builder);
+        self.stats.borrow_mut().frame_tx.connection_close += 1;
+    }
+
     /// Build a datagram, possibly from multiple packets (for different PN
     /// spaces) and each containing 1+ frames.
     #[allow(clippy::too_many_lines)] // Yeah, that's just the way it is.
@@ -2252,17 +2286,7 @@ impl Connection {
             let payload_start = builder.len();
             let (mut tokens, mut ack_eliciting, mut padded) = (Vec::new(), false, false);
             if let Some(ref close) = closing_frame {
-                // ConnectionError::Application is only allowed at 1RTT.
-                let sanitized = if *space == PacketNumberSpace::ApplicationData {
-                    None
-                } else {
-                    close.sanitize()
-                };
-                sanitized
-                    .as_ref()
-                    .unwrap_or(close)
-                    .write_frame(&mut builder);
-                self.stats.borrow_mut().frame_tx.connection_close += 1;
+                self.write_closing_frames(close, &mut builder, *space, now, path, &mut tokens);
             } else {
                 (tokens, ack_eliciting, padded) =
                     self.write_frames(path, *space, &profile, &mut builder, now);

neqo-transport/src/connection/state.rs

@@ -156,10 +156,13 @@ impl ClosingFrame {
         }
     }
 
+    /// Length of a closing frame with a truncated `reason_length`. Allow 8 bytes for the reason
+    /// phrase to ensure that if it needs to be truncated there is still at least a few bytes of
+    /// the value.
+    pub const MIN_LENGTH: usize = 1 + 8 + 8 + 2 + 8;
+
     pub fn write_frame(&self, builder: &mut PacketBuilder) {
-        // Allow 8 bytes for the reason phrase to ensure that if it needs to be
-        // truncated there is still at least a few bytes of the value.
-        if builder.remaining() < 1 + 8 + 8 + 2 + 8 {
+        if builder.remaining() < ClosingFrame::MIN_LENGTH {
             return;
         }
         match &self.error {

neqo-transport/src/connection/tests/close.rs

@@ -40,7 +40,14 @@ fn connection_close() {
 
     client.close(now, 42, "");
 
+    let stats_before = client.stats().frame_tx;
     let out = client.process(None, now);
+    let stats_after = client.stats().frame_tx;
+    assert_eq!(
+        stats_after.connection_close,
+        stats_before.connection_close + 1
+    );
+    assert_eq!(stats_after.ack, stats_before.ack + 1);
 
     server.process_input(&out.dgram().unwrap(), now);
     assert_draining(&server, &Error::PeerApplicationError(42));
@@ -57,7 +64,14 @@ fn connection_close_with_long_reason_string() {
     let long_reason = String::from_utf8([0x61; 2048].to_vec()).unwrap();
     client.close(now, 42, long_reason);
 
+    let stats_before = client.stats().frame_tx;
     let out = client.process(None, now);
+    let stats_after = client.stats().frame_tx;
+    assert_eq!(
+        stats_after.connection_close,
+        stats_before.connection_close + 1
+    );
+    assert_eq!(stats_after.ack, stats_before.ack + 1);
 
     server.process_input(&out.dgram().unwrap(), now);
     assert_draining(&server, &Error::PeerApplicationError(42));

neqo-transport/src/tracking.rs

@@ -559,6 +559,10 @@ impl RecvdPackets {
         }
     }
 
+    /// Length of the worst possible ACK frame, assuming only one range and ECN counts.
+    /// Note that this assumes one byte for the type and count of extra ranges.
+    pub const USEFUL_ACK_LEN: usize = 1 + 8 + 8 + 1 + 8 + 3 * 8;
+
     /// Generate an ACK frame for this packet number space.
     ///
     /// Unlike other frame generators this doesn't modify the underlying instance
@@ -577,10 +581,6 @@ impl RecvdPackets {
         tokens: &mut Vec<RecoveryToken>,
         stats: &mut FrameStats,
     ) {
-        // The worst possible ACK frame, assuming only one range.
-        // Note that this assumes one byte for the type and count of extra ranges.
-        const LONGEST_ACK_HEADER: usize = 1 + 8 + 8 + 1 + 8;
-
         // Check that we aren't delaying ACKs.
         if !self.ack_now(now, rtt) {
             return;
@@ -592,7 +592,10 @@ impl RecvdPackets {
         // When congestion limited, ACK-only packets are 255 bytes at most
         // (`recovery::ACK_ONLY_SIZE_LIMIT - 1`).  This results in limiting the
         // ranges to 13 here.
-        let max_ranges = if let Some(avail) = builder.remaining().checked_sub(LONGEST_ACK_HEADER) {
+        let max_ranges = if let Some(avail) = builder
+            .remaining()
+            .checked_sub(RecvdPackets::USEFUL_ACK_LEN)
+        {
             // Apply a hard maximum to keep plenty of space for other stuff.
             min(1 + (avail / 16), MAX_ACKS_PER_FRAME)
         } else {
@@ -1158,7 +1161,9 @@ mod tests {
             .is_some());
 
         let mut builder = PacketBuilder::short(Encoder::new(), false, []);
-        builder.set_limit(32);
+        // The code pessimistically assumes that each range needs 16 bytes to express.
+        // So this won't be enough for a second range.
+        builder.set_limit(RecvdPackets::USEFUL_ACK_LEN + 8);
 
         let mut stats = FrameStats::default();
         tracker.write_frame(