From 33264a05c84564663b2629ed7a81a8fe9318e144 Mon Sep 17 00:00:00 2001
From: Lars Eggert <lars@eggert.org>
Date: Mon, 16 Sep 2024 16:28:36 +0300
Subject: [PATCH] fix: Resend more data during handshake

We don't track which packets are coalesced with others, so when we
detect a loss in one packet number space, we cannot resend any coalesced
packets in other packet number space. This PR tries to approximate this
behavior by scheduling un-ACKed Handshake and 0-RTT packets for RTX when
packets are lost in the Initial space.

Broken out of #1998
---
 neqo-transport/src/connection/mod.rs | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/neqo-transport/src/connection/mod.rs b/neqo-transport/src/connection/mod.rs
index 29bc85ff73..53264df694 100644
--- a/neqo-transport/src/connection/mod.rs
+++ b/neqo-transport/src/connection/mod.rs
@@ -1396,12 +1396,14 @@ impl Connection {
                 // that Initial packets were lost.
                 // Resend Initial CRYPTO frames immediately a few times just
                 // in case.  As we don't have an RTT estimate yet, this helps
-                // when there is a short RTT and losses.
+                // when there is a short RTT and losses. Also mark all 0-RTT
+                // data as lost.
                 if dcid.is_none()
                     && self.cid_manager.is_valid(packet.dcid())
                     && self.stats.borrow().saved_datagrams <= EXTRA_INITIALS
                 {
                     self.crypto.resend_unacked(PacketNumberSpace::Initial);
+                    self.resend_0rtt(now);
                 }
             }
             (PacketType::VersionNegotiation | PacketType::Retry | PacketType::OtherVersion, ..) => {
@@ -2849,8 +2851,13 @@ impl Connection {
                     self.handshake(now, packet_version, space, Some(&buf))?;
                     self.create_resumption_token(now);
                 } else {
-                    // If we get a useless CRYPTO frame send outstanding CRYPTO frames again.
+                    // If we get a useless CRYPTO frame send outstanding CRYPTO frames and 0-RTT
+                    // data again.
                     self.crypto.resend_unacked(space);
+                    if space == PacketNumberSpace::Initial {
+                        self.crypto.resend_unacked(PacketNumberSpace::Handshake);
+                        self.resend_0rtt(now);
+                    }
                 }
             }
             Frame::NewToken { token } => {
@@ -3062,21 +3069,22 @@ impl Connection {
         stats.largest_acknowledged = max(stats.largest_acknowledged, largest_acknowledged);
     }
 
+    /// Tell 0-RTT packets that they were "lost".
+    fn resend_0rtt(&mut self, now: Instant) {
+        if let Some(path) = self.paths.primary() {
+            let dropped = self.loss_recovery.drop_0rtt(&path, now);
+            self.handle_lost_packets(&dropped);
+        }
+    }
+
     /// When the server rejects 0-RTT we need to drop a bunch of stuff.
     fn client_0rtt_rejected(&mut self, now: Instant) {
         if !matches!(self.zero_rtt_state, ZeroRttState::Sending) {
             return;
         }
         qdebug!([self], "0-RTT rejected");
-
-        // Tell 0-RTT packets that they were "lost".
-        if let Some(path) = self.paths.primary() {
-            let dropped = self.loss_recovery.drop_0rtt(&path, now);
-            self.handle_lost_packets(&dropped);
-        }
-
+        self.resend_0rtt(now);
         self.streams.zero_rtt_rejected();
-
         self.crypto.states.discard_0rtt_keys();
         self.events.client_0rtt_rejected();
     }