Internal QUICHE change

PiperOrigin-RevId: 259859589
Change-Id: I33f879f5422c0fad5d3570dbb48b82bac23cefff
diff --git a/quic/test_tools/quic_test_client.cc b/quic/test_tools/quic_test_client.cc
index e2e33c7..ef92cb0 100644
--- a/quic/test_tools/quic_test_client.cc
+++ b/quic/test_tools/quic_test_client.cc
@@ -233,9 +233,14 @@
 }
 
 QuicConnectionId MockableQuicClient::GenerateNewConnectionId() {
-  return server_connection_id_overridden_
-             ? override_server_connection_id_
-             : QuicClient::GenerateNewConnectionId();
+  if (server_connection_id_overridden_) {
+    return override_server_connection_id_;
+  }
+  if (override_server_connection_id_length_ >= 0) {
+    return QuicUtils::CreateRandomConnectionId(
+        override_server_connection_id_length_);
+  }
+  return QuicClient::GenerateNewConnectionId();
 }
 
 void MockableQuicClient::UseConnectionId(
@@ -244,9 +249,20 @@
   override_server_connection_id_ = server_connection_id;
 }
 
+void MockableQuicClient::UseConnectionIdLength(
+    int server_connection_id_length) {
+  override_server_connection_id_length_ = server_connection_id_length;
+}
+
 QuicConnectionId MockableQuicClient::GetClientConnectionId() {
-  return client_connection_id_overridden_ ? override_client_connection_id_
-                                          : QuicClient::GetClientConnectionId();
+  if (client_connection_id_overridden_) {
+    return override_client_connection_id_;
+  }
+  if (override_client_connection_id_length_ >= 0) {
+    return QuicUtils::CreateRandomConnectionId(
+        override_client_connection_id_length_);
+  }
+  return QuicClient::GetClientConnectionId();
 }
 
 void MockableQuicClient::UseClientConnectionId(
@@ -255,6 +271,11 @@
   override_client_connection_id_ = client_connection_id;
 }
 
+void MockableQuicClient::UseClientConnectionIdLength(
+    int client_connection_id_length) {
+  override_client_connection_id_length_ = client_connection_id_length;
+}
+
 void MockableQuicClient::UseWriter(QuicPacketWriterWrapper* writer) {
   mockable_network_helper()->UseWriter(writer);
 }
@@ -769,12 +790,23 @@
   client_->UseConnectionId(server_connection_id);
 }
 
+void QuicTestClient::UseConnectionIdLength(int server_connection_id_length) {
+  DCHECK(!connected());
+  client_->UseConnectionIdLength(server_connection_id_length);
+}
+
 void QuicTestClient::UseClientConnectionId(
     QuicConnectionId client_connection_id) {
   DCHECK(!connected());
   client_->UseClientConnectionId(client_connection_id);
 }
 
+void QuicTestClient::UseClientConnectionIdLength(
+    int client_connection_id_length) {
+  DCHECK(!connected());
+  client_->UseClientConnectionIdLength(client_connection_id_length);
+}
+
 bool QuicTestClient::MigrateSocket(const QuicIpAddress& new_host) {
   return client_->MigrateSocket(new_host);
 }
diff --git a/quic/test_tools/quic_test_client.h b/quic/test_tools/quic_test_client.h
index 7b83562..cd5b09f 100644
--- a/quic/test_tools/quic_test_client.h
+++ b/quic/test_tools/quic_test_client.h
@@ -56,8 +56,10 @@
 
   QuicConnectionId GenerateNewConnectionId() override;
   void UseConnectionId(QuicConnectionId server_connection_id);
+  void UseConnectionIdLength(int server_connection_id_length);
   QuicConnectionId GetClientConnectionId() override;
   void UseClientConnectionId(QuicConnectionId client_connection_id);
+  void UseClientConnectionIdLength(int client_connection_id_length);
 
   void UseWriter(QuicPacketWriterWrapper* writer);
   void set_peer_address(const QuicSocketAddress& address);
@@ -74,9 +76,11 @@
   // Server connection ID to use, if server_connection_id_overridden_
   QuicConnectionId override_server_connection_id_;
   bool server_connection_id_overridden_;
+  int override_server_connection_id_length_ = -1;
   // Client connection ID to use, if client_connection_id_overridden_
   QuicConnectionId override_client_connection_id_;
   bool client_connection_id_overridden_;
+  int override_client_connection_id_length_ = -1;
   CachedNetworkParameters cached_network_paramaters_;
 };
 
@@ -227,9 +231,15 @@
   // Configures client_ to use a specific server connection ID instead of a
   // random one.
   void UseConnectionId(QuicConnectionId server_connection_id);
+  // Configures client_ to use a specific server connection ID length instead
+  // of the default of kQuicDefaultConnectionIdLength.
+  void UseConnectionIdLength(int server_connection_id_length);
   // Configures client_ to use a specific client connection ID instead of an
   // empty one.
   void UseClientConnectionId(QuicConnectionId client_connection_id);
+  // Configures client_ to use a specific client connection ID length instead
+  // of the default of zero.
+  void UseClientConnectionIdLength(int client_connection_id_length);
 
   // Returns nullptr if the maximum number of streams have already been created.
   QuicSpdyClientStream* GetOrCreateStream();