In IETF quic, start peer migration by switching peer address as soon as detecting a non-probing frame in the incoming packet rather than after finishing parsing the whole packet. Protected by --gfe2_reloadable_flag_quic_start_peer_migration_earlier.

Doing so will fix the bug in IETF QUIC that a server may send in-cache response to old address while the end of current packet processing isn't reached

PiperOrigin-RevId: 328429809
Change-Id: Ib07274bdef3944240d43defc4183ac03cdd47bba

quic/core/quic_connection.cc

diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index f0769fd..c44e5cf 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -1775,7 +1775,8 @@
   }
   // Server starts to migrate connection upon receiving of non-probing packet
   // from a new peer address.
-  if (last_header_.packet_number == GetLargestReceivedPacket()) {
+  if (!start_peer_migration_earlier_ &&
+      last_header_.packet_number == GetLargestReceivedPacket()) {
     direct_peer_address_ = last_packet_source_address_;
     if (current_effective_peer_migration_type_ != NO_CHANGE) {
       // TODO(fayang): When multiple packet number spaces is supported, only
@@ -4125,14 +4126,14 @@
 }
 
 void QuicConnection::UpdatePacketContent(QuicFrameType type) {
+  if (version().HasIetfQuicFrames()) {
+    MaybeStartIetfPeerMigration(type);
+    return;
+  }
   // Packet content is tracked to identify connectivity probe in non-IETF
   // version, where a connectivity probe is defined as
   // - a padded PING packet with peer address change received by server,
   // - a padded PING packet on new path received by client.
-  if (version().HasIetfQuicFrames()) {
-    // TODO(b/149315151) detect IETF non-probing packet.
-    return;
-  }
 
   if (current_packet_content_ == NOT_PADDED_PING) {
     // We have already learned the current packet is not a connectivity
@@ -4184,6 +4185,24 @@
     if (current_effective_peer_migration_type_ != NO_CHANGE) {
       // Start effective peer migration immediately when the current packet is
       // confirmed not a connectivity probing packet.
+      StartEffectivePeerMigration(current_effective_peer_migration_type_);
+    }
+  }
+  current_effective_peer_migration_type_ = NO_CHANGE;
+}
+
+void QuicConnection::MaybeStartIetfPeerMigration(QuicFrameType type) {
+  DCHECK(version().HasIetfQuicFrames());
+  if (!start_peer_migration_earlier_ || QuicUtils::IsProbingFrame(type)) {
+    return;
+  }
+  QUIC_CODE_COUNT(quic_start_peer_migration_earlier);
+  if (GetLargestReceivedPacket().IsInitialized() &&
+      last_header_.packet_number == GetLargestReceivedPacket()) {
+    direct_peer_address_ = last_packet_source_address_;
+    if (current_effective_peer_migration_type_ != NO_CHANGE) {
+      // Start effective peer migration when the current packet contains a
+      // non-probing frame.
       // TODO(fayang): When multiple packet number spaces is supported, only
       // start peer migration for the application data.
       StartEffectivePeerMigration(current_effective_peer_migration_type_);

quic/core/quic_connection.h

diff --git a/quic/core/quic_connection.h b/quic/core/quic_connection.h
index 30d91d1..f6584ff 100644
--- a/quic/core/quic_connection.h
+++ b/quic/core/quic_connection.h
@@ -1368,6 +1368,10 @@
   // deprecated.
   void MaybeRespondToConnectivityProbingOrMigration();
 
+  // Called in IETF QUIC. Start peer migration if a non-probing frame is
+  // received and the current packet number is largest received so far.
+  void MaybeStartIetfPeerMigration(QuicFrameType type);
+
   QuicFramer framer_;
 
   // Contents received in the current packet, especially used to identify
@@ -1723,6 +1727,9 @@
 
   size_t anti_amplification_factor_ =
       GetQuicFlag(FLAGS_quic_anti_amplification_factor);
+
+  bool start_peer_migration_earlier_ =
+      GetQuicReloadableFlag(quic_start_peer_migration_earlier);
 };
 
 }  // namespace quic

quic/core/quic_connection_test.cc

diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index 0f1a8dc..e4cc987 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -375,8 +375,9 @@
   WriteResult WritePacket(const char* buffer,
                           size_t buf_len,
                           const QuicIpAddress& /*self_address*/,
-                          const QuicSocketAddress& /*peer_address*/,
+                          const QuicSocketAddress& peer_address,
                           PerPacketOptions* /*options*/) override {
+    last_write_peer_address_ = peer_address;
     // If the buffer is allocated from the pool, return it back to the pool.
     // Note the buffer content doesn't change.
     if (packet_buffer_pool_index_.find(const_cast<char*>(buffer)) !=
@@ -621,6 +622,10 @@
 
   SimpleQuicFramer* framer() { return &framer_; }
 
+  QuicSocketAddress last_write_peer_address() const {
+    return last_write_peer_address_;
+  }
+
  private:
   char* AllocPacketBuffer() {
     PacketBuffer* p = packet_buffer_free_list_.front();
@@ -673,6 +678,8 @@
   QuicHashMap<char*, PacketBuffer*> packet_buffer_pool_index_;
   // Indices in packet_buffer_pool_ that are not allocated.
   std::list<PacketBuffer*> packet_buffer_free_list_;
+  // The peer address passed into WritePacket().
+  QuicSocketAddress last_write_peer_address_;
 };
 
 class TestConnection : public QuicConnection {
@@ -1958,6 +1965,11 @@
   set_perspective(Perspective::IS_SERVER);
   QuicPacketCreatorPeer::SetSendVersionInPacket(creator_, false);
   EXPECT_EQ(Perspective::IS_SERVER, connection_.perspective());
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_FORWARD_SECURE);
+  peer_creator_.set_encryption_level(ENCRYPTION_FORWARD_SECURE);
+  // Prevent packets from being coalesced.
+  EXPECT_CALL(visitor_, GetHandshakeState())
+      .WillRepeatedly(Return(HANDSHAKE_CONFIRMED));
 
   // Clear direct_peer_address.
   QuicConnectionPeer::SetDirectPeerAddress(&connection_, QuicSocketAddress());
@@ -1967,23 +1979,30 @@
                                               QuicSocketAddress());
   EXPECT_FALSE(connection_.effective_peer_address().IsInitialized());
 
-  if (QuicVersionUsesCryptoFrames(connection_.transport_version())) {
-    EXPECT_CALL(visitor_, OnCryptoFrame(_)).Times(AnyNumber());
-  } else {
-    EXPECT_CALL(visitor_, OnStreamFrame(_)).Times(AnyNumber());
-  }
-  ProcessFramePacketWithAddresses(MakeCryptoFrame(), kSelfAddress,
-                                  kPeerAddress);
+  const QuicSocketAddress kNewPeerAddress =
+      QuicSocketAddress(QuicIpAddress::Loopback6(), /*port=*/23456);
+  EXPECT_CALL(visitor_, OnStreamFrame(_))
+      .WillOnce(Invoke(
+          [=]() { EXPECT_EQ(kPeerAddress, connection_.peer_address()); }))
+      .WillOnce(Invoke([=]() {
+        EXPECT_EQ((GetQuicReloadableFlag(quic_start_peer_migration_earlier) ||
+                           !GetParam().version.HasIetfQuicFrames()
+                       ? kNewPeerAddress
+                       : kPeerAddress),
+                  connection_.peer_address());
+      }));
+  QuicFrames frames;
+  frames.push_back(QuicFrame(frame1_));
+  ProcessFramesPacketWithAddresses(frames, kSelfAddress, kPeerAddress);
   EXPECT_EQ(kPeerAddress, connection_.peer_address());
   EXPECT_EQ(kPeerAddress, connection_.effective_peer_address());
 
   // Process another packet with a different peer address on server side will
   // start connection migration.
-  const QuicSocketAddress kNewPeerAddress =
-      QuicSocketAddress(QuicIpAddress::Loopback6(), /*port=*/23456);
   EXPECT_CALL(visitor_, OnConnectionMigration(PORT_CHANGE)).Times(1);
-  ProcessFramePacketWithAddresses(MakeCryptoFrame(), kSelfAddress,
-                                  kNewPeerAddress);
+  QuicFrames frames2;
+  frames2.push_back(QuicFrame(frame2_));
+  ProcessFramesPacketWithAddresses(frames2, kSelfAddress, kNewPeerAddress);
   EXPECT_EQ(kNewPeerAddress, connection_.peer_address());
   EXPECT_EQ(kNewPeerAddress, connection_.effective_peer_address());
 }

quic/core/quic_utils.cc

diff --git a/quic/core/quic_utils.cc b/quic/core/quic_utils.cc
index 367db32..13f6a2a 100644
--- a/quic/core/quic_utils.cc
+++ b/quic/core/quic_utils.cc
@@ -654,5 +654,18 @@
   }
 }
 
+// static
+bool QuicUtils::IsProbingFrame(QuicFrameType type) {
+  switch (type) {
+    case PATH_CHALLENGE_FRAME:
+    case PATH_RESPONSE_FRAME:
+    case NEW_CONNECTION_ID_FRAME:
+    case PADDING_FRAME:
+      return true;
+    default:
+      return false;
+  }
+}
+
 #undef RETURN_STRING_LITERAL  // undef for jumbo builds
 }  // namespace quic

quic/core/quic_utils.h

diff --git a/quic/core/quic_utils.h b/quic/core/quic_utils.h
index 9e190e0..05c7124 100644
--- a/quic/core/quic_utils.h
+++ b/quic/core/quic_utils.h
@@ -230,6 +230,9 @@
   // exceeds this value, it will result in a stream ID that exceeds the
   // implementation limit on stream ID size.
   static QuicStreamCount GetMaxStreamCount();
+
+  // Return true if this frame is an IETF probing frame.
+  static bool IsProbingFrame(QuicFrameType type);
 };
 
 template <typename Mask>