Neuter initial packet in the coalescer when discarding initial keys.

PiperOrigin-RevId: 328808677
Change-Id: I34a3dbc33c52146fb88d65c60cdfc5ca8bcf1cf7
diff --git a/quic/core/quic_coalesced_packet.cc b/quic/core/quic_coalesced_packet.cc
index b92cf75..e266a76 100644
--- a/quic/core/quic_coalesced_packet.cc
+++ b/quic/core/quic_coalesced_packet.cc
@@ -97,6 +97,26 @@
   initial_packet_ = nullptr;
 }
 
+void QuicCoalescedPacket::NeuterInitialPacket() {
+  if (initial_packet_ == nullptr) {
+    return;
+  }
+  if (length_ < initial_packet_->encrypted_length) {
+    QUIC_BUG << "length_: " << length_
+             << ", is less than initial packet length: "
+             << initial_packet_->encrypted_length;
+    Clear();
+    return;
+  }
+  length_ -= initial_packet_->encrypted_length;
+  if (length_ == 0) {
+    Clear();
+    return;
+  }
+  transmission_types_[ENCRYPTION_INITIAL] = NOT_RETRANSMISSION;
+  initial_packet_ = nullptr;
+}
+
 bool QuicCoalescedPacket::CopyEncryptedBuffers(char* buffer,
                                                size_t buffer_len,
                                                size_t* length_copied) const {
diff --git a/quic/core/quic_coalesced_packet.h b/quic/core/quic_coalesced_packet.h
index 6cf712d..64275f6 100644
--- a/quic/core/quic_coalesced_packet.h
+++ b/quic/core/quic_coalesced_packet.h
@@ -27,6 +27,9 @@
   // Clears this coalesced packet.
   void Clear();
 
+  // Clears all state associated with initial_packet_.
+  void NeuterInitialPacket();
+
   // Copies encrypted_buffers_ to |buffer| and sets |length_copied| to the
   // copied amount. Returns false if copy fails (i.e., |buffer_len| is not
   // enough).
diff --git a/quic/core/quic_coalesced_packet_test.cc b/quic/core/quic_coalesced_packet_test.cc
index 81e6b27..dd1db44 100644
--- a/quic/core/quic_coalesced_packet_test.cc
+++ b/quic/core/quic_coalesced_packet_test.cc
@@ -125,6 +125,84 @@
                                               length_copied, expected, 1000);
 }
 
+TEST(QuicCoalescedPacketTest, NeuterInitialPacket) {
+  QuicCoalescedPacket coalesced;
+  EXPECT_EQ("total_length: 0 padding_size: 0 packets: {}",
+            coalesced.ToString(0));
+  // Noop when neutering initial packet on a empty coalescer.
+  coalesced.NeuterInitialPacket();
+  EXPECT_EQ("total_length: 0 padding_size: 0 packets: {}",
+            coalesced.ToString(0));
+
+  SimpleBufferAllocator allocator;
+  EXPECT_EQ(0u, coalesced.length());
+  char buffer[1000];
+  QuicSocketAddress self_address(QuicIpAddress::Loopback4(), 1);
+  QuicSocketAddress peer_address(QuicIpAddress::Loopback4(), 2);
+  SerializedPacket packet1(QuicPacketNumber(1), PACKET_4BYTE_PACKET_NUMBER,
+                           buffer, 500, false, false);
+  packet1.transmission_type = PTO_RETRANSMISSION;
+  QuicAckFrame ack_frame(InitAckFrame(1));
+  packet1.nonretransmittable_frames.push_back(QuicFrame(&ack_frame));
+  packet1.retransmittable_frames.push_back(
+      QuicFrame(QuicStreamFrame(1, true, 0, 100)));
+  ASSERT_TRUE(coalesced.MaybeCoalescePacket(packet1, self_address, peer_address,
+                                            &allocator, 1500));
+  EXPECT_EQ(PTO_RETRANSMISSION,
+            coalesced.TransmissionTypeOfPacket(ENCRYPTION_INITIAL));
+  EXPECT_EQ(1500u, coalesced.max_packet_length());
+  EXPECT_EQ(500u, coalesced.length());
+  EXPECT_EQ(
+      "total_length: 1500 padding_size: 1000 packets: {ENCRYPTION_INITIAL}",
+      coalesced.ToString(1500));
+  // Neuter initial packet.
+  coalesced.NeuterInitialPacket();
+  EXPECT_EQ(0u, coalesced.max_packet_length());
+  EXPECT_EQ(0u, coalesced.length());
+  EXPECT_EQ("total_length: 0 padding_size: 0 packets: {}",
+            coalesced.ToString(0));
+
+  // Coalesce initial packet again.
+  ASSERT_TRUE(coalesced.MaybeCoalescePacket(packet1, self_address, peer_address,
+                                            &allocator, 1500));
+
+  SerializedPacket packet2(QuicPacketNumber(3), PACKET_4BYTE_PACKET_NUMBER,
+                           buffer, 500, false, false);
+  packet2.nonretransmittable_frames.push_back(QuicFrame(QuicPaddingFrame(100)));
+  packet2.encryption_level = ENCRYPTION_ZERO_RTT;
+  packet2.transmission_type = LOSS_RETRANSMISSION;
+  ASSERT_TRUE(coalesced.MaybeCoalescePacket(packet2, self_address, peer_address,
+                                            &allocator, 1500));
+  EXPECT_EQ(1500u, coalesced.max_packet_length());
+  EXPECT_EQ(1000u, coalesced.length());
+  EXPECT_EQ(LOSS_RETRANSMISSION,
+            coalesced.TransmissionTypeOfPacket(ENCRYPTION_ZERO_RTT));
+  EXPECT_EQ(
+      "total_length: 1500 padding_size: 500 packets: {ENCRYPTION_INITIAL, "
+      "ENCRYPTION_ZERO_RTT}",
+      coalesced.ToString(1500));
+
+  // Neuter initial packet.
+  coalesced.NeuterInitialPacket();
+  EXPECT_EQ(1500u, coalesced.max_packet_length());
+  EXPECT_EQ(500u, coalesced.length());
+  EXPECT_EQ(
+      "total_length: 1500 padding_size: 1000 packets: {ENCRYPTION_ZERO_RTT}",
+      coalesced.ToString(1500));
+
+  SerializedPacket packet3(QuicPacketNumber(5), PACKET_4BYTE_PACKET_NUMBER,
+                           buffer, 501, false, false);
+  packet3.encryption_level = ENCRYPTION_FORWARD_SECURE;
+  EXPECT_TRUE(coalesced.MaybeCoalescePacket(packet3, self_address, peer_address,
+                                            &allocator, 1500));
+  EXPECT_EQ(1500u, coalesced.max_packet_length());
+  EXPECT_EQ(1001u, coalesced.length());
+  // Neuter initial packet.
+  coalesced.NeuterInitialPacket();
+  EXPECT_EQ(1500u, coalesced.max_packet_length());
+  EXPECT_EQ(1001u, coalesced.length());
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index 7aff953..5d32a5f 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -2462,6 +2462,13 @@
 
 void QuicConnection::NeuterUnencryptedPackets() {
   sent_packet_manager_.NeuterUnencryptedPackets();
+  if (GetQuicReloadableFlag(
+          quic_neuter_initial_packet_in_coalescer_with_initial_key_discarded) &&
+      version().CanSendCoalescedPackets()) {
+    QUIC_RELOADABLE_FLAG_COUNT(
+        quic_neuter_initial_packet_in_coalescer_with_initial_key_discarded);
+    coalesced_packet_.NeuterInitialPacket();
+  }
   // This may have changed the retransmission timer, so re-arm it.
   SetRetransmissionAlarm();
   if (default_enable_5rto_blackhole_detection_) {
diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index 1b8c0c6..d33cb35 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -11812,6 +11812,45 @@
   EXPECT_EQ(clock_.Now() + timeout, connection_.GetTimeoutAlarm()->deadline());
 }
 
+// Regression test for b/166255274
+TEST_P(QuicConnectionTest,
+       ReserializeInitialPacketInCoalescerAfterDiscardingInitialKey) {
+  SetQuicReloadableFlag(
+      quic_neuter_initial_packet_in_coalescer_with_initial_key_discarded, true);
+  if (!connection_.version().CanSendCoalescedPackets()) {
+    return;
+  }
+  use_tagging_decrypter();
+  connection_.SetEncrypter(ENCRYPTION_INITIAL,
+                           std::make_unique<TaggingEncrypter>(0x01));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_INITIAL);
+  EXPECT_CALL(visitor_, OnCryptoFrame(_)).Times(1);
+  ProcessCryptoPacketAtLevel(1, ENCRYPTION_INITIAL);
+  EXPECT_TRUE(connection_.HasPendingAcks());
+  connection_.SetEncrypter(ENCRYPTION_HANDSHAKE,
+                           std::make_unique<TaggingEncrypter>(0x02));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_HANDSHAKE);
+  {
+    QuicConnection::ScopedPacketFlusher flusher(&connection_);
+    connection_.GetAckAlarm()->Fire();
+    // Verify this ACK packet is on hold.
+    EXPECT_EQ(0u, writer_->packets_write_attempts());
+
+    // Discard INITIAL key while there is an INITIAL packet in the coalescer.
+    connection_.RemoveEncrypter(ENCRYPTION_INITIAL);
+    connection_.NeuterUnencryptedPackets();
+  }
+  // If not setting
+  // quic_neuter_initial_packet_in_coalescer_with_initial_key_discarded, there
+  // will be pending frames in the creator.
+  EXPECT_FALSE(connection_.packet_creator().HasPendingFrames());
+  // The ACK frame is deleted along with initial_packet_ in coalescer. Sending
+  // connection close would cause this (released) ACK frame be serialized (and
+  // crashes).
+  EXPECT_CALL(visitor_, OnStreamFrame(_)).Times(1);
+  ProcessDataPacketAtLevel(1000, false, ENCRYPTION_FORWARD_SECURE);
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic
diff --git a/quic/core/quic_packet_creator.cc b/quic/core/quic_packet_creator.cc
index f842538..5140df8 100644
--- a/quic/core/quic_packet_creator.cc
+++ b/quic/core/quic_packet_creator.cc
@@ -549,6 +549,7 @@
     }
   }
   SerializePacket(QuicOwnedPacketBuffer(buffer, nullptr), buffer_len);
+  // TODO(b/166255274): report unrecoverable error on serialization failures.
   const size_t encrypted_length = packet_.encrypted_length;
   // Clear frames in packet_. No need to DeleteFrames since frames are owned by
   // initial_packet.
@@ -754,6 +755,10 @@
 
 void QuicPacketCreator::SerializePacket(QuicOwnedPacketBuffer encrypted_buffer,
                                         size_t encrypted_buffer_len) {
+  const bool use_queued_frames_cleaner = GetQuicReloadableFlag(
+      quic_neuter_initial_packet_in_coalescer_with_initial_key_discarded);
+  ScopedQueuedFramesCleaner cleaner(use_queued_frames_cleaner ? this : nullptr);
+
   DCHECK_LT(0u, encrypted_buffer_len);
   QUIC_BUG_IF(queued_frames_.empty() && pending_padding_bytes_ == 0)
       << "Attempt to serialize empty packet";
@@ -826,7 +831,9 @@
   }
 
   packet_size_ = 0;
-  queued_frames_.clear();
+  if (!use_queued_frames_cleaner) {
+    queued_frames_.clear();
+  }
   packet_.encrypted_buffer = encrypted_buffer.buffer;
   packet_.encrypted_length = encrypted_length;
 
@@ -2011,5 +2018,16 @@
   return latched_hard_max_packet_length_ != 0;
 }
 
+QuicPacketCreator::ScopedQueuedFramesCleaner::ScopedQueuedFramesCleaner(
+    QuicPacketCreator* creator)
+    : creator_(creator) {}
+
+QuicPacketCreator::ScopedQueuedFramesCleaner::~ScopedQueuedFramesCleaner() {
+  if (creator_ == nullptr) {
+    return;
+  }
+  creator_->queued_frames_.clear();
+}
+
 #undef ENDPOINT  // undef for jumbo builds
 }  // namespace quic
diff --git a/quic/core/quic_packet_creator.h b/quic/core/quic_packet_creator.h
index a225793..aefa975 100644
--- a/quic/core/quic_packet_creator.h
+++ b/quic/core/quic_packet_creator.h
@@ -466,6 +466,16 @@
  private:
   friend class test::QuicPacketCreatorPeer;
 
+  // Used to clear queued_frames_ of creator upon exiting the scope.
+  class QUIC_EXPORT_PRIVATE ScopedQueuedFramesCleaner {
+   public:
+    explicit ScopedQueuedFramesCleaner(QuicPacketCreator* creator);
+    ~ScopedQueuedFramesCleaner();
+
+   private:
+    QuicPacketCreator* creator_;  // Unowned.
+  };
+
   // Creates a stream frame which fits into the current open packet. If
   // |data_size| is 0 and fin is true, the expected behavior is to consume
   // the fin.