Allow connections to request unroutable connection IDs when Maglev is out of sync.

Protected by FLAGS_quic_restart_flag_quic_request_unroutable_cids.

PiperOrigin-RevId: 574849171
diff --git a/quiche/quic/core/quic_buffered_packet_store.cc b/quiche/quic/core/quic_buffered_packet_store.cc
index 650a08f..48690ce 100644
--- a/quiche/quic/core/quic_buffered_packet_store.cc
+++ b/quiche/quic/core/quic_buffered_packet_store.cc
@@ -4,14 +4,30 @@
 
 #include "quiche/quic/core/quic_buffered_packet_store.h"
 
+#include <cstddef>
+#include <list>
+#include <memory>
 #include <string>
+#include <utility>
 
 #include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
+#include "quiche/quic/core/connection_id_generator.h"
+#include "quiche/quic/core/quic_alarm.h"
+#include "quiche/quic/core/quic_alarm_factory.h"
+#include "quiche/quic/core/quic_clock.h"
 #include "quiche/quic/core/quic_connection_id.h"
+#include "quiche/quic/core/quic_constants.h"
+#include "quiche/quic/core/quic_error_codes.h"
+#include "quiche/quic/core/quic_framer.h"
+#include "quiche/quic/core/quic_packets.h"
+#include "quiche/quic/core/quic_time.h"
 #include "quiche/quic/core/quic_types.h"
 #include "quiche/quic/core/quic_versions.h"
 #include "quiche/quic/platform/api/quic_bug_tracker.h"
 #include "quiche/quic/platform/api/quic_flags.h"
+#include "quiche/quic/platform/api/quic_socket_address.h"
+#include "quiche/common/platform/api/quiche_logging.h"
 
 namespace quic {
 
@@ -89,7 +105,8 @@
     QuicConnectionId connection_id, bool ietf_quic,
     const QuicReceivedPacket& packet, QuicSocketAddress self_address,
     QuicSocketAddress peer_address, const ParsedQuicVersion& version,
-    absl::optional<ParsedClientHello> parsed_chlo) {
+    absl::optional<ParsedClientHello> parsed_chlo,
+    ConnectionIdGeneratorInterface* connection_id_generator) {
   const bool is_chlo = parsed_chlo.has_value();
   QUIC_BUG_IF(quic_bug_12410_1, !GetQuicFlag(quic_allow_chlo_buffering))
       << "Shouldn't buffer packets if disabled via flag.";
@@ -144,6 +161,7 @@
     connections_with_chlo_[connection_id] = false;  // Dummy value.
     // Set the version of buffered packets of this connection on CHLO.
     queue.version = version;
+    queue.connection_id_generator = connection_id_generator;
   } else {
     // Buffer non-CHLO packets in arrival order.
     queue.buffered_packets.push_back(std::move(new_entry));
diff --git a/quiche/quic/core/quic_buffered_packet_store.h b/quiche/quic/core/quic_buffered_packet_store.h
index 0973d5c..012374a 100644
--- a/quiche/quic/core/quic_buffered_packet_store.h
+++ b/quiche/quic/core/quic_buffered_packet_store.h
@@ -5,14 +5,22 @@
 #ifndef QUICHE_QUIC_CORE_QUIC_BUFFERED_PACKET_STORE_H_
 #define QUICHE_QUIC_CORE_QUIC_BUFFERED_PACKET_STORE_H_
 
+#include <cstdint>
 #include <list>
+#include <memory>
 #include <string>
+#include <vector>
 
+#include "absl/types/optional.h"
+#include "quiche/quic/core/connection_id_generator.h"
 #include "quiche/quic/core/quic_alarm.h"
 #include "quiche/quic/core/quic_alarm_factory.h"
 #include "quiche/quic/core/quic_clock.h"
+#include "quiche/quic/core/quic_connection_id.h"
 #include "quiche/quic/core/quic_packets.h"
 #include "quiche/quic/core/quic_time.h"
+#include "quiche/quic/core/quic_types.h"
+#include "quiche/quic/core/quic_versions.h"
 #include "quiche/quic/core/tls_chlo_extractor.h"
 #include "quiche/quic/platform/api/quic_export.h"
 #include "quiche/quic/platform/api/quic_socket_address.h"
@@ -75,6 +83,11 @@
     // Otherwise, it is the version of the first packet in |buffered_packets|.
     ParsedQuicVersion version;
     TlsChloExtractor tls_chlo_extractor;
+    // Only one reference to the generator is stored per connection, and this is
+    // stored when the CHLO is buffered. The connection needs a stable,
+    // consistent way to generate IDs. Fixing it on the CHLO is a
+    // straightforward way to enforce that.
+    ConnectionIdGeneratorInterface* connection_id_generator = nullptr;
   };
 
   using BufferedPacketMap =
@@ -101,12 +114,15 @@
 
   // Adds a copy of packet into the packet queue for given connection. If the
   // packet is the last one of the CHLO, |parsed_chlo| will contain a parsed
-  // version of the CHLO.
+  // version of the CHLO. |connection_id_generator| is the Connection ID
+  // Generator to use with the connection. It is ignored if |parsed_chlo| is
+  // absent.
   EnqueuePacketResult EnqueuePacket(
       QuicConnectionId connection_id, bool ietf_quic,
       const QuicReceivedPacket& packet, QuicSocketAddress self_address,
       QuicSocketAddress peer_address, const ParsedQuicVersion& version,
-      absl::optional<ParsedClientHello> parsed_chlo);
+      absl::optional<ParsedClientHello> parsed_chlo,
+      ConnectionIdGeneratorInterface* connection_id_generator);
 
   // Returns true if there are any packets buffered for |connection_id|.
   bool HasBufferedPackets(QuicConnectionId connection_id) const;
diff --git a/quiche/quic/core/quic_buffered_packet_store_test.cc b/quiche/quic/core/quic_buffered_packet_store_test.cc
index 562b516..640ac5d8 100644
--- a/quiche/quic/core/quic_buffered_packet_store_test.cc
+++ b/quiche/quic/core/quic_buffered_packet_store_test.cc
@@ -4,18 +4,31 @@
 
 #include "quiche/quic/core/quic_buffered_packet_store.h"
 
+#include <cstddef>
+#include <cstdint>
 #include <list>
 #include <memory>
 #include <string>
+#include <utility>
 
+#include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
+#include "quiche/quic/core/connection_id_generator.h"
 #include "quiche/quic/core/crypto/transport_parameters.h"
 #include "quiche/quic/core/quic_connection_id.h"
+#include "quiche/quic/core/quic_constants.h"
 #include "quiche/quic/core/quic_error_codes.h"
+#include "quiche/quic/core/quic_framer.h"
+#include "quiche/quic/core/quic_packets.h"
+#include "quiche/quic/core/quic_time.h"
 #include "quiche/quic/core/quic_types.h"
 #include "quiche/quic/core/quic_versions.h"
+#include "quiche/quic/platform/api/quic_ip_address.h"
+#include "quiche/quic/platform/api/quic_socket_address.h"
 #include "quiche/quic/platform/api/quic_test.h"
 #include "quiche/quic/test_tools/first_flight.h"
 #include "quiche/quic/test_tools/mock_clock.h"
+#include "quiche/quic/test_tools/mock_connection_id_generator.h"
 #include "quiche/quic/test_tools/quic_buffered_packet_store_peer.h"
 #include "quiche/quic/test_tools/quic_test_utils.h"
 
@@ -82,12 +95,13 @@
   QuicReceivedPacket packet_;
   const ParsedQuicVersion invalid_version_;
   const ParsedQuicVersion valid_version_;
+  MockConnectionIdGenerator connection_id_generator_;
 };
 
 TEST_F(QuicBufferedPacketStoreTest, SimpleEnqueueAndDeliverPacket) {
   QuicConnectionId connection_id = TestConnectionId(1);
   store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                       peer_address_, invalid_version_, kNoParsedChlo);
+                       peer_address_, invalid_version_, kNoParsedChlo, nullptr);
   EXPECT_TRUE(store_.HasBufferedPackets(connection_id));
   auto packets = store_.DeliverPackets(connection_id);
   const std::list<BufferedPacket>& queue = packets.buffered_packets;
@@ -109,9 +123,10 @@
   QuicSocketAddress addr_with_new_port(QuicIpAddress::Any4(), 256);
   QuicConnectionId connection_id = TestConnectionId(1);
   store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                       peer_address_, invalid_version_, kNoParsedChlo);
+                       peer_address_, invalid_version_, kNoParsedChlo, nullptr);
   store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                       addr_with_new_port, invalid_version_, kNoParsedChlo);
+                       addr_with_new_port, invalid_version_, kNoParsedChlo,
+                       nullptr);
   std::list<BufferedPacket> queue =
       store_.DeliverPackets(connection_id).buffered_packets;
   ASSERT_EQ(2u, queue.size());
@@ -126,9 +141,11 @@
   for (uint64_t conn_id = 1; conn_id <= num_connections; ++conn_id) {
     QuicConnectionId connection_id = TestConnectionId(conn_id);
     store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                         peer_address_, invalid_version_, kNoParsedChlo);
+                         peer_address_, invalid_version_, kNoParsedChlo,
+                         nullptr);
     store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                         peer_address_, invalid_version_, kNoParsedChlo);
+                         peer_address_, invalid_version_, kNoParsedChlo,
+                         nullptr);
   }
 
   // Deliver packets in reversed order.
@@ -148,15 +165,15 @@
   QuicConnectionId connection_id = TestConnectionId(1);
   // Arrived CHLO packet shouldn't affect how many non-CHLO pacekts store can
   // keep.
-  EXPECT_EQ(
-      QuicBufferedPacketStore::SUCCESS,
-      store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                           peer_address_, valid_version_, kDefaultParsedChlo));
+  EXPECT_EQ(QuicBufferedPacketStore::SUCCESS,
+            store_.EnqueuePacket(connection_id, false, packet_, self_address_,
+                                 peer_address_, valid_version_,
+                                 kDefaultParsedChlo, nullptr));
   for (size_t i = 1; i <= num_packets; ++i) {
     // Only first |kDefaultMaxUndecryptablePackets packets| will be buffered.
-    EnqueuePacketResult result =
-        store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                             peer_address_, invalid_version_, kNoParsedChlo);
+    EnqueuePacketResult result = store_.EnqueuePacket(
+        connection_id, false, packet_, self_address_, peer_address_,
+        invalid_version_, kNoParsedChlo, nullptr);
     if (i <= kDefaultMaxUndecryptablePackets) {
       EXPECT_EQ(EnqueuePacketResult::SUCCESS, result);
     } else {
@@ -176,9 +193,9 @@
   const size_t kNumConnections = kMaxConnectionsWithoutCHLO + 1;
   for (uint64_t conn_id = 1; conn_id <= kNumConnections; ++conn_id) {
     QuicConnectionId connection_id = TestConnectionId(conn_id);
-    EnqueuePacketResult result =
-        store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                             peer_address_, invalid_version_, kNoParsedChlo);
+    EnqueuePacketResult result = store_.EnqueuePacket(
+        connection_id, false, packet_, self_address_, peer_address_,
+        invalid_version_, kNoParsedChlo, nullptr);
     if (conn_id <= kMaxConnectionsWithoutCHLO) {
       EXPECT_EQ(EnqueuePacketResult::SUCCESS, result);
     } else {
@@ -208,7 +225,7 @@
     EXPECT_EQ(EnqueuePacketResult::SUCCESS,
               store_.EnqueuePacket(TestConnectionId(conn_id), false, packet_,
                                    self_address_, peer_address_, valid_version_,
-                                   kDefaultParsedChlo));
+                                   kDefaultParsedChlo, nullptr));
   }
 
   // Send data packets on another |kMaxConnectionsWithoutCHLO| connections.
@@ -216,9 +233,9 @@
   for (uint64_t conn_id = num_chlos + 1;
        conn_id <= (kDefaultMaxConnectionsInStore + 1); ++conn_id) {
     QuicConnectionId connection_id = TestConnectionId(conn_id);
-    EnqueuePacketResult result =
-        store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                             peer_address_, valid_version_, kDefaultParsedChlo);
+    EnqueuePacketResult result = store_.EnqueuePacket(
+        connection_id, false, packet_, self_address_, peer_address_,
+        valid_version_, kDefaultParsedChlo, nullptr);
     if (conn_id <= kDefaultMaxConnectionsInStore) {
       EXPECT_EQ(EnqueuePacketResult::SUCCESS, result);
     } else {
@@ -227,23 +244,71 @@
   }
 }
 
+TEST_F(QuicBufferedPacketStoreTest, BasicGeneratorBuffering) {
+  EXPECT_EQ(
+      EnqueuePacketResult::SUCCESS,
+      store_.EnqueuePacket(TestConnectionId(1), false, packet_, self_address_,
+                           peer_address_, valid_version_, kDefaultParsedChlo,
+                           &connection_id_generator_));
+  QuicConnectionId delivered_conn_id;
+  BufferedPacketList packet_list =
+      store_.DeliverPacketsForNextConnection(&delivered_conn_id);
+  EXPECT_EQ(1u, packet_list.buffered_packets.size());
+  EXPECT_EQ(delivered_conn_id, TestConnectionId(1));
+  EXPECT_EQ(&connection_id_generator_, packet_list.connection_id_generator);
+}
+
+TEST_F(QuicBufferedPacketStoreTest, NullGeneratorOk) {
+  EXPECT_EQ(EnqueuePacketResult::SUCCESS,
+            store_.EnqueuePacket(TestConnectionId(1), false, packet_,
+                                 self_address_, peer_address_, valid_version_,
+                                 kDefaultParsedChlo, nullptr));
+  QuicConnectionId delivered_conn_id;
+  BufferedPacketList packet_list =
+      store_.DeliverPacketsForNextConnection(&delivered_conn_id);
+  EXPECT_EQ(1u, packet_list.buffered_packets.size());
+  EXPECT_EQ(delivered_conn_id, TestConnectionId(1));
+  EXPECT_EQ(packet_list.connection_id_generator, nullptr);
+}
+
+TEST_F(QuicBufferedPacketStoreTest, GeneratorIgnoredForNonChlo) {
+  MockConnectionIdGenerator generator2;
+  EXPECT_EQ(
+      EnqueuePacketResult::SUCCESS,
+      store_.EnqueuePacket(TestConnectionId(1), false, packet_, self_address_,
+                           peer_address_, valid_version_, kDefaultParsedChlo,
+                           &connection_id_generator_));
+  EXPECT_EQ(EnqueuePacketResult::SUCCESS,
+            store_.EnqueuePacket(TestConnectionId(1), false, packet_,
+                                 self_address_, peer_address_, valid_version_,
+                                 kNoParsedChlo, &generator2));
+  QuicConnectionId delivered_conn_id;
+  BufferedPacketList packet_list =
+      store_.DeliverPacketsForNextConnection(&delivered_conn_id);
+  EXPECT_EQ(2u, packet_list.buffered_packets.size());
+  EXPECT_EQ(delivered_conn_id, TestConnectionId(1));
+  EXPECT_EQ(packet_list.connection_id_generator, &connection_id_generator_);
+}
+
 TEST_F(QuicBufferedPacketStoreTest, EnqueueChloOnTooManyDifferentConnections) {
   // Buffer data packets on different connections upto limit.
   for (uint64_t conn_id = 1; conn_id <= kMaxConnectionsWithoutCHLO; ++conn_id) {
     QuicConnectionId connection_id = TestConnectionId(conn_id);
-    EXPECT_EQ(
-        EnqueuePacketResult::SUCCESS,
-        store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                             peer_address_, invalid_version_, kNoParsedChlo));
+    EXPECT_EQ(EnqueuePacketResult::SUCCESS,
+              // connection_id_generator_ will be ignored because the chlo has
+              // not been parsed.
+              store_.EnqueuePacket(connection_id, false, packet_, self_address_,
+                                   peer_address_, invalid_version_,
+                                   kNoParsedChlo, &connection_id_generator_));
   }
 
   // Buffer CHLOs on other connections till store is full.
   for (size_t i = kMaxConnectionsWithoutCHLO + 1;
        i <= kDefaultMaxConnectionsInStore + 1; ++i) {
     QuicConnectionId connection_id = TestConnectionId(i);
-    EnqueuePacketResult rs =
-        store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                             peer_address_, valid_version_, kDefaultParsedChlo);
+    EnqueuePacketResult rs = store_.EnqueuePacket(
+        connection_id, false, packet_, self_address_, peer_address_,
+        valid_version_, kDefaultParsedChlo, &connection_id_generator_);
     if (i <= kDefaultMaxConnectionsInStore) {
       EXPECT_EQ(EnqueuePacketResult::SUCCESS, rs);
       EXPECT_TRUE(store_.HasChloForConnection(connection_id));
@@ -257,11 +322,11 @@
   // But buffering a CHLO belonging to a connection already has data packet
   // buffered in the store should success. This is the connection should be
   // delivered at last.
-  EXPECT_EQ(
-      EnqueuePacketResult::SUCCESS,
-      store_.EnqueuePacket(
-          /*connection_id=*/TestConnectionId(1), false, packet_, self_address_,
-          peer_address_, valid_version_, kDefaultParsedChlo));
+  EXPECT_EQ(EnqueuePacketResult::SUCCESS,
+            store_.EnqueuePacket(
+                /*connection_id=*/TestConnectionId(1), false, packet_,
+                self_address_, peer_address_, valid_version_,
+                kDefaultParsedChlo, &connection_id_generator_));
   EXPECT_TRUE(store_.HasChloForConnection(
       /*connection_id=*/TestConnectionId(1)));
 
@@ -269,17 +334,18 @@
   for (size_t i = 0;
        i < kDefaultMaxConnectionsInStore - kMaxConnectionsWithoutCHLO + 1;
        ++i) {
+    BufferedPacketList packet_list =
+        store_.DeliverPacketsForNextConnection(&delivered_conn_id);
     if (i < kDefaultMaxConnectionsInStore - kMaxConnectionsWithoutCHLO) {
       // Only CHLO is buffered.
-      EXPECT_EQ(1u, store_.DeliverPacketsForNextConnection(&delivered_conn_id)
-                        .buffered_packets.size());
+      EXPECT_EQ(1u, packet_list.buffered_packets.size());
       EXPECT_EQ(TestConnectionId(i + kMaxConnectionsWithoutCHLO + 1),
                 delivered_conn_id);
     } else {
-      EXPECT_EQ(2u, store_.DeliverPacketsForNextConnection(&delivered_conn_id)
-                        .buffered_packets.size());
+      EXPECT_EQ(2u, packet_list.buffered_packets.size());
       EXPECT_EQ(TestConnectionId(1u), delivered_conn_id);
     }
+    EXPECT_EQ(packet_list.connection_id_generator, &connection_id_generator_);
   }
   EXPECT_FALSE(store_.HasChlosBuffered());
 }
@@ -289,26 +355,27 @@
 TEST_F(QuicBufferedPacketStoreTest, PacketQueueExpiredBeforeDelivery) {
   QuicConnectionId connection_id = TestConnectionId(1);
   store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                       peer_address_, invalid_version_, kNoParsedChlo);
-  EXPECT_EQ(
-      EnqueuePacketResult::SUCCESS,
-      store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                           peer_address_, valid_version_, kDefaultParsedChlo));
+                       peer_address_, invalid_version_, kNoParsedChlo,
+                       &connection_id_generator_);
+  EXPECT_EQ(EnqueuePacketResult::SUCCESS,
+            store_.EnqueuePacket(
+                connection_id, false, packet_, self_address_, peer_address_,
+                valid_version_, kDefaultParsedChlo, &connection_id_generator_));
   QuicConnectionId connection_id2 = TestConnectionId(2);
-  EXPECT_EQ(
-      EnqueuePacketResult::SUCCESS,
-      store_.EnqueuePacket(connection_id2, false, packet_, self_address_,
-                           peer_address_, invalid_version_, kNoParsedChlo));
+  EXPECT_EQ(EnqueuePacketResult::SUCCESS,
+            store_.EnqueuePacket(connection_id2, false, packet_, self_address_,
+                                 peer_address_, invalid_version_, kNoParsedChlo,
+                                 &connection_id_generator_));
 
   // CHLO on connection 3 arrives 1ms later.
   clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(1));
   QuicConnectionId connection_id3 = TestConnectionId(3);
-  // Use different client address to differetiate packets from different
+  // Use different client address to differentiate packets from different
   // connections.
   QuicSocketAddress another_client_address(QuicIpAddress::Any4(), 255);
   store_.EnqueuePacket(connection_id3, false, packet_, self_address_,
                        another_client_address, valid_version_,
-                       kDefaultParsedChlo);
+                       kDefaultParsedChlo, &connection_id_generator_);
 
   // Advance clock to the time when connection 1 and 2 expires.
   clock_.AdvanceTime(
@@ -328,21 +395,24 @@
   ASSERT_EQ(0u, store_.DeliverPackets(connection_id).buffered_packets.size());
   ASSERT_EQ(0u, store_.DeliverPackets(connection_id2).buffered_packets.size());
   QuicConnectionId delivered_conn_id;
-  auto queue = store_.DeliverPacketsForNextConnection(&delivered_conn_id)
-                   .buffered_packets;
+  BufferedPacketList packet_list =
+      store_.DeliverPacketsForNextConnection(&delivered_conn_id);
+
   // Connection 3 is the next to be delivered as connection 1 already expired.
   EXPECT_EQ(connection_id3, delivered_conn_id);
-  ASSERT_EQ(1u, queue.size());
+  EXPECT_EQ(packet_list.connection_id_generator, &connection_id_generator_);
+  ASSERT_EQ(1u, packet_list.buffered_packets.size());
   // Packets in connection 3 should use another peer address.
-  EXPECT_EQ(another_client_address, queue.front().peer_address);
+  EXPECT_EQ(another_client_address,
+            packet_list.buffered_packets.front().peer_address);
 
   // Test the alarm is reset by enqueueing 2 packets for 4th connection and wait
   // for them to expire.
   QuicConnectionId connection_id4 = TestConnectionId(4);
   store_.EnqueuePacket(connection_id4, false, packet_, self_address_,
-                       peer_address_, invalid_version_, kNoParsedChlo);
+                       peer_address_, invalid_version_, kNoParsedChlo, nullptr);
   store_.EnqueuePacket(connection_id4, false, packet_, self_address_,
-                       peer_address_, invalid_version_, kNoParsedChlo);
+                       peer_address_, invalid_version_, kNoParsedChlo, nullptr);
   clock_.AdvanceTime(
       QuicBufferedPacketStorePeer::expiration_alarm(&store_)->deadline() -
       clock_.ApproximateNow());
@@ -357,9 +427,9 @@
 
   // Enqueue some packets
   store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                       peer_address_, invalid_version_, kNoParsedChlo);
+                       peer_address_, invalid_version_, kNoParsedChlo, nullptr);
   store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                       peer_address_, invalid_version_, kNoParsedChlo);
+                       peer_address_, invalid_version_, kNoParsedChlo, nullptr);
   EXPECT_TRUE(store_.HasBufferedPackets(connection_id));
   EXPECT_FALSE(store_.HasChlosBuffered());
 
@@ -383,11 +453,12 @@
 
   // Enqueue some packets, which include a CHLO
   store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                       peer_address_, invalid_version_, kNoParsedChlo);
+                       peer_address_, invalid_version_, kNoParsedChlo, nullptr);
   store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                       peer_address_, valid_version_, kDefaultParsedChlo);
+                       peer_address_, valid_version_, kDefaultParsedChlo,
+                       nullptr);
   store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                       peer_address_, invalid_version_, kNoParsedChlo);
+                       peer_address_, invalid_version_, kNoParsedChlo, nullptr);
   EXPECT_TRUE(store_.HasBufferedPackets(connection_id));
   EXPECT_TRUE(store_.HasChlosBuffered());
 
@@ -412,15 +483,15 @@
 
   // Enqueue some packets for two connection IDs
   store_.EnqueuePacket(connection_id_1, false, packet_, self_address_,
-                       peer_address_, invalid_version_, kNoParsedChlo);
+                       peer_address_, invalid_version_, kNoParsedChlo, nullptr);
   store_.EnqueuePacket(connection_id_1, false, packet_, self_address_,
-                       peer_address_, invalid_version_, kNoParsedChlo);
+                       peer_address_, invalid_version_, kNoParsedChlo, nullptr);
 
   ParsedClientHello parsed_chlo;
   parsed_chlo.alpns.push_back("h3");
   parsed_chlo.sni = TestHostname();
   store_.EnqueuePacket(connection_id_2, false, packet_, self_address_,
-                       peer_address_, valid_version_, parsed_chlo);
+                       peer_address_, valid_version_, parsed_chlo, nullptr);
   EXPECT_TRUE(store_.HasBufferedPackets(connection_id_1));
   EXPECT_TRUE(store_.HasBufferedPackets(connection_id_2));
   EXPECT_TRUE(store_.HasChlosBuffered());
@@ -472,7 +543,7 @@
 
   EXPECT_FALSE(store_.HasBufferedPackets(connection_id));
   store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                       peer_address_, valid_version_, kNoParsedChlo);
+                       peer_address_, valid_version_, kNoParsedChlo, nullptr);
   EXPECT_TRUE(store_.HasBufferedPackets(connection_id));
 
   // The packet in 'packet_' is not a TLS CHLO packet.
@@ -492,9 +563,9 @@
   ASSERT_EQ(packets.size(), 2u);
 
   store_.EnqueuePacket(connection_id, false, *packets[0], self_address_,
-                       peer_address_, valid_version_, kNoParsedChlo);
+                       peer_address_, valid_version_, kNoParsedChlo, nullptr);
   store_.EnqueuePacket(connection_id, false, *packets[1], self_address_,
-                       peer_address_, valid_version_, kNoParsedChlo);
+                       peer_address_, valid_version_, kNoParsedChlo, nullptr);
 
   EXPECT_TRUE(store_.HasBufferedPackets(connection_id));
   EXPECT_FALSE(store_.IngestPacketForTlsChloExtraction(
@@ -575,11 +646,11 @@
   EXPECT_NE(long_packet_type, INITIAL);
 
   store_.EnqueuePacket(connection_id, false, packet_, self_address_,
-                       peer_address_, valid_version_, kNoParsedChlo);
+                       peer_address_, valid_version_, kNoParsedChlo, nullptr);
   store_.EnqueuePacket(connection_id, false, *initial_packets[0], self_address_,
-                       peer_address_, valid_version_, kNoParsedChlo);
+                       peer_address_, valid_version_, kNoParsedChlo, nullptr);
   store_.EnqueuePacket(connection_id, false, *initial_packets[1], self_address_,
-                       peer_address_, valid_version_, kNoParsedChlo);
+                       peer_address_, valid_version_, kNoParsedChlo, nullptr);
 
   BufferedPacketList delivered_packets = store_.DeliverPackets(connection_id);
   EXPECT_THAT(delivered_packets.buffered_packets, SizeIs(3));
diff --git a/quiche/quic/core/quic_dispatcher.cc b/quiche/quic/core/quic_dispatcher.cc
index 09ece44..81ea4ef 100644
--- a/quiche/quic/core/quic_dispatcher.cc
+++ b/quiche/quic/core/quic_dispatcher.cc
@@ -1153,9 +1153,10 @@
           << server_connection_id;
       continue;
     }
-    auto session_ptr = QuicDispatcher::CreateSessionFromChlo(
+    auto session_ptr = CreateSessionFromChlo(
         server_connection_id, *packet_list.parsed_chlo, packet_list.version,
-        packets.front().self_address, packets.front().peer_address);
+        packets.front().self_address, packets.front().peer_address,
+        packet_list.connection_id_generator);
     if (session_ptr != nullptr) {
       DeliverPacketsToSession(packets, session_ptr.get());
     }
@@ -1183,11 +1184,13 @@
 }
 
 void QuicDispatcher::BufferEarlyPacket(const ReceivedPacketInfo& packet_info) {
+  // The connection ID generator will only be set for CHLOs, not for early
+  // packets.
   EnqueuePacketResult rs = buffered_packets_.EnqueuePacket(
       packet_info.destination_connection_id,
       packet_info.form != GOOGLE_QUIC_PACKET, packet_info.packet,
       packet_info.self_address, packet_info.peer_address, packet_info.version,
-      /*parsed_chlo=*/absl::nullopt);
+      /*parsed_chlo=*/absl::nullopt, /*connection_id_generator=*/nullptr);
   if (rs != EnqueuePacketResult::SUCCESS) {
     OnBufferPacketFailure(rs, packet_info.destination_connection_id);
   }
@@ -1204,16 +1207,17 @@
         packet_info->destination_connection_id,
         packet_info->form != GOOGLE_QUIC_PACKET, packet_info->packet,
         packet_info->self_address, packet_info->peer_address,
-        packet_info->version, std::move(parsed_chlo));
+        packet_info->version, std::move(parsed_chlo), &ConnectionIdGenerator());
     if (rs != EnqueuePacketResult::SUCCESS) {
       OnBufferPacketFailure(rs, packet_info->destination_connection_id);
     }
     return;
   }
 
-  auto session_ptr = QuicDispatcher::CreateSessionFromChlo(
+  auto session_ptr = CreateSessionFromChlo(
       packet_info->destination_connection_id, parsed_chlo, packet_info->version,
-      packet_info->self_address, packet_info->peer_address);
+      packet_info->self_address, packet_info->peer_address,
+      &ConnectionIdGenerator());
   if (session_ptr == nullptr) {
     return;
   }
@@ -1285,10 +1289,13 @@
 std::shared_ptr<QuicSession> QuicDispatcher::CreateSessionFromChlo(
     const QuicConnectionId original_connection_id,
     const ParsedClientHello& parsed_chlo, const ParsedQuicVersion version,
-    const QuicSocketAddress self_address,
-    const QuicSocketAddress peer_address) {
+    const QuicSocketAddress self_address, const QuicSocketAddress peer_address,
+    ConnectionIdGeneratorInterface* connection_id_generator) {
+  if (connection_id_generator == nullptr) {
+    connection_id_generator = &ConnectionIdGenerator();
+  }
   absl::optional<QuicConnectionId> server_connection_id =
-      connection_id_generator_.MaybeReplaceConnectionId(original_connection_id,
+      connection_id_generator->MaybeReplaceConnectionId(original_connection_id,
                                                         version);
   const bool replaced_connection_id = server_connection_id.has_value();
   if (!replaced_connection_id) {
@@ -1314,7 +1321,7 @@
   std::string alpn = SelectAlpn(parsed_chlo.alpns);
   std::unique_ptr<QuicSession> session =
       CreateQuicSession(*server_connection_id, self_address, peer_address, alpn,
-                        version, parsed_chlo, ConnectionIdGenerator());
+                        version, parsed_chlo, *connection_id_generator);
   if (ABSL_PREDICT_FALSE(session == nullptr)) {
     QUIC_BUG(quic_bug_10287_8)
         << "CreateQuicSession returned nullptr for " << *server_connection_id
diff --git a/quiche/quic/core/quic_dispatcher.h b/quiche/quic/core/quic_dispatcher.h
index 183b335..7224712 100644
--- a/quiche/quic/core/quic_dispatcher.h
+++ b/quiche/quic/core/quic_dispatcher.h
@@ -16,6 +16,7 @@
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/optional.h"
 #include "quiche/quic/core/connection_id_generator.h"
@@ -406,7 +407,8 @@
   std::shared_ptr<QuicSession> CreateSessionFromChlo(
       QuicConnectionId original_connection_id,
       const ParsedClientHello& parsed_chlo, ParsedQuicVersion version,
-      QuicSocketAddress self_address, QuicSocketAddress peer_address);
+      QuicSocketAddress self_address, QuicSocketAddress peer_address,
+      ConnectionIdGeneratorInterface* connection_id_generator);
 
   const QuicConfig* config_;
 
diff --git a/quiche/quic/core/quic_dispatcher_test.cc b/quiche/quic/core/quic_dispatcher_test.cc
index 35b96ff..7ba892c 100644
--- a/quiche/quic/core/quic_dispatcher_test.cc
+++ b/quiche/quic/core/quic_dispatcher_test.cc
@@ -654,6 +654,9 @@
       .WillRepeatedly(ReturnRef(mock_connection_id_generator));
   ConnectionIdGeneratorInterface& expected_generator =
       mock_connection_id_generator;
+  EXPECT_CALL(mock_connection_id_generator,
+              MaybeReplaceConnectionId(TestConnectionId(1), version_))
+      .WillOnce(Return(absl::nullopt));
   EXPECT_CALL(*dispatcher_,
               CreateQuicSession(TestConnectionId(1), _, client_address,
                                 Eq(ExpectedAlpn()), _, MatchParsedClientHello(),
@@ -662,6 +665,7 @@
           dispatcher_.get(), config_, TestConnectionId(1), client_address,
           &mock_helper_, &mock_alarm_factory_, &crypto_config_,
           QuicDispatcherPeer::GetCache(dispatcher_.get()), &session1_))));
+  expect_generator_is_called_ = false;
   ProcessFirstFlight(client_address, TestConnectionId(1));
 }
 
@@ -2798,7 +2802,7 @@
     }
   }
 
-  // Graduately consume buffered CHLOs. The buffered connections should be
+  // Gradually consume buffered CHLOs. The buffered connections should be
   // created but the dropped one shouldn't.
   for (uint64_t conn_id = kMaxNumSessionsToCreate + 1;
        conn_id <= kMaxNumSessionsToCreate + kDefaultMaxConnectionsInStore;
@@ -2840,6 +2844,73 @@
             session1_->connection_id());
 }
 
+TEST_P(BufferedPacketStoreTest,
+       ProcessCHLOsUptoLimitAndBufferWithDifferentConnectionIdGenerator) {
+  if (!GetQuicRestartFlag(quic_request_unroutable_cid)) {
+    return;
+  }
+  // Process (|kMaxNumSessionsToCreate| + 1) CHLOs,
+  // the first |kMaxNumSessionsToCreate| should create connections immediately,
+  // the last should be buffered.
+  QuicBufferedPacketStore* store =
+      QuicDispatcherPeer::GetBufferedPackets(dispatcher_.get());
+  const size_t kNumCHLOs = kMaxNumSessionsToCreate + 1;
+  MockConnectionIdGenerator generator2;
+  for (uint64_t conn_id = 1; conn_id < kNumCHLOs; ++conn_id) {
+    EXPECT_CALL(
+        *dispatcher_,
+        CreateQuicSession(TestConnectionId(conn_id), _, client_addr_,
+                          Eq(ExpectedAlpn()), _, MatchParsedClientHello(), _))
+        .WillOnce(Return(ByMove(CreateSession(
+            dispatcher_.get(), config_, TestConnectionId(conn_id), client_addr_,
+            &mock_helper_, &mock_alarm_factory_, &crypto_config_,
+            QuicDispatcherPeer::GetCache(dispatcher_.get()), &session1_))));
+    EXPECT_CALL(*reinterpret_cast<MockQuicConnection*>(session1_->connection()),
+                ProcessUdpPacket(_, _, _))
+        .WillOnce(WithArg<2>(
+            Invoke([this, conn_id](const QuicEncryptedPacket& packet) {
+              if (version_.UsesQuicCrypto()) {
+                ValidatePacket(TestConnectionId(conn_id), packet);
+              }
+            })));
+    ProcessFirstFlight(TestConnectionId(conn_id));
+  }
+  uint64_t conn_id = kNumCHLOs;
+  expect_generator_is_called_ = false;
+  EXPECT_CALL(*dispatcher_, ConnectionIdGenerator())
+      .WillRepeatedly(ReturnRef(generator2));
+  ProcessFirstFlight(TestConnectionId(conn_id));
+  EXPECT_TRUE(store->HasChloForConnection(TestConnectionId(conn_id)));
+  // Change the generator back so that the session can only access generator2
+  // by using the buffer entry.
+  EXPECT_CALL(*dispatcher_, ConnectionIdGenerator())
+      .WillRepeatedly(ReturnRef(connection_id_generator_));
+
+  // Consume the buffered CHLO. The buffered connection should be
+  // created using generator2.
+  EXPECT_CALL(generator2,
+              MaybeReplaceConnectionId(TestConnectionId(conn_id), version_))
+      .WillOnce(Return(absl::nullopt));
+  EXPECT_CALL(*dispatcher_, CreateQuicSession(TestConnectionId(conn_id), _,
+                                              client_addr_, Eq(ExpectedAlpn()),
+                                              _, MatchParsedClientHello(), _))
+      .WillOnce(Return(ByMove(CreateSession(
+          dispatcher_.get(), config_, TestConnectionId(conn_id), client_addr_,
+          &mock_helper_, &mock_alarm_factory_, &crypto_config_,
+          QuicDispatcherPeer::GetCache(dispatcher_.get()), &session1_))));
+  EXPECT_CALL(*reinterpret_cast<MockQuicConnection*>(session1_->connection()),
+              ProcessUdpPacket(_, _, _))
+      .WillOnce(
+          WithArg<2>(Invoke([this, conn_id](const QuicEncryptedPacket& packet) {
+            if (version_.UsesQuicCrypto()) {
+              ValidatePacket(TestConnectionId(conn_id), packet);
+            }
+          })));
+  while (store->HasChlosBuffered()) {
+    dispatcher_->ProcessBufferedChlos(kMaxNumSessionsToCreate);
+  }
+}
+
 // Duplicated CHLO shouldn't be buffered.
 TEST_P(BufferedPacketStoreTest, BufferDuplicatedCHLO) {
   for (uint64_t conn_id = 1; conn_id <= kMaxNumSessionsToCreate + 1;