Drop INITIAL packet if INITIAL key has been dropped because of coalesce packet gets flushed.

Protected by FLAGS_quic_reloadable_flag_quic_discard_initial_packet_with_key_dropped.

PiperOrigin-RevId: 331542696
Change-Id: I6c24f39f920135d573aa0c64aa5fa6c7acd50a04
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index b30e564..f658999 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -2713,6 +2713,18 @@
           // Failed to flush coalesced packet, write error has been handled.
           return false;
         }
+        if (GetQuicReloadableFlag(
+                quic_discard_initial_packet_with_key_dropped)) {
+          QUIC_RELOADABLE_FLAG_COUNT(
+              quic_discard_initial_packet_with_key_dropped);
+          if (packet->encryption_level == ENCRYPTION_INITIAL &&
+              !framer_.HasEncrypterOfEncryptionLevel(ENCRYPTION_INITIAL)) {
+            // Discard initial packet since flush of coalesce packet could
+            // cause initial keys to be dropped.
+            ++stats_.packets_discarded;
+            return true;
+          }
+        }
         if (!coalesced_packet_.MaybeCoalescePacket(
                 *packet, self_address(), send_to_address,
                 helper_->GetStreamSendBufferAllocator(),
diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index c5a4b9d..638c515 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -11945,6 +11945,39 @@
   }
 }
 
+// Regression test for b/168294218.
+TEST_P(QuicConnectionTest, InitialPacketCausedCoalescerToBeFlushed) {
+  if (!connection_.version().CanSendCoalescedPackets()) {
+    return;
+  }
+  SetQuicReloadableFlag(quic_discard_initial_packet_with_key_dropped, true);
+  // Verify only one HANDSHAKE packet gets sent.
+  EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, _, _, _)).Times(1);
+  EXPECT_CALL(visitor_, OnHandshakePacketSent()).WillOnce(Invoke([this]() {
+    connection_.RemoveEncrypter(ENCRYPTION_INITIAL);
+    connection_.NeuterUnencryptedPackets();
+  }));
+  EXPECT_CALL(visitor_, OnCryptoFrame(_)).Times(AnyNumber());
+
+  EXPECT_EQ(0u, connection_.GetStats().packets_discarded);
+  {
+    QuicConnection::ScopedPacketFlusher flusher(&connection_);
+    use_tagging_decrypter();
+    ProcessCryptoPacketAtLevel(1000, ENCRYPTION_INITIAL);
+    clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(1));
+    connection_.SetEncrypter(ENCRYPTION_INITIAL,
+                             std::make_unique<TaggingEncrypter>(0x01));
+    connection_.SetEncrypter(ENCRYPTION_HANDSHAKE,
+                             std::make_unique<TaggingEncrypter>(0x02));
+    connection_.SetDefaultEncryptionLevel(ENCRYPTION_HANDSHAKE);
+    connection_.SendCryptoDataWithString(std::string(1300, 'a'), 0);
+    // 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);
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic