Make QuicSimpleServerStream supports multipart response body

net/tools/quic/quic_simple_server_stream.cc

 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "net/tools/quic/quic_simple_server_stream.h"
 
 #include "base/logging.h"
 #include "base/stl_util.h"
 #include "base/strings/string_number_conversions.h"
 #include "base/strings/string_piece.h"
 #include "base/strings/string_split.h"
 #include "net/quic/quic_bug_tracker.h"
 #include "net/quic/quic_flags.h"
 #include "net/quic/quic_spdy_stream.h"
 #include "net/quic/spdy_utils.h"
 #include "net/spdy/spdy_protocol.h"
 #include "net/tools/quic/quic_in_memory_cache.h"
 #include "net/tools/quic/quic_simple_server_session.h"
 
 using base::StringPiece;
 using base::StringToInt;
 using std::string;
 
 namespace net {
 
 QuicSimpleServerStream::QuicSimpleServerStream(QuicStreamId id,
                                                QuicSpdySession* session)
-    : QuicSpdyStream(id, session), content_length_(-1) {}
+    : QuicSpdyStream(id, session),
+      next_state_(INIT),
+      content_length_(-1),
+      body_chunks_consumed_(0) {}
 
 QuicSimpleServerStream::~QuicSimpleServerStream() {}
 
 void QuicSimpleServerStream::OnInitialHeadersComplete(bool fin,
                                                       size_t frame_len) {
   QuicSpdyStream::OnInitialHeadersComplete(fin, frame_len);
   if (!SpdyUtils::ParseHeaders(decompressed_headers().data(),
                                decompressed_headers().length(),
                                &content_length_, &request_headers_)) {
     DVLOG(1) << "Invalid headers";
     SendErrorResponse();
   }
+
+  const QuicInMemoryCache::Response* response =
+      QuicInMemoryCache::GetInstance()->GetResponse(
+          request_headers_[":authority"], request_headers_[":path"]);
+  if (response != nullptr) {
+    response_.set_response_type(response->response_type());
+  }
+
   MarkHeadersConsumed(decompressed_headers().length());
 }
 
 void QuicSimpleServerStream::OnTrailingHeadersComplete(bool fin,
                                                        size_t frame_len) {
   QUIC_BUG << "Server does not support receiving Trailers.";
   SendErrorResponse();
 }
 
 void QuicSimpleServerStream::OnDataAvailable() {
@@ skipping 10 matching lines @@
         static_cast<int>(body_.size()) > content_length_) {
       DVLOG(1) << "Body size (" << body_.size() << ") > content length ("
                << content_length_ << ").";
       SendErrorResponse();
       return;
     }
     MarkConsumed(iov.iov_len);
   }
   if (!sequencer()->IsClosed()) {
     sequencer()->SetUnblocked();
+    if (response_.response_type() == QuicInMemoryCache::BIDIRECTIONAL) {
+      // This will continue send multipart body or trailers.
+      if (next_state_ != NONE)
+        DoLoop();
+    }
     return;
   }
 
   // If the sequencer is closed, then all the body, including the fin, has been
   // consumed.
   OnFinRead();
 
   if (write_side_closed() || fin_buffered()) {
     return;
   }
 
   if (request_headers_.empty()) {
     DVLOG(1) << "Request headers empty.";
     SendErrorResponse();
     return;
   }
 
   if (content_length_ > 0 &&
       content_length_ != static_cast<int>(body_.size())) {
     DVLOG(1) << "Content length (" << content_length_ << ") != body size ("
              << body_.size() << ").";
     SendErrorResponse();
     return;
   }
 
-  SendResponse();
+  if (next_state_ != NONE)
+    DoLoop();
 }
 
 void QuicSimpleServerStream::PushResponse(
     SpdyHeaderBlock push_request_headers) {
   if (id() % 2 != 0) {
     QUIC_BUG << "Client initiated stream shouldn't be used as promised stream.";
     return;
   }
   // Change the stream state to emulate a client request.
   request_headers_ = push_request_headers;
   content_length_ = 0;
   DVLOG(1) << "Stream " << id() << ": Ready to receive server push response.";
 
   // Set as if stream decompresed the headers and received fin.
   QuicSpdyStream::OnInitialHeadersComplete(/*fin=*/true, 0);
 }
 
-void QuicSimpleServerStream::SendResponse() {
+void QuicSimpleServerStream::SendNotFoundResponse() {
+  DVLOG(1) << "Sending not found response for stream " << id();
+
+  SpdyHeaderBlock response_headers;
+  response_headers[":status"] = "404";
+  response_headers["content-length"] =
+      base::IntToString(strlen(kNotFoundResponseBody));
+  response_.set_response_type(QuicInMemoryCache::REGULAR_RESPONSE);
+  response_.set_headers(response_headers);
+  response_.set_body(kNotFoundResponseBody);
+
+  next_state_ = SEND_HEADERS;
+  DoLoop();
+}
+
+void QuicSimpleServerStream::SendErrorResponse() {
+  DVLOG(1) << "Sending error response for stream " << id();
+
+  SpdyHeaderBlock response_headers;
+  response_headers[":status"] = "500";
+  response_headers["content-length"] =
+      base::UintToString(strlen(kErrorResponseBody));
+  response_.set_response_type(QuicInMemoryCache::REGULAR_RESPONSE);
+  response_.set_headers(response_headers);
+  response_.set_body(kErrorResponseBody);
+
+  next_state_ = SEND_HEADERS;
+  DoLoop();
+}
+
+void QuicSimpleServerStream::DoLoop() {
+  DCHECK_NE(NONE, next_state_);
+  do {
+    State state = next_state_;
+    next_state_ = NONE;
+    switch (state) {
+      case INIT:
+        Init();
+        break;
+      case SEND_HEADERS:
+        SendHeaders();
+        break;
+      case SEND_BODY:
+        SendBody();
+        break;
+      case SEND_TRAILERS:
+        SendTrailers();
+        break;
+      default:
+        NOTREACHED();
+    }
+  } while (next_state_ != NONE &&
+           (sequencer()->IsClosed() ||
+            response_.response_type() != QuicInMemoryCache::BIDIRECTIONAL));
+}
+
+void QuicSimpleServerStream::Init() {
   if (!ContainsKey(request_headers_, ":authority") ||
       !ContainsKey(request_headers_, ":path")) {
     DVLOG(1) << "Request headers do not contain :authority or :path.";
     SendErrorResponse();
     return;
   }
 
   // Find response in cache. If not found, send error response.
   const QuicInMemoryCache::Response* response =
       QuicInMemoryCache::GetInstance()->GetResponse(
@@ skipping 54 matching lines @@
            << id();
 
   if (!resources.empty()) {
     QuicSimpleServerSession* session =
         static_cast<QuicSimpleServerSession*>(spdy_session());
     session->PromisePushResources(request_url, resources, id(),
                                   request_headers_);
   }
 
   DVLOG(1) << "Sending response for stream " << id();
-  SendHeadersAndBodyAndTrailers(response->headers(), response->body(),
-                                response->trailers());
+  response_.set_response_type(response->response_type());
+  response_.set_headers(response->headers());
+  response_.set_body(response->body());
+  response_.set_trailers(response->trailers());
+  response_.SetBodyChunks(response->body_chunks());
+  next_state_ = SEND_HEADERS;
 }
 
-void QuicSimpleServerStream::SendNotFoundResponse() {
-  DVLOG(1) << "Sending not found response for stream " << id();
-  SpdyHeaderBlock headers;
-  headers[":status"] = "404";
-  headers["content-length"] = base::IntToString(strlen(kNotFoundResponseBody));
-  SendHeadersAndBody(headers, kNotFoundResponseBody);
-}
-
-void QuicSimpleServerStream::SendErrorResponse() {
-  DVLOG(1) << "Sending error response for stream " << id();
-  SpdyHeaderBlock headers;
-  headers[":status"] = "500";
-  headers["content-length"] = base::UintToString(strlen(kErrorResponseBody));
-  SendHeadersAndBody(headers, kErrorResponseBody);
-}
-
-void QuicSimpleServerStream::SendHeadersAndBody(
-    const SpdyHeaderBlock& response_headers,
-    StringPiece body) {
-  SendHeadersAndBodyAndTrailers(response_headers, body, SpdyHeaderBlock());
-}
-
-void QuicSimpleServerStream::SendHeadersAndBodyAndTrailers(
-    const SpdyHeaderBlock& response_headers,
-    StringPiece body,
-    const SpdyHeaderBlock& response_trailers) {
-  // This server only supports SPDY and HTTP, and neither handles bidirectional
-  // streaming.
-  if (!reading_stopped()) {
-    StopReading();
+void QuicSimpleServerStream::SendHeaders() {
+  // Stop reading if this response is not bidirectional.
+  if (response_.response_type() != QuicInMemoryCache::BIDIRECTIONAL) {
+    if (!reading_stopped()) {
+      StopReading();
+    }
   }
-
   // Send the headers, with a FIN if there's nothing else to send.
-  bool send_fin = (body.empty() && response_trailers.empty());
+  bool send_fin =
+      (response_.body().empty() && response_.body_chunks().empty() &&
+       response_.trailers().empty());
   DVLOG(1) << "Writing headers (fin = " << send_fin
-           << ") : " << response_headers.DebugString();
-  WriteHeaders(response_headers, send_fin, nullptr);
+           << ") : " << response_.headers().DebugString();
+  WriteHeaders(response_.headers(), send_fin, nullptr);
   if (send_fin) {
     // Nothing else to send.
     return;
   }
+  next_state_ = SEND_BODY;
+}
 
-  // Send the body, with a FIN if there's nothing else to send.
-  send_fin = response_trailers.empty();
-  DVLOG(1) << "Writing body (fin = " << send_fin
-           << ") with size: " << body.size();
-  WriteOrBufferData(body, send_fin, nullptr);
-  if (send_fin) {
-    // Nothing else to send.
-    return;
+void QuicSimpleServerStream::SendBody() {
+  if (response_.response_type() == QuicInMemoryCache::BIDIRECTIONAL) {
+    if (response_.body_chunks().size() > body_chunks_consumed_) {
+      WriteOrBufferData(
+          response_.body_chunks().at(body_chunks_consumed_),
+          response_.trailers().empty() &&
+              body_chunks_consumed_ + 1 == response_.body_chunks().size(),
+          nullptr);
+      body_chunks_consumed_++;
+      next_state_ = SEND_BODY;
+    } else if (response_.body_chunks().size() == body_chunks_consumed_) {
+      if (response_.trailers().empty())
+        return;
+      next_state_ = SEND_TRAILERS;
+    }
+  } else {
+    // Send the body, with a FIN if there's nothing else to send.
+    bool send_fin = response_.trailers().empty();
+    DVLOG(1) << "Writing body (fin = " << send_fin
+             << ") with size: " << response_.body().size();
+    WriteOrBufferData(response_.body(), send_fin, nullptr);
+    if (send_fin) {
+      // Nothing else to send.
+      return;
+    }
+    next_state_ = SEND_TRAILERS;
   }
+}
 
+void QuicSimpleServerStream::SendTrailers() {
   // Send the trailers. A FIN is always sent with trailers.
   DVLOG(1) << "Writing trailers (fin = true): "
-           << response_trailers.DebugString();
-  WriteTrailers(response_trailers, nullptr);
+           << response_.trailers().DebugString();
+  WriteTrailers(response_.trailers(), nullptr);
+  next_state_ = NONE;
 }
 
 const char* const QuicSimpleServerStream::kErrorResponseBody = "bad";
 const char* const QuicSimpleServerStream::kNotFoundResponseBody =
     "file not found";
 
 }  // namespace net