fix: Drop of Initial space only after sending 1st Handshake

This is what RFC9000 says, and it helps with servers that expect ACKs in
the Initial packet number space to happen. It also eliminates a
difference in behavior we had when a client received an Initial packet
containing an ACK and a CRYPTO frame vs. a CRYPTO frame and an ACK,
where that CRYPTO frame caused the Initial packet number space to be
dropped.

Broken out of mozilla#1998

neqo-transport/src/connection/mod.rs

@@ -2401,13 +2401,15 @@ impl Connection {
                 self.loss_recovery.on_packet_sent(path, sent);
             }
 
-            if *space == PacketNumberSpace::Handshake
-                && self.role == Role::Server
-                && self.state == State::Confirmed
-            {
-                // We could discard handshake keys in set_state,
-                // but wait until after sending an ACK.
-                self.discard_keys(PacketNumberSpace::Handshake, now);
+            if *space == PacketNumberSpace::Handshake {
+                if self.role == Role::Client {
+                    // We're sending a Handshake packet, so we can discard Initial keys.
+                    self.discard_keys(PacketNumberSpace::Initial, now);
+                } else if self.role == Role::Server && self.state == State::Confirmed {
+                    // We could discard handshake keys in set_state,
+                    // but wait until after sending an ACK.
+                    self.discard_keys(PacketNumberSpace::Handshake, now);
+                }
             }
         }
 
@@ -2758,11 +2760,6 @@ impl Connection {
                 self.set_initial_limits();
             }
             if self.crypto.install_keys(self.role)? {
-                if self.role == Role::Client {
-                    // We won't acknowledge Initial packets as a result of this, but the
-                    // server can rely on implicit acknowledgment.
-                    self.discard_keys(PacketNumberSpace::Initial, now);
-                }
                 self.saved_datagrams.make_available(CryptoSpace::Handshake);
             }
         }

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

@@ -591,8 +591,9 @@ fn reorder_1rtt() {
     now += RTT / 2;
     let s2 = server.process(c2.as_ref(), now).dgram();
     // The server has now received those packets, and saved them.
-    // The two additional are a Handshake and a 1-RTT (w/ NEW_CONNECTION_ID).
-    assert_eq!(server.stats().packets_rx, PACKETS * 2 + 4);
+    // The two additional are an Initial w/ACK, a Handshake w/ACK and a 1-RTT (w/
+    // NEW_CONNECTION_ID).
+    assert_eq!(server.stats().packets_rx, PACKETS * 2 + 5);
     assert_eq!(server.stats().saved_datagrams, PACKETS);
     assert_eq!(server.stats().dropped_rx, 1);
     assert_eq!(*server.state(), State::Confirmed);
@@ -802,9 +803,9 @@ fn anti_amplification() {
     let ack = client.process(Some(&s_init3), now).dgram().unwrap();
     assert!(!maybe_authenticate(&mut client)); // No need yet.
 
-    // The client sends a padded datagram, with just ACK for Handshake.
-    assert_eq!(client.stats().frame_tx.ack, ack_count + 1);
-    assert_eq!(client.stats().frame_tx.all(), frame_count + 1);
+    // The client sends a padded datagram, with just ACKs for Initial and Handshake.
+    assert_eq!(client.stats().frame_tx.ack, ack_count + 2);
+    assert_eq!(client.stats().frame_tx.all(), frame_count + 2);
     assert_ne!(ack.len(), client.plpmtu()); // Not padded (it includes Handshake).
 
     now += DEFAULT_RTT / 2;
@@ -1051,25 +1052,24 @@ fn only_server_initial() {
     let (initial, handshake) = split_datagram(&server_dgram1.unwrap());
     assert!(handshake.is_some());
 
-    // The client will not acknowledge the Initial as it discards keys.
-    // It sends a Handshake probe instead, containing just a PING frame.
-    assert_eq!(client.stats().frame_tx.ping, 0);
+    // The client sends an Initial ACK.
+    assert_eq!(client.stats().frame_tx.ack, 0);
     let probe = client.process(Some(&initial), now).dgram();
-    assertions::assert_handshake(&probe.unwrap());
+    assertions::assert_initial(&probe.unwrap(), false);
     assert_eq!(client.stats().dropped_rx, 0);
-    assert_eq!(client.stats().frame_tx.ping, 1);
+    assert_eq!(client.stats().frame_tx.ack, 1);
 
     let (initial, handshake) = split_datagram(&server_dgram2.unwrap());
     assert!(handshake.is_some());
 
-    // The same happens after a PTO, even though the client will discard the Initial packet.
+    // The same happens after a PTO.
     now += AT_LEAST_PTO;
-    assert_eq!(client.stats().frame_tx.ping, 1);
+    assert_eq!(client.stats().frame_tx.ack, 1);
     let discarded = client.stats().dropped_rx;
     let probe = client.process(Some(&initial), now).dgram();
-    assertions::assert_handshake(&probe.unwrap());
-    assert_eq!(client.stats().frame_tx.ping, 2);
-    assert_eq!(client.stats().dropped_rx, discarded + 1);
+    assertions::assert_initial(&probe.unwrap(), false);
+    assert_eq!(client.stats().frame_tx.ack, 2);
+    assert_eq!(client.stats().dropped_rx, discarded);
 
     // Pass the Handshake packet and complete the handshake.
     client.process_input(&handshake.unwrap(), now);
@@ -1136,7 +1136,9 @@ fn implicit_rtt_server() {
     let dgram = server.process(dgram.as_ref(), now).dgram();
     now += RTT / 2;
     let dgram = client.process(dgram.as_ref(), now).dgram();
-    assertions::assert_handshake(dgram.as_ref().unwrap());
+    let (initial, handshake) = split_datagram(dgram.as_ref().unwrap());
+    assertions::assert_initial(&initial, false);
+    assertions::assert_handshake(handshake.as_ref().unwrap());
     now += RTT / 2;
     server.process_input(&dgram.unwrap(), now);
 

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

@@ -224,7 +224,9 @@ fn pto_handshake_complete() {
 
     now += HALF_RTT;
     let pkt = client.process(pkt.as_ref(), now).dgram();
-    assert_handshake(pkt.as_ref().unwrap());
+    let (initial, handshake) = split_datagram(&pkt.clone().unwrap());
+    assert_initial(&initial, false);
+    assert_handshake(handshake.as_ref().unwrap());
 
     let cb = client.process(None, now).callback();
     // The client now has a single RTT estimate (20ms), so