From a2b673d9424f84514b0bf4f9a8cdffb4feffc57d Mon Sep 17 00:00:00 2001
From: Benjamin Saunders <ben.e.saunders@gmail.com>
Date: Sun, 18 Feb 2024 18:30:46 -0800
Subject: [PATCH] Limit memory used to track sent ACKs when purely receiving

---
 quinn-proto/src/connection/packet_builder.rs |  2 +-
 quinn-proto/src/connection/spaces.rs         | 52 +++++++++++++++++++-
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/quinn-proto/src/connection/packet_builder.rs b/quinn-proto/src/connection/packet_builder.rs
index c53d09a2db..81191cfbe0 100644
--- a/quinn-proto/src/connection/packet_builder.rs
+++ b/quinn-proto/src/connection/packet_builder.rs
@@ -200,7 +200,7 @@ impl PacketBuilder {
         };
 
         conn.in_flight.insert(&packet);
-        conn.spaces[space_id].sent(exact_number, packet);
+        conn.in_flight.bytes -= conn.spaces[space_id].sent(exact_number, packet);
         conn.stats.path.sent_packets += 1;
         conn.reset_keep_alive(now);
         if size != 0 {
diff --git a/quinn-proto/src/connection/spaces.rs b/quinn-proto/src/connection/spaces.rs
index 5cfc76b287..1b3d459691 100644
--- a/quinn-proto/src/connection/spaces.rs
+++ b/quinn-proto/src/connection/spaces.rs
@@ -2,7 +2,7 @@ use std::{
     cmp,
     collections::{BTreeMap, VecDeque},
     mem,
-    ops::{Index, IndexMut},
+    ops::{Bound, Index, IndexMut},
     time::{Duration, Instant},
 };
 
@@ -33,6 +33,10 @@ pub(super) struct PacketSpace {
     /// The largest packet number the remote peer acknowledged in an ACK frame.
     pub(super) largest_acked_packet: Option<u64>,
     pub(super) largest_acked_packet_sent: Instant,
+    /// The highest-numbered ACK-eliciting packet we've sent
+    pub(super) largest_ack_eliciting_sent: u64,
+    /// Number of packets in `sent_packets` with numbers above `largest_ack_eliciting_sent`
+    pub(super) unacked_non_ack_eliciting_tail: u64,
     /// Transmitted but not acked
     // We use a BTreeMap here so we can efficiently query by range on ACK and for loss detection
     pub(super) sent_packets: BTreeMap<u64, SentPacket>,
@@ -80,6 +84,8 @@ impl PacketSpace {
             next_packet_number: 0,
             largest_acked_packet: None,
             largest_acked_packet_sent: now,
+            largest_ack_eliciting_sent: 0,
+            unacked_non_ack_eliciting_tail: 0,
             sent_packets: BTreeMap::new(),
             ecn_counters: frame::EcnCounts::ZERO,
             ecn_feedback: frame::EcnCounts::ZERO,
@@ -203,12 +209,54 @@ impl PacketSpace {
     pub(super) fn take(&mut self, number: u64) -> Option<SentPacket> {
         let packet = self.sent_packets.remove(&number)?;
         self.in_flight -= u64::from(packet.size);
+        if !packet.ack_eliciting && number > self.largest_ack_eliciting_sent {
+            self.unacked_non_ack_eliciting_tail =
+                self.unacked_non_ack_eliciting_tail.checked_sub(1).unwrap();
+        }
         Some(packet)
     }
 
-    pub(super) fn sent(&mut self, number: u64, packet: SentPacket) {
+    /// Returns the number of bytes to *remove* from the connection's in-flight count
+    pub(super) fn sent(&mut self, number: u64, packet: SentPacket) -> u64 {
+        // Retain state for at most this many non-ACK-eliciting packets sent after the most recently
+        // sent ACK-eliciting packet. We're never guaranteed to receive an ACK for those, and we
+        // can't judge them as lost without an ACK, so to limit memory in applications which receive
+        // packets but don't send ACK-eliciting data for long periods use we must eventually start
+        // forgetting about them, although it might also be reasonable to just kill the connection
+        // due to weird peer behavior.
+        const MAX_UNACKED_NON_ACK_ELICTING_TAIL: u64 = 1_000;
+
+        let mut forgotten_bytes = 0;
+        if packet.ack_eliciting {
+            self.unacked_non_ack_eliciting_tail = 0;
+            self.largest_ack_eliciting_sent = number;
+        } else if self.unacked_non_ack_eliciting_tail > MAX_UNACKED_NON_ACK_ELICTING_TAIL {
+            let oldest_after_ack_eliciting = *self
+                .sent_packets
+                .range((
+                    Bound::Excluded(self.largest_ack_eliciting_sent),
+                    Bound::Unbounded,
+                ))
+                .next()
+                .unwrap()
+                .0;
+            // Per https://www.rfc-editor.org/rfc/rfc9000.html#name-frames-and-frame-types,
+            // non-ACK-eliciting packets must only contain PADDING, ACK, and CONNECTION_CLOSE
+            // frames, which require no special handling on ACK or loss beyond removal from
+            // in-flight counters if padded.
+            let packet = self
+                .sent_packets
+                .remove(&oldest_after_ack_eliciting)
+                .unwrap();
+            forgotten_bytes = u64::from(packet.size);
+            self.in_flight -= forgotten_bytes;
+        } else {
+            self.unacked_non_ack_eliciting_tail += 1;
+        }
+
         self.in_flight += u64::from(packet.size);
         self.sent_packets.insert(number, packet);
+        forgotten_bytes
     }
 }