QuicConnection::ShouldGeneratePacket returns false when peer connection ID is required but unavailable.

Protected by FLAGS_quic_reloadable_flag_quic_connection_support_multiple_cids_v3.

PiperOrigin-RevId: 368481673
Change-Id: I92be6afd3aa243a0cd9d4fbdb8324a92fdaf4520
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index bf1f6c2..62e9f4d 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -393,7 +393,7 @@
       GetQuicRestartFlag(quic_time_wait_list_support_multiple_cid_v2) &&
       GetQuicRestartFlag(
           quic_dispatcher_support_multiple_cid_per_connection_v2) &&
-      GetQuicReloadableFlag(quic_connection_support_multiple_cids_v3);
+      GetQuicReloadableFlag(quic_connection_support_multiple_cids_v4);
 
   QUIC_DLOG(INFO) << ENDPOINT << "Created connection with server connection ID "
                   << server_connection_id
@@ -1707,7 +1707,8 @@
         group_path_response_and_challenge_sending_closer_
             ? nullptr
             : std::make_unique<QuicPacketCreator::ScopedPeerAddressContext>(
-                  &packet_creator_, last_packet_source_address_);
+                  &packet_creator_, last_packet_source_address_,
+                  /*update_connection_id=*/false);
     if (!OnPathChallengeFrameInternal(frame)) {
       return false;
     }
@@ -1729,7 +1730,8 @@
   std::unique_ptr<QuicPacketCreator::ScopedPeerAddressContext> context;
   if (group_path_response_and_challenge_sending_closer_) {
     context = std::make_unique<QuicPacketCreator::ScopedPeerAddressContext>(
-        &packet_creator_, last_packet_source_address_);
+        &packet_creator_, last_packet_source_address_,
+        /*update_connection_id=*/false);
   }
   if (should_proactively_validate_peer_address_on_path_challenge_) {
     QUIC_RELOADABLE_FLAG_COUNT(
@@ -1993,7 +1995,7 @@
       perspective_ == Perspective::IS_SERVER) {
     OnClientConnectionIdAvailable();
   }
-  QUIC_RELOADABLE_FLAG_COUNT_N(quic_connection_support_multiple_cids_v3, 1, 2);
+  QUIC_RELOADABLE_FLAG_COUNT_N(quic_connection_support_multiple_cids_v4, 1, 2);
   return true;
 }
 
@@ -2051,7 +2053,7 @@
                     ConnectionCloseBehavior::SEND_CONNECTION_CLOSE_PACKET);
     return false;
   }
-  QUIC_RELOADABLE_FLAG_COUNT_N(quic_connection_support_multiple_cids_v3, 2, 2);
+  QUIC_RELOADABLE_FLAG_COUNT_N(quic_connection_support_multiple_cids_v4, 2, 2);
   return true;
 }
 
@@ -3186,6 +3188,15 @@
                 QuicVersionUsesCryptoFrames(transport_version()))
       << ENDPOINT
       << "Handshake in STREAM frames should not check ShouldGeneratePacket";
+  if (support_multiple_connection_ids_ && peer_issued_cid_manager_ != nullptr &&
+      packet_creator_.GetDestinationConnectionId().IsEmpty()) {
+    QUIC_CODE_COUNT(quic_generate_packet_blocked_by_no_connection_id);
+    QUIC_BUG_IF(quic_bug_90265_1, perspective_ == Perspective::IS_CLIENT);
+    QUIC_DLOG(INFO) << ENDPOINT
+                    << "There is no destination connection ID available to "
+                       "generate packet.";
+    return false;
+  }
   if (!count_bytes_on_alternative_path_separately_) {
     return CanWrite(retransmittable);
   }
@@ -4457,8 +4468,8 @@
     QuicIetfTransportErrorCodes ietf_error,
     const std::string& details) {
   // Always use the current path to send CONNECTION_CLOSE.
-  QuicPacketCreator::ScopedPeerAddressContext context(&packet_creator_,
-                                                      peer_address());
+  QuicPacketCreator::ScopedPeerAddressContext context(
+      &packet_creator_, peer_address(), /*update_connection_id=*/false);
   if (!SupportsMultiplePacketNumberSpaces()) {
     QUIC_DLOG(INFO) << ENDPOINT << "Sending connection close packet.";
     if (!use_encryption_level_context_) {
@@ -6425,8 +6436,8 @@
     {
       // It's on current path, add the PATH_CHALLENGE the same way as other
       // frames.
-      QuicPacketCreator::ScopedPeerAddressContext context(&packet_creator_,
-                                                          peer_address);
+      QuicPacketCreator::ScopedPeerAddressContext context(
+          &packet_creator_, peer_address, /*update_connection_id=*/false);
       // This may cause connection to be closed.
       packet_creator_.AddPathChallengeFrame(data_buffer);
     }
@@ -6478,8 +6489,8 @@
   // Send PATH_RESPONSE using the provided peer address. If the creator has been
   // using a different peer address, it will flush before and after serializing
   // the current PATH_RESPONSE.
-  QuicPacketCreator::ScopedPeerAddressContext context(&packet_creator_,
-                                                      peer_address_to_send);
+  QuicPacketCreator::ScopedPeerAddressContext context(
+      &packet_creator_, peer_address_to_send, /*update_connection_id=*/false);
   QUIC_DVLOG(1) << ENDPOINT << "Send PATH_RESPONSE to " << peer_address_to_send;
   if (default_path_.self_address == last_packet_destination_address_) {
     // The PATH_CHALLENGE is received on the default socket. Respond on the same
diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index 1e81672..fcdcf4f 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -25,6 +25,7 @@
 #include "quic/core/quic_connection_id.h"
 #include "quic/core/quic_constants.h"
 #include "quic/core/quic_error_codes.h"
+#include "quic/core/quic_packet_creator.h"
 #include "quic/core/quic_packets.h"
 #include "quic/core/quic_path_validator.h"
 #include "quic/core/quic_simple_buffer_allocator.h"
@@ -33,6 +34,7 @@
 #include "quic/core/quic_versions.h"
 #include "quic/platform/api/quic_expect_bug.h"
 #include "quic/platform/api/quic_flags.h"
+#include "quic/platform/api/quic_ip_address.h"
 #include "quic/platform/api/quic_logging.h"
 #include "quic/platform/api/quic_reference_counted.h"
 #include "quic/platform/api/quic_socket_address.h"
@@ -14402,6 +14404,40 @@
   ASSERT_EQ(packet_creator->GetDestinationConnectionId(), frame.connection_id);
 }
 
+TEST_P(QuicConnectionTest, ShouldGeneratePacketBlockedByMissingConnectionId) {
+  if (!version().HasIetfQuicFrames() ||
+      !connection_.support_multiple_connection_ids()) {
+    return;
+  }
+  set_perspective(Perspective::IS_SERVER);
+  connection_.set_client_connection_id(TestConnectionId(1));
+  connection_.CreateConnectionIdManager();
+  if (version().SupportsAntiAmplificationLimit()) {
+    QuicConnectionPeer::SetAddressValidated(&connection_);
+  }
+
+  ASSERT_TRUE(
+      connection_.ShouldGeneratePacket(NO_RETRANSMITTABLE_DATA, NOT_HANDSHAKE));
+
+  QuicPacketCreator* packet_creator =
+      QuicConnectionPeer::GetPacketCreator(&connection_);
+  QuicIpAddress peer_host1;
+  peer_host1.FromString("12.12.12.12");
+  QuicSocketAddress peer_address1(peer_host1, 1235);
+
+  {
+    // No connection ID is available as context is created without any.
+    QuicPacketCreator::ScopedPeerAddressContext context(
+        packet_creator, peer_address1, EmptyQuicConnectionId(),
+        EmptyQuicConnectionId(),
+        /*update_connection_id=*/true);
+    ASSERT_FALSE(connection_.ShouldGeneratePacket(NO_RETRANSMITTABLE_DATA,
+                                                  NOT_HANDSHAKE));
+  }
+  ASSERT_TRUE(
+      connection_.ShouldGeneratePacket(NO_RETRANSMITTABLE_DATA, NOT_HANDSHAKE));
+}
+
 // Regression test for b/182571515
 TEST_P(QuicConnectionTest, LostDataThenGetAcknowledged) {
   set_perspective(Perspective::IS_SERVER);
diff --git a/quic/core/quic_flags_list.h b/quic/core/quic_flags_list.h
index 582a1cc..099043c 100644
--- a/quic/core/quic_flags_list.h
+++ b/quic/core/quic_flags_list.h
@@ -17,7 +17,7 @@
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_bbr2_fix_bw_lo_mode, false)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_can_send_ack_frequency, true)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_close_connection_with_too_many_outstanding_packets, true)
-QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_connection_support_multiple_cids_v3, false)
+QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_connection_support_multiple_cids_v4, true)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_conservative_bursts, false)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_conservative_cwnd_and_pacing_gains, false)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_count_bytes_on_alternative_path_seperately, true)
diff --git a/quic/core/quic_packet_creator.cc b/quic/core/quic_packet_creator.cc
index 42aebe5..9fb2f5d 100644
--- a/quic/core/quic_packet_creator.cc
+++ b/quic/core/quic_packet_creator.cc
@@ -2007,17 +2007,44 @@
 
 QuicPacketCreator::ScopedPeerAddressContext::ScopedPeerAddressContext(
     QuicPacketCreator* creator,
-    QuicSocketAddress address)
-    : creator_(creator), old_peer_address_(creator_->packet_.peer_address) {
-  QUIC_BUG_IF(quic_bug_12398_19,
-              !creator_->packet_.peer_address.IsInitialized())
+    QuicSocketAddress address,
+    bool update_connection_id)
+    : ScopedPeerAddressContext(creator,
+                               address,
+                               EmptyQuicConnectionId(),
+                               EmptyQuicConnectionId(),
+                               update_connection_id) {}
+
+QuicPacketCreator::ScopedPeerAddressContext::ScopedPeerAddressContext(
+    QuicPacketCreator* creator,
+    QuicSocketAddress address,
+    const QuicConnectionId& client_connection_id,
+    const QuicConnectionId& server_connection_id,
+    bool update_connection_id)
+    : creator_(creator),
+      old_peer_address_(creator_->packet_.peer_address),
+      old_client_connection_id_(creator_->GetClientConnectionId()),
+      old_server_connection_id_(creator_->GetServerConnectionId()),
+      update_connection_id_(update_connection_id) {
+  QUIC_BUG_IF(quic_bug_12398_19, !old_peer_address_.IsInitialized())
       << "Context is used before seralized packet's peer address is "
          "initialized.";
   creator_->SetDefaultPeerAddress(address);
+  if (update_connection_id_) {
+    QUICHE_DCHECK(address != old_peer_address_ ||
+                  ((client_connection_id == old_client_connection_id_) &&
+                   (server_connection_id == old_server_connection_id_)));
+    creator_->SetClientConnectionId(client_connection_id);
+    creator_->SetServerConnectionId(server_connection_id);
+  }
 }
 
 QuicPacketCreator::ScopedPeerAddressContext::~ScopedPeerAddressContext() {
   creator_->SetDefaultPeerAddress(old_peer_address_);
+  if (update_connection_id_) {
+    creator_->SetClientConnectionId(old_client_connection_id_);
+    creator_->SetServerConnectionId(old_server_connection_id_);
+  }
 }
 
 QuicPacketCreator::ScopedSerializationFailureHandler::
diff --git a/quic/core/quic_packet_creator.h b/quic/core/quic_packet_creator.h
index 5e0c65e..2fc32f6 100644
--- a/quic/core/quic_packet_creator.h
+++ b/quic/core/quic_packet_creator.h
@@ -25,6 +25,7 @@
 #include "quic/core/frames/quic_stream_frame.h"
 #include "quic/core/quic_circular_deque.h"
 #include "quic/core/quic_coalesced_packet.h"
+#include "quic/core/quic_connection_id.h"
 #include "quic/core/quic_framer.h"
 #include "quic/core/quic_packets.h"
 #include "quic/core/quic_types.h"
@@ -83,18 +84,28 @@
     virtual void OnStreamFrameCoalesced(const QuicStreamFrame& /*frame*/) {}
   };
 
-  // Set the peer address which the serialized packet will be sent to during the
-  // scope of this object. Upon exiting the scope, the original peer address is
-  // restored.
+  // Set the peer address and connection IDs with which the serialized packet
+  // will be sent to during the scope of this object. Upon exiting the scope,
+  // the original peer address and connection IDs are restored.
   class QUIC_EXPORT_PRIVATE ScopedPeerAddressContext {
    public:
     ScopedPeerAddressContext(QuicPacketCreator* creator,
-                             QuicSocketAddress address);
+                             QuicSocketAddress address,
+                             bool update_connection_id);
+
+    ScopedPeerAddressContext(QuicPacketCreator* creator,
+                             QuicSocketAddress address,
+                             const QuicConnectionId& client_connection_id,
+                             const QuicConnectionId& server_connection_id,
+                             bool update_connection_id);
     ~ScopedPeerAddressContext();
 
    private:
     QuicPacketCreator* creator_;
     QuicSocketAddress old_peer_address_;
+    QuicConnectionId old_client_connection_id_;
+    QuicConnectionId old_server_connection_id_;
+    bool update_connection_id_;
   };
 
   QuicPacketCreator(QuicConnectionId server_connection_id,
@@ -261,6 +272,16 @@
   // Returns a dummy packet that is valid but contains no useful information.
   static SerializedPacket NoPacket();
 
+  // Returns the server connection ID to send over the wire.
+  const QuicConnectionId& GetServerConnectionId() const {
+    return server_connection_id_;
+  }
+
+  // Returns the client connection ID to send over the wire.
+  const QuicConnectionId& GetClientConnectionId() const {
+    return client_connection_id_;
+  }
+
   // Returns the destination connection ID to send over the wire.
   QuicConnectionId GetDestinationConnectionId() const;
 
diff --git a/quic/core/quic_packet_creator_test.cc b/quic/core/quic_packet_creator_test.cc
index 5986bd2..3ccf3d8 100644
--- a/quic/core/quic_packet_creator_test.cc
+++ b/quic/core/quic_packet_creator_test.cc
@@ -19,6 +19,7 @@
 #include "quic/core/crypto/quic_decrypter.h"
 #include "quic/core/crypto/quic_encrypter.h"
 #include "quic/core/frames/quic_stream_frame.h"
+#include "quic/core/quic_connection_id.h"
 #include "quic/core/quic_data_writer.h"
 #include "quic/core/quic_simple_buffer_allocator.h"
 #include "quic/core/quic_types.h"
@@ -3811,8 +3812,12 @@
 
 TEST_F(QuicPacketCreatorMultiplePacketsTest,
        PeerAddressContextWithSameAddress) {
+  QuicConnectionId client_connection_id = TestConnectionId(1);
+  QuicConnectionId server_connection_id = TestConnectionId(2);
   QuicSocketAddress peer_addr(QuicIpAddress::Any4(), 12345);
   creator_.SetDefaultPeerAddress(peer_addr);
+  creator_.SetClientConnectionId(client_connection_id);
+  creator_.SetServerConnectionId(server_connection_id);
   // Send some stream data.
   MakeIOVector("foo", &iov_);
   EXPECT_CALL(delegate_, ShouldGeneratePacket(_, _))
@@ -3824,8 +3829,12 @@
   EXPECT_EQ(3u, consumed.bytes_consumed);
   EXPECT_TRUE(creator_.HasPendingFrames());
   {
-    // Set a different address via context which should trigger flush.
-    QuicPacketCreator::ScopedPeerAddressContext context(&creator_, peer_addr);
+    // Set the same address via context which should not trigger flush.
+    QuicPacketCreator::ScopedPeerAddressContext context(
+        &creator_, peer_addr, client_connection_id, server_connection_id,
+        /*update_connection_id=*/true);
+    ASSERT_EQ(client_connection_id, creator_.GetClientConnectionId());
+    ASSERT_EQ(server_connection_id, creator_.GetServerConnectionId());
     EXPECT_TRUE(creator_.HasPendingFrames());
     // Queue another STREAM_FRAME.
     QuicConsumedData consumed = creator_.ConsumeData(
@@ -3874,8 +3883,14 @@
       }));
   EXPECT_TRUE(creator_.HasPendingFrames());
   {
+    QuicConnectionId client_connection_id = TestConnectionId(1);
+    QuicConnectionId server_connection_id = TestConnectionId(2);
     // Set a different address via context which should trigger flush.
-    QuicPacketCreator::ScopedPeerAddressContext context(&creator_, peer_addr1);
+    QuicPacketCreator::ScopedPeerAddressContext context(
+        &creator_, peer_addr1, client_connection_id, server_connection_id,
+        /*update_connection_id=*/true);
+    ASSERT_EQ(client_connection_id, creator_.GetClientConnectionId());
+    ASSERT_EQ(server_connection_id, creator_.GetServerConnectionId());
     EXPECT_FALSE(creator_.HasPendingFrames());
     // Queue another STREAM_FRAME.
     QuicConsumedData consumed = creator_.ConsumeData(
@@ -3891,9 +3906,15 @@
 
 TEST_F(QuicPacketCreatorMultiplePacketsTest,
        NestedPeerAddressContextWithDifferentAddress) {
+  QuicConnectionId client_connection_id1 = creator_.GetClientConnectionId();
+  QuicConnectionId server_connection_id1 = creator_.GetServerConnectionId();
   QuicSocketAddress peer_addr(QuicIpAddress::Any4(), 12345);
   creator_.SetDefaultPeerAddress(peer_addr);
-  QuicPacketCreator::ScopedPeerAddressContext context(&creator_, peer_addr);
+  QuicPacketCreator::ScopedPeerAddressContext context(
+      &creator_, peer_addr, client_connection_id1, server_connection_id1,
+      /*update_connection_id=*/true);
+  ASSERT_EQ(client_connection_id1, creator_.GetClientConnectionId());
+  ASSERT_EQ(server_connection_id1, creator_.GetServerConnectionId());
 
   // Send some stream data.
   MakeIOVector("foo", &iov_);
@@ -3913,9 +3934,14 @@
         ASSERT_EQ(1u, packet.retransmittable_frames.size());
         EXPECT_EQ(STREAM_FRAME, packet.retransmittable_frames.front().type);
 
+        QuicConnectionId client_connection_id2 = TestConnectionId(3);
+        QuicConnectionId server_connection_id2 = TestConnectionId(4);
         // Set up another context with a different address.
-        QuicPacketCreator::ScopedPeerAddressContext context(&creator_,
-                                                            peer_addr1);
+        QuicPacketCreator::ScopedPeerAddressContext context(
+            &creator_, peer_addr1, client_connection_id2, server_connection_id2,
+            /*update_connection_id=*/true);
+        ASSERT_EQ(client_connection_id2, creator_.GetClientConnectionId());
+        ASSERT_EQ(server_connection_id2, creator_.GetServerConnectionId());
         MakeIOVector("foo", &iov_);
         EXPECT_CALL(delegate_, ShouldGeneratePacket(_, _))
             .WillRepeatedly(Return(true));