Discard initial key right after a HANDSHAKE packet has been sent and unify the fixes for missing initial keys. This goal is to make sure an endpoint never tries to serialize initial frames (could be pending frames or re-serialize initial packet in the coalescer)
while initial key gets dropped.

Protected by FLAGS_quic_reloadable_flag_quic_fix_missing_initial_keys.

PiperOrigin-RevId: 333389633
Change-Id: I7126a9fcd96e0ac3b2a70bd695b9885c2c137e88
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index 2348191..bb7dd47 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -2503,7 +2503,8 @@
 
 void QuicConnection::NeuterUnencryptedPackets() {
   sent_packet_manager_.NeuterUnencryptedPackets();
-  if (GetQuicReloadableFlag(
+  if (!fix_missing_initial_keys_ &&
+      GetQuicReloadableFlag(
           quic_neuter_initial_packet_in_coalescer_with_initial_key_discarded) &&
       version().CanSendCoalescedPackets()) {
     QUIC_RELOADABLE_FLAG_COUNT(
@@ -2725,7 +2726,8 @@
           // Failed to flush coalesced packet, write error has been handled.
           return false;
         }
-        if (GetQuicReloadableFlag(
+        if (!fix_missing_initial_keys_ &&
+            GetQuicReloadableFlag(
                 quic_discard_initial_packet_with_key_dropped)) {
           QUIC_RELOADABLE_FLAG_COUNT(
               quic_discard_initial_packet_with_key_dropped);
@@ -2957,6 +2959,11 @@
           packet->nonretransmittable_frames, packet_send_time);
     }
   }
+  if (fix_missing_initial_keys_ &&
+      packet->encryption_level == ENCRYPTION_HANDSHAKE) {
+    QUIC_RELOADABLE_FLAG_COUNT_N(quic_fix_missing_initial_keys, 1, 2);
+    visitor_->OnHandshakePacketSent();
+  }
 
   if (in_flight || !retransmission_alarm_->IsSet()) {
     SetRetransmissionAlarm();
@@ -4628,6 +4635,14 @@
     QUIC_BUG_IF(coalesced_packet_.length() > 0);
     return true;
   }
+  if (fix_missing_initial_keys_) {
+    QUIC_RELOADABLE_FLAG_COUNT_N(quic_fix_missing_initial_keys, 2, 2);
+    if (!framer_.HasEncrypterOfEncryptionLevel(ENCRYPTION_INITIAL)) {
+      // Initial packet will be re-serialized. Neuter it in case initial key has
+      // been dropped.
+      coalesced_packet_.NeuterInitialPacket();
+    }
+  }
   if (coalesced_packet_.length() == 0) {
     return true;
   }
@@ -4650,7 +4665,8 @@
     if (debug_visitor_ != nullptr) {
       debug_visitor_->OnCoalescedPacketSent(coalesced_packet_, length);
     }
-    if (coalesced_packet_.ContainsPacketOfEncryptionLevel(
+    if (!fix_missing_initial_keys_ &&
+        coalesced_packet_.ContainsPacketOfEncryptionLevel(
             ENCRYPTION_HANDSHAKE)) {
       // This is only called in coalescer because all ENCRYPTION_HANDSHAKE
       // packets go through the coalescer.
@@ -4679,7 +4695,8 @@
   if (debug_visitor_ != nullptr) {
     debug_visitor_->OnCoalescedPacketSent(coalesced_packet_, length);
   }
-  if (coalesced_packet_.ContainsPacketOfEncryptionLevel(ENCRYPTION_HANDSHAKE)) {
+  if (!fix_missing_initial_keys_ &&
+      coalesced_packet_.ContainsPacketOfEncryptionLevel(ENCRYPTION_HANDSHAKE)) {
     // This is only called in coalescer because all ENCRYPTION_HANDSHAKE
     // packets go through the coalescer.
     visitor_->OnHandshakePacketSent();
diff --git a/quic/core/quic_connection.h b/quic/core/quic_connection.h
index de01c03..994a8df 100644
--- a/quic/core/quic_connection.h
+++ b/quic/core/quic_connection.h
@@ -1782,6 +1782,9 @@
                              GetQuicReloadableFlag(quic_send_path_response);
   // True if AckFrequencyFrame is supported.
   bool can_receive_ack_frequency_frame_ = false;
+
+  const bool fix_missing_initial_keys_ =
+      GetQuicReloadableFlag(quic_fix_missing_initial_keys);
 };
 
 }  // namespace quic
diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index 85b6224..17a7daa 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -11544,7 +11544,10 @@
        ReserializeInitialPacketInCoalescerAfterDiscardingInitialKey) {
   SetQuicReloadableFlag(
       quic_neuter_initial_packet_in_coalescer_with_initial_key_discarded, true);
-  if (!connection_.version().CanSendCoalescedPackets()) {
+  if (!connection_.version().CanSendCoalescedPackets() ||
+      !GetQuicReloadableFlag(quic_fix_missing_initial_keys)) {
+    // Cannot set quic_fix_missing_initial_keys in the test since connection_ is
+    // created since the setup.
     return;
   }
   use_tagging_decrypter();
@@ -11576,6 +11579,7 @@
   // crashes).
   EXPECT_CALL(visitor_, OnStreamFrame(_)).Times(1);
   ProcessDataPacketAtLevel(1000, false, ENCRYPTION_FORWARD_SECURE);
+  EXPECT_TRUE(connection_.connected());
 }
 
 // Check that if there are two PATH_CHALLENGE frames in the packet, the latter
@@ -12009,8 +12013,9 @@
 }
 
 // Regression test for b/168294218.
-TEST_P(QuicConnectionTest, InitialPacketCausedCoalescerToBeFlushed) {
-  if (!connection_.version().CanSendCoalescedPackets()) {
+TEST_P(QuicConnectionTest, CoalescerHandlesInitialKeyDiscard) {
+  if (!connection_.version().CanSendCoalescedPackets() ||
+      !GetQuicReloadableFlag(quic_fix_missing_initial_keys)) {
     return;
   }
   SetQuicReloadableFlag(quic_discard_initial_packet_with_key_dropped, true);
@@ -12037,8 +12042,7 @@
     // Verify this packet is on hold.
     EXPECT_EQ(0u, writer_->packets_write_attempts());
   }
-  // Verify the INITIAL ACK packet gets discarded.
-  EXPECT_EQ(1u, connection_.GetStats().packets_discarded);
+  EXPECT_TRUE(connection_.connected());
 }
 
 }  // namespace