Make the toy QUIC server send responses of arbitrary length based a number specified in the path.

gfe-relnote: n/a - Tools-only code
PiperOrigin-RevId: 274840831
Change-Id: Ia22584f9ae202896f268a47172e15e8d1a9a75f8
diff --git a/quic/tools/quic_backend_response.h b/quic/tools/quic_backend_response.h
index b578915..f0b7b89 100644
--- a/quic/tools/quic_backend_response.h
+++ b/quic/tools/quic_backend_response.h
@@ -42,6 +42,8 @@
     STOP_SENDING,          // Acts like INCOMPLETE_RESPONSE in that the entire
                            // response is not sent. After sending what is sent,
                            // the server will send a STOP_SENDING.
+    GENERATE_BYTES         // Sends a response with a length equal to the number
+                           // of bytes in the URL path.
   };
   QuicBackendResponse();
 
diff --git a/quic/tools/quic_memory_cache_backend.cc b/quic/tools/quic_memory_cache_backend.cc
index fb3dd1f..9fe9dc5 100644
--- a/quic/tools/quic_memory_cache_backend.cc
+++ b/quic/tools/quic_memory_cache_backend.cc
@@ -136,6 +136,15 @@
 
   auto it = responses_.find(GetKey(host, path));
   if (it == responses_.end()) {
+    uint64_t ignored = 0;
+    if (generate_bytes_response_) {
+      if (QuicTextUtils::StringToUint64(
+              QuicStringPiece(&path[1], path.size() - 1), &ignored)) {
+        // The actual parsed length is ignored here and will be recomputed
+        // by the caller.
+        return generate_bytes_response_.get();
+      }
+    }
     QUIC_DVLOG(1) << "Get response for resource failed: host " << host
                   << " path " << path;
     if (default_response_) {
@@ -271,10 +280,23 @@
     MaybeAddServerPushResources(resource_file->host(), resource_file->path(),
                                 push_resources);
   }
+
   cache_initialized_ = true;
   return true;
 }
 
+void QuicMemoryCacheBackend::GenerateDynamicResponses() {
+  QuicWriterMutexLock lock(&response_mutex_);
+  // Add a generate bytes response.
+  spdy::SpdyHeaderBlock response_headers;
+  response_headers[":version"] = "HTTP/1.1";
+  response_headers[":status"] = "200";
+  generate_bytes_response_ = std::make_unique<QuicBackendResponse>();
+  generate_bytes_response_->set_headers(std::move(response_headers));
+  generate_bytes_response_->set_response_type(
+      QuicBackendResponse::GENERATE_BYTES);
+}
+
 bool QuicMemoryCacheBackend::IsBackendInitialized() const {
   return cache_initialized_;
 }
diff --git a/quic/tools/quic_memory_cache_backend.h b/quic/tools/quic_memory_cache_backend.h
index 8e202c0..e0ff400 100644
--- a/quic/tools/quic_memory_cache_backend.h
+++ b/quic/tools/quic_memory_cache_backend.h
@@ -137,6 +137,10 @@
   // |cache_cirectory| can be generated using `wget -p --save-headers <url>`.
   void InitializeFromDirectory(const std::string& cache_directory);
 
+  // Once called, URLs which have a numeric path will send a dynamically
+  // generated response of that many bytes.
+  void GenerateDynamicResponses();
+
   // Find all the server push resources associated with |request_url|.
   std::list<QuicBackendResponse::ServerPushInfo> GetServerPushResources(
       std::string request_url);
@@ -183,6 +187,10 @@
   std::unique_ptr<QuicBackendResponse> default_response_
       QUIC_GUARDED_BY(response_mutex_);
 
+  // The generate bytes response, if set.
+  std::unique_ptr<QuicBackendResponse> generate_bytes_response_
+      QUIC_GUARDED_BY(response_mutex_);
+
   // A map from request URL to associated server push responses (if any).
   std::multimap<std::string, QuicBackendResponse::ServerPushInfo>
       server_push_resources_ QUIC_GUARDED_BY(response_mutex_);
diff --git a/quic/tools/quic_simple_server_stream.cc b/quic/tools/quic_simple_server_stream.cc
index 5badfe1..f821a5d 100644
--- a/quic/tools/quic_simple_server_stream.cc
+++ b/quic/tools/quic_simple_server_stream.cc
@@ -29,6 +29,7 @@
     QuicSimpleServerBackend* quic_simple_server_backend)
     : QuicSpdyServerStreamBase(id, session, type),
       content_length_(-1),
+      generate_bytes_length_(0),
       quic_simple_server_backend_(quic_simple_server_backend) {
   DCHECK(quic_simple_server_backend_);
 }
@@ -40,6 +41,7 @@
     QuicSimpleServerBackend* quic_simple_server_backend)
     : QuicSpdyServerStreamBase(pending, session, type),
       content_length_(-1),
+      generate_bytes_length_(0),
       quic_simple_server_backend_(quic_simple_server_backend) {
   DCHECK(quic_simple_server_backend_);
 }
@@ -260,11 +262,46 @@
     return;
   }
 
+  if (response->response_type() == QuicBackendResponse::GENERATE_BYTES) {
+    QUIC_DVLOG(1) << "Stream " << id() << " sending a generate bytes response.";
+    std::string path = request_headers_[":path"].as_string().substr(1);
+    if (!QuicTextUtils::StringToUint64(path, &generate_bytes_length_)) {
+      QUIC_LOG(ERROR) << "Path is not a number.";
+      SendNotFoundResponse();
+      return;
+    }
+    SpdyHeaderBlock headers = response->headers().Clone();
+    headers["content-length"] =
+        QuicTextUtils::Uint64ToString(generate_bytes_length_);
+
+    WriteHeaders(std::move(headers), false, nullptr);
+
+    WriteGeneratedBytes();
+
+    return;
+  }
+
   QUIC_DVLOG(1) << "Stream " << id() << " sending response.";
   SendHeadersAndBodyAndTrailers(response->headers().Clone(), response->body(),
                                 response->trailers().Clone());
 }
 
+void QuicSimpleServerStream::OnCanWrite() {
+  QuicSpdyStream::OnCanWrite();
+  WriteGeneratedBytes();
+}
+
+void QuicSimpleServerStream::WriteGeneratedBytes() {
+  static size_t kChunkSize = 1024;
+  while (!HasBufferedData() && generate_bytes_length_ > 0) {
+    size_t len = std::min<size_t>(kChunkSize, generate_bytes_length_);
+    std::string data(len, 'a');
+    generate_bytes_length_ -= len;
+    bool fin = generate_bytes_length_ == 0;
+    WriteOrBufferBody(data, fin);
+  }
+}
+
 void QuicSimpleServerStream::SendNotFoundResponse() {
   QUIC_DVLOG(1) << "Stream " << id() << " sending not found response.";
   SpdyHeaderBlock headers;
diff --git a/quic/tools/quic_simple_server_stream.h b/quic/tools/quic_simple_server_stream.h
index ace1571..594ca02 100644
--- a/quic/tools/quic_simple_server_stream.h
+++ b/quic/tools/quic_simple_server_stream.h
@@ -38,6 +38,7 @@
   void OnTrailingHeadersComplete(bool fin,
                                  size_t frame_len,
                                  const QuicHeaderList& header_list) override;
+  void OnCanWrite() override;
 
   // QuicStream implementation called by the sequencer when there is
   // data (or a FIN) to be read.
@@ -88,12 +89,17 @@
 
   const std::string& body() { return body_; }
 
+  // Writes the body bytes for the GENERATE_BYTES response type.
+  void WriteGeneratedBytes();
+
   // The parsed headers received from the client.
   spdy::SpdyHeaderBlock request_headers_;
   int64_t content_length_;
   std::string body_;
 
  private:
+  uint64_t generate_bytes_length_;
+
   QuicSimpleServerBackend* quic_simple_server_backend_;  // Not owned.
 };
 
diff --git a/quic/tools/quic_toy_server.cc b/quic/tools/quic_toy_server.cc
index 0b0cd3d..fdb1bb6 100644
--- a/quic/tools/quic_toy_server.cc
+++ b/quic/tools/quic_toy_server.cc
@@ -26,6 +26,13 @@
     "construction to seed the cache. Cache directory can be "
     "generated using `wget -p --save-headers <url>`");
 
+DEFINE_QUIC_COMMAND_LINE_FLAG(
+    bool,
+    generate_dynamic_responses,
+    false,
+    "If true, then URLs which have a numeric path will send a dynamically "
+    "generated response of that many bytes.");
+
 DEFINE_QUIC_COMMAND_LINE_FLAG(bool,
                               quic_ietf_draft,
                               false,
@@ -37,6 +44,9 @@
 std::unique_ptr<quic::QuicSimpleServerBackend>
 QuicToyServer::MemoryCacheBackendFactory::CreateBackend() {
   auto memory_cache_backend = std::make_unique<QuicMemoryCacheBackend>();
+  if (GetQuicFlag(FLAGS_generate_dynamic_responses)) {
+    memory_cache_backend->GenerateDynamicResponses();
+  }
   if (!GetQuicFlag(FLAGS_quic_response_cache_dir).empty()) {
     memory_cache_backend->InitializeBackend(
         GetQuicFlag(FLAGS_quic_response_cache_dir));