From 94ca9705ef543d6b9150b920d679a60a4fec5dd0 Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Fri, 6 Mar 2015 21:36:40 +0900
Subject: [PATCH] asio: Rewrite response body handling if response body is not
 expected

---
 src/asio_server_http2_handler.cc   | 25 ++++++++++++++-----------
 src/asio_server_request_handler.cc |  7 +++++++
 src/asio_server_response_impl.cc   | 14 +++-----------
 src/http2.cc                       |  8 ++++++++
 src/http2.h                        |  7 +++++++
 5 files changed, 39 insertions(+), 22 deletions(-)

diff --git a/src/asio_server_http2_handler.cc b/src/asio_server_http2_handler.cc
index 48685df9..af7bdce7 100644
--- a/src/asio_server_http2_handler.cc
+++ b/src/asio_server_http2_handler.cc
@@ -325,18 +325,21 @@ int http2_handler::start_response(stream &strm) {
                                           hd.second.sensitive));
   }
 
-  nghttp2_data_provider prd;
-  prd.source.ptr = &strm;
-  prd.read_callback =
-      [](nghttp2_session *session, int32_t stream_id, uint8_t *buf,
-         size_t length, uint32_t *data_flags, nghttp2_data_source *source,
-         void *user_data) -> ssize_t {
-    auto &strm = *static_cast<stream *>(source->ptr);
-    return strm.response().impl().call_read(buf, length, data_flags);
-  };
-
+  nghttp2_data_provider *prd_ptr = nullptr, prd;
+  auto &req = strm.request().impl();
+  if (::nghttp2::http2::expect_response_body(req.method(), res.status_code())) {
+    prd.source.ptr = &strm;
+    prd.read_callback =
+        [](nghttp2_session *session, int32_t stream_id, uint8_t *buf,
+           size_t length, uint32_t *data_flags, nghttp2_data_source *source,
+           void *user_data) -> ssize_t {
+      auto &strm = *static_cast<stream *>(source->ptr);
+      return strm.response().impl().call_read(buf, length, data_flags);
+    };
+    prd_ptr = &prd;
+  }
   rv = nghttp2_submit_response(session_, strm.get_stream_id(), nva.data(),
-                               nva.size(), &prd);
+                               nva.size(), prd_ptr);
 
   if (rv != 0) {
     return -1;
diff --git a/src/asio_server_request_handler.cc b/src/asio_server_request_handler.cc
index 14eec64e..7f113bd4 100644
--- a/src/asio_server_request_handler.cc
+++ b/src/asio_server_request_handler.cc
@@ -62,6 +62,13 @@ request_cb redirect_handler(int status_code, std::string uri) {
 
 request_cb status_handler(int status_code) {
   return [status_code](const request &req, const response &res) {
+    if (!::nghttp2::http2::expect_response_body(status_code)) {
+      res.write_head(status_code);
+      res.end();
+      return;
+    }
+    // we supply content-length for HEAD request, but body will not be
+    // sent.
     auto html = create_html(status_code);
     header_map h;
     h.emplace("content-length", header_value{util::utos(html.size())});
diff --git a/src/asio_server_response_impl.cc b/src/asio_server_response_impl.cc
index 814a988c..e9122dca 100644
--- a/src/asio_server_response_impl.cc
+++ b/src/asio_server_response_impl.cc
@@ -29,6 +29,8 @@
 #include "asio_server_http2_handler.h"
 #include "asio_common.h"
 
+#include "http2.h"
+
 namespace nghttp2 {
 namespace asio_http2 {
 namespace server {
@@ -79,23 +81,13 @@ void response_impl::end(generator_cb cb) {
   state_ = response_state::BODY_STARTED;
 }
 
-namespace {
-bool expect_response_body(const std::string &method, int status_code) {
-  return method != "HEAD" && status_code / 100 != 1 && status_code != 304 &&
-         status_code != 204;
-}
-} // namespace
-
 void response_impl::start_response() {
   auto handler = strm_->handler();
 
   auto &req = strm_->request().impl();
 
-  // if response body is not expected, nullify it so that HEADERS has
-  // END_STREAM flag set.
-  if (!expect_response_body(req.method(), status_code_)) {
+  if (!::nghttp2::http2::expect_response_body(req.method(), status_code_)) {
     state_ = response_state::BODY_STARTED;
-    generator_cb_ = generator_cb();
   }
 
   if (handler->start_response(*strm_) != 0) {
diff --git a/src/http2.cc b/src/http2.cc
index 0aeb0370..0abde9c4 100644
--- a/src/http2.cc
+++ b/src/http2.cc
@@ -1092,6 +1092,14 @@ std::string path_join(const char *base_path, size_t base_pathlen,
   return res;
 }
 
+bool expect_response_body(int status_code) {
+  return status_code / 100 != 1 && status_code != 304 && status_code != 204;
+}
+
+bool expect_response_body(const std::string &method, int status_code) {
+  return method != "HEAD" && expect_response_body(status_code);
+}
+
 } // namespace http2
 
 } // namespace nghttp2
diff --git a/src/http2.h b/src/http2.h
index 067a2f9b..af8fd101 100644
--- a/src/http2.h
+++ b/src/http2.h
@@ -282,6 +282,13 @@ std::string path_join(const char *base_path, size_t base_pathlen,
                       const char *rel_path, size_t rel_pathlen,
                       const char *rel_query, size_t rel_querylen);
 
+// true if response has body, taking into account the request method
+// and status code.
+bool expect_response_body(const std::string &method, int status_code);
+
+// true if response has body, taking into account status code only.
+bool expect_response_body(int status_code);
+
 } // namespace http2
 
 } // namespace nghttp2