asio: Rewrite response body handling if response body is not expected

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;

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())});

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) {

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

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