Make QuicSimpleServerStream supports multipart response body

net/tools/quic/quic_simple_server_stream.cc

Index: net/tools/quic/quic_simple_server_stream.cc
diff --git a/net/tools/quic/quic_simple_server_stream.cc b/net/tools/quic/quic_simple_server_stream.cc
index 8609999935b17f636347225b8e6cf126801ef9a3..11c310e7e2a3a556a4ff818a1c5dd0bd1a0e617b 100644
--- a/net/tools/quic/quic_simple_server_stream.cc
+++ b/net/tools/quic/quic_simple_server_stream.cc
@@ -25,7 +25,10 @@ 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() {}
 
@@ -38,6 +41,14 @@ void QuicSimpleServerStream::OnInitialHeadersComplete(bool fin,
     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());
 }
 
@@ -68,6 +79,11 @@ void QuicSimpleServerStream::OnDataAvailable() {
   }
   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;
   }
 
@@ -93,7 +109,8 @@ void QuicSimpleServerStream::OnDataAvailable() {
     return;
   }
 
-  SendResponse();
+  if (next_state_ != NONE)
+    DoLoop();
 }
 
 void QuicSimpleServerStream::PushResponse(
@@ -111,7 +128,63 @@ void QuicSimpleServerStream::PushResponse(
   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.";
@@ -186,66 +259,70 @@ void QuicSimpleServerStream::SendResponse() {
   }
 
   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";