Flush pending padding bytes before default_path_ is moved in connection migration.

Protected by FLAGS_quic_reloadable_flag_quic_flush_pending_frames_and_padding_bytes_on_migration.

PiperOrigin-RevId: 405888129
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index a93715d..aa598e4 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -5272,13 +5272,25 @@
         << "EffectivePeerMigration started without address change.";
     return;
   }
-  // There could be pending NEW_TOKEN_FRAME triggered by non-probing
-  // PATH_RESPONSE_FRAME in the same packet.
-  if (packet_creator_.HasPendingFrames()) {
+  if (GetQuicReloadableFlag(
+          quic_flush_pending_frames_and_padding_bytes_on_migration)) {
+    QUIC_RELOADABLE_FLAG_COUNT(
+        quic_flush_pending_frames_and_padding_bytes_on_migration);
+    // There could be pending NEW_TOKEN_FRAME triggered by non-probing
+    // PATH_RESPONSE_FRAME in the same packet or pending padding bytes in the
+    // packet creator.
     packet_creator_.FlushCurrentPacket();
+    packet_creator_.SendRemainingPendingPadding();
     if (!connected_) {
       return;
     }
+  } else {
+    if (packet_creator_.HasPendingFrames()) {
+      packet_creator_.FlushCurrentPacket();
+      if (!connected_) {
+        return;
+      }
+    }
   }
 
   // Action items:
diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index 09b4ebe..1b37830 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -512,6 +512,15 @@
     return false;
   }
 
+  void SendOrQueuePacket(SerializedPacket packet) override {
+    QuicConnection::SendOrQueuePacket(std::move(packet));
+    self_address_on_default_path_while_sending_packet_ = self_address();
+  }
+
+  QuicSocketAddress self_address_on_default_path_while_sending_packet() {
+    return self_address_on_default_path_while_sending_packet_;
+  }
+
   SimpleDataProducer* producer() { return &producer_; }
 
   using QuicConnection::active_effective_peer_migration_type;
@@ -538,6 +547,8 @@
   SimpleSessionNotifier* notifier_;
 
   std::unique_ptr<QuicSocketAddress> next_effective_peer_addr_;
+
+  QuicSocketAddress self_address_on_default_path_while_sending_packet_;
 };
 
 enum class AckResponse { kDefer, kImmediate };
@@ -1320,15 +1331,15 @@
   void set_perspective(Perspective perspective) {
     connection_.set_perspective(perspective);
     if (perspective == Perspective::IS_SERVER) {
+      QuicConfig config;
       if (!GetQuicReloadableFlag(
               quic_remove_connection_migration_connection_option)) {
-        QuicConfig config;
         QuicTagVector connection_options;
         connection_options.push_back(kRVCM);
         config.SetInitialReceivedConnectionOptions(connection_options);
-        EXPECT_CALL(*send_algorithm_, SetFromConfig(_, _));
-        connection_.SetFromConfig(config);
       }
+      EXPECT_CALL(*send_algorithm_, SetFromConfig(_, _));
+      connection_.SetFromConfig(config);
 
       connection_.set_can_truncate_connection_ids(true);
       QuicConnectionPeer::SetNegotiatedVersion(&connection_);
@@ -1509,8 +1520,7 @@
 };
 
 // Run all end to end tests with all supported versions.
-INSTANTIATE_TEST_SUITE_P(QuicConnectionTests,
-                         QuicConnectionTest,
+INSTANTIATE_TEST_SUITE_P(QuicConnectionTests, QuicConnectionTest,
                          ::testing::ValuesIn(GetTestParams()),
                          ::testing::PrintToStringParamName());
 
@@ -2067,6 +2077,58 @@
   }
 }
 
+// Regression test for b/200020764.
+TEST_P(QuicConnectionTest, ConnetcionMigrationWithPendingPaddingBytes) {
+  // TODO(haoyuewang) Move these test setup code to a common member function.
+  set_perspective(Perspective::IS_SERVER);
+  if (!connection_.connection_migration_use_new_cid()) {
+    return;
+  }
+  QuicPacketCreatorPeer::SetSendVersionInPacket(creator_, false);
+  connection_.CreateConnectionIdManager();
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_FORWARD_SECURE);
+  QuicConnectionPeer::SetPeerAddress(&connection_, kPeerAddress);
+  QuicConnectionPeer::SetEffectivePeerAddress(&connection_, kPeerAddress);
+  QuicConnectionPeer::SetAddressValidated(&connection_);
+
+  // Sends new server CID to client.
+  QuicConnectionId new_cid;
+  EXPECT_CALL(visitor_, OnServerConnectionIdIssued(_))
+      .WillOnce(Invoke([&](const QuicConnectionId& cid) { new_cid = cid; }));
+  EXPECT_CALL(visitor_, SendNewConnectionId(_));
+  // Discard INITIAL key.
+  connection_.RemoveEncrypter(ENCRYPTION_INITIAL);
+  connection_.NeuterUnencryptedPackets();
+  connection_.OnHandshakeComplete();
+  EXPECT_CALL(visitor_, GetHandshakeState())
+      .WillRepeatedly(Return(HANDSHAKE_CONFIRMED));
+
+  auto* packet_creator = QuicConnectionPeer::GetPacketCreator(&connection_);
+  packet_creator->FlushCurrentPacket();
+  packet_creator->AddPendingPadding(50u);
+  const QuicSocketAddress kPeerAddress3 =
+      QuicSocketAddress(QuicIpAddress::Loopback6(), /*port=*/56789);
+  auto ack_frame = InitAckFrame(1);
+  EXPECT_CALL(*send_algorithm_, OnCongestionEvent(_, _, _, _, _));
+  EXPECT_CALL(visitor_, OnConnectionMigration(PORT_CHANGE)).Times(1);
+  ProcessFramesPacketWithAddresses({QuicFrame(&ack_frame)}, kSelfAddress,
+                                   kPeerAddress3, ENCRYPTION_FORWARD_SECURE);
+  if (GetQuicReloadableFlag(
+          quic_flush_pending_frames_and_padding_bytes_on_migration)) {
+    // Any pending frames/padding should be flushed before default_path_ is
+    // temporarily reset.
+    ASSERT_EQ(connection_.self_address_on_default_path_while_sending_packet()
+                  .host()
+                  .address_family(),
+              IpAddressFamily::IP_V6);
+  } else {
+    ASSERT_EQ(connection_.self_address_on_default_path_while_sending_packet()
+                  .host()
+                  .address_family(),
+              IpAddressFamily::IP_UNSPEC);
+  }
+}
+
 // Regression test for b/196208556.
 TEST_P(QuicConnectionTest,
        ReversePathValidationResponseReceivedFromUnexpectedPeerAddress) {
@@ -10595,10 +10657,8 @@
 }
 
 void QuicConnectionTest::TestClientRetryHandling(
-    bool invalid_retry_tag,
-    bool missing_original_id_in_config,
-    bool wrong_original_id_in_config,
-    bool missing_retry_id_in_config,
+    bool invalid_retry_tag, bool missing_original_id_in_config,
+    bool wrong_original_id_in_config, bool missing_retry_id_in_config,
     bool wrong_retry_id_in_config) {
   if (invalid_retry_tag) {
     ASSERT_FALSE(missing_original_id_in_config);
diff --git a/quic/core/quic_flags_list.h b/quic/core/quic_flags_list.h
index 509bae5..77d96ec 100644
--- a/quic/core/quic_flags_list.h
+++ b/quic/core/quic_flags_list.h
@@ -75,6 +75,8 @@
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_drop_unsent_path_response, true)
 // If true, enable server retransmittable on wire PING.
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_enable_server_on_wire_ping, true)
+// If true, flush pending frames as well as pending padding bytes on connection migration.
+QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_flush_pending_frames_and_padding_bytes_on_migration, false)
 // If true, ietf connection migration is no longer conditioned on connection option RVCM.
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_remove_connection_migration_connection_option, true)
 // If true, ignore peer_max_ack_delay during handshake.