nghttpx: More freedom for custom response headers

src/shrpx_http2_upstream.cc

@@ -1138,6 +1138,61 @@ ssize_t downstream_data_read_callback(nghttp2_session *session,
 }
 } // namespace
 
+int Http2Upstream::send_reply(Downstream *downstream, const uint8_t *body,
+                              size_t bodylen) {
+  int rv;
+
+  nghttp2_data_provider data_prd, *data_prd_ptr = nullptr;
+
+  if (bodylen) {
+    data_prd.source.ptr = downstream;
+    data_prd.read_callback = downstream_data_read_callback;
+    data_prd_ptr = &data_prd;
+  }
+
+  auto status_code_str = util::utos(downstream->get_response_http_status());
+  auto &headers = downstream->get_response_headers();
+  // 2 for :status and server
+  auto nva = std::vector<nghttp2_nv>(2 + headers.size());
+  nva.push_back(http2::make_nv_ls(":status", status_code_str));
+
+  for (auto &kv : headers) {
+    if (kv.name.empty() || kv.name[0] == ':') {
+      continue;
+    }
+    switch (kv.token) {
+    case http2::HD_CONNECTION:
+    case http2::HD_KEEP_ALIVE:
+    case http2::HD_PROXY_CONNECTION:
+    case http2::HD_TE:
+    case http2::HD_TRANSFER_ENCODING:
+    case http2::HD_UPGRADE:
+      continue;
+    }
+    nva.push_back(http2::make_nv(kv.name, kv.value, kv.no_index));
+  }
+
+  if (!downstream->get_response_header(http2::HD_SERVER)) {
+    nva.push_back(http2::make_nv_lc("server", get_config()->server_name));
+  }
+
+  rv = nghttp2_submit_response(session_, downstream->get_stream_id(),
+                               nva.data(), nva.size(), data_prd_ptr);
+  if (nghttp2_is_fatal(rv)) {
+    ULOG(FATAL, this) << "nghttp2_submit_response() failed: "
+                      << nghttp2_strerror(rv);
+    return -1;
+  }
+
+  auto buf = downstream->get_response_buf();
+
+  buf->append(body, bodylen);
+
+  downstream->set_response_state(Downstream::MSG_COMPLETE);
+
+  return 0;
+}
+
 int Http2Upstream::error_reply(Downstream *downstream,
                                unsigned int status_code) {
   int rv;

src/shrpx_http2_upstream.h

@@ -78,6 +78,8 @@ public:
 
   virtual void on_handler_delete();
   virtual int on_downstream_reset(bool no_retry);
+  virtual int send_reply(Downstream *downstream, const uint8_t *body,
+                         size_t bodylen);
 
   bool get_flow_control() const;
   // Perform HTTP/2 upgrade from |upstream|. On success, this object

src/shrpx_https_upstream.cc

@@ -725,6 +725,62 @@ int HttpsUpstream::downstream_error(DownstreamConnection *dconn, int events) {
   return 0;
 }
 
+int HttpsUpstream::send_reply(Downstream *downstream, const uint8_t *body,
+                              size_t bodylen) {
+  auto major = downstream->get_request_major();
+  auto minor = downstream->get_request_minor();
+
+  auto connection_close = false;
+  if (major <= 0 || (major == 1 && minor == 0)) {
+    connection_close = true;
+  } else {
+    auto c = downstream->get_response_header(http2::HD_CONNECTION);
+    if (c && util::strieq_l("close", c->value)) {
+      connection_close = true;
+    }
+  }
+
+  if (connection_close) {
+    downstream->set_response_connection_close(true);
+    handler_->set_should_close_after_write(true);
+  }
+
+  auto output = downstream->get_response_buf();
+
+  output->append("HTTP/1.1 ");
+  auto status_str =
+      http2::get_status_string(downstream->get_response_http_status());
+  output->append(status_str.c_str(), status_str.size());
+
+  for (auto &kv : downstream->get_response_headers()) {
+    if (kv.name.empty() || kv.name[0] == ':') {
+      continue;
+    }
+    auto name = kv.name;
+    http2::capitalize(name, 0);
+    output->append(name.c_str(), name.size());
+    output->append(": ");
+    output->append(kv.value.c_str(), kv.value.size());
+    output->append("\r\n");
+  }
+
+  if (!downstream->get_response_header(http2::HD_SERVER)) {
+    output->append("Server: ");
+    output->append(get_config()->server_name,
+                   strlen(get_config()->server_name));
+    output->append("\r\n");
+  }
+
+  output->append("\r\n");
+
+  output->append(body, bodylen);
+
+  downstream->add_response_sent_bodylen(bodylen);
+  downstream->set_response_state(Downstream::MSG_COMPLETE);
+
+  return 0;
+}
+
 void HttpsUpstream::error_reply(unsigned int status_code) {
   auto html = http::create_error_html(status_code);
   auto downstream = get_downstream();

src/shrpx_https_upstream.h

@@ -74,6 +74,8 @@ public:
 
   virtual void on_handler_delete();
   virtual int on_downstream_reset(bool no_retry);
+  virtual int send_reply(Downstream *downstream, const uint8_t *body,
+                         size_t bodylen);
 
   void reset_current_header_length();
   void log_response_headers(const std::string &hdrs) const;

src/shrpx_mruby.cc

@@ -62,7 +62,10 @@ int MRubyContext::run_request_proc(Downstream *downstream, RProc *proc) {
   (void)res;
 
   if (mrb_->exc) {
-    rv = -1;
+    // If response has been committed, ignore error
+    if (downstream->get_response_state() != Downstream::MSG_COMPLETE) {
+      rv = -1;
+    }
     auto error =
         mrb_str_ptr(mrb_funcall(mrb_, mrb_obj_value(mrb_->exc), "inspect", 0));
 

src/shrpx_mruby_module_response.cc

@@ -31,6 +31,7 @@
 
 #include "shrpx_downstream.h"
 #include "shrpx_upstream.h"
+#include "shrpx_client_handler.h"
 #include "shrpx_mruby.h"
 #include "shrpx_mruby_module.h"
 #include "util.h"
@@ -139,7 +140,7 @@ mrb_value response_set_header(mrb_state *mrb, mrb_value self) {
 } // namespace
 
 namespace {
-mrb_value response_end(mrb_state *mrb, mrb_value self) {
+mrb_value response_return(mrb_state *mrb, mrb_value self) {
   auto data = static_cast<MRubyAssocData *>(mrb->ud);
   auto downstream = data->downstream;
   int rv;
@@ -151,14 +152,9 @@ mrb_value response_end(mrb_state *mrb, mrb_value self) {
   mrb_value val;
   mrb_get_args(mrb, "|o", &val);
 
-  const char *body = nullptr;
+  const uint8_t *body = nullptr;
   size_t bodylen = 0;
 
-  if (!mrb_nil_p(val)) {
-    body = RSTRING_PTR(val);
-    bodylen = RSTRING_LEN(val);
-  }
-
   if (downstream->get_response_http_status() == 0) {
     downstream->set_response_http_status(200);
   }
@@ -168,6 +164,11 @@ mrb_value response_end(mrb_state *mrb, mrb_value self) {
     data->response_headers_dirty = false;
   }
 
+  if (downstream->expect_response_body() && !mrb_nil_p(val)) {
+    body = reinterpret_cast<const uint8_t *>(RSTRING_PTR(val));
+    bodylen = RSTRING_LEN(val);
+  }
+
   auto cl = downstream->get_response_header(http2::HD_CONTENT_LENGTH);
   if (cl) {
     cl->value = util::utos(bodylen);
@@ -179,17 +180,14 @@ mrb_value response_end(mrb_state *mrb, mrb_value self) {
 
   auto upstream = downstream->get_upstream();
 
-  rv = upstream->on_downstream_header_complete(downstream);
+  rv = upstream->send_reply(downstream, body, bodylen);
   if (rv != 0) {
     mrb_raise(mrb, E_RUNTIME_ERROR, "could not send response");
   }
 
-  if (downstream->expect_response_body()) {
-    auto output = downstream->get_response_buf();
-    output->append(body, bodylen);
-  }
+  auto handler = upstream->get_client_handler();
 
-  downstream->set_response_state(Downstream::MSG_COMPLETE);
+  handler->signal_write();
 
   return self;
 }
@@ -213,7 +211,8 @@ void init_response_class(mrb_state *mrb, RClass *module) {
                     MRB_ARGS_NONE());
   mrb_define_method(mrb, response_class, "set_header", response_set_header,
                     MRB_ARGS_REQ(2));
-  mrb_define_method(mrb, response_class, "end", response_end, MRB_ARGS_OPT(1));
+  mrb_define_method(mrb, response_class, "return", response_return,
+                    MRB_ARGS_OPT(1));
 }
 
 } // namespace mruby

src/shrpx_spdy_upstream.cc

@@ -793,6 +793,69 @@ ssize_t spdy_data_read_callback(spdylay_session *session, int32_t stream_id,
 }
 } // namespace
 
+int SpdyUpstream::send_reply(Downstream *downstream, const uint8_t *body,
+                             size_t bodylen) {
+  int rv;
+
+  spdylay_data_provider data_prd, *data_prd_ptr = nullptr;
+  if (bodylen) {
+    data_prd.source.ptr = downstream;
+    data_prd.read_callback = spdy_data_read_callback;
+    data_prd_ptr = &data_prd;
+  }
+
+  auto status_string =
+      http2::get_status_string(downstream->get_response_http_status());
+
+  auto &headers = downstream->get_response_headers();
+
+  // 3 for :status, :version and server
+  auto nva = std::vector<const char *>(3 + headers.size());
+
+  nva.push_back(":status");
+  nva.push_back(status_string.c_str());
+  nva.push_back(":version");
+  nva.push_back("HTTP/1.1");
+
+  for (auto &kv : headers) {
+    if (kv.name.empty() || kv.name[0] == ':') {
+      continue;
+    }
+    switch (kv.token) {
+    case http2::HD_CONNECTION:
+    case http2::HD_KEEP_ALIVE:
+    case http2::HD_PROXY_CONNECTION:
+    case http2::HD_TRANSFER_ENCODING:
+      continue;
+    }
+    nva.push_back(kv.name.c_str());
+    nva.push_back(kv.value.c_str());
+  }
+
+  if (!downstream->get_response_header(http2::HD_SERVER)) {
+    nva.push_back("server");
+    nva.push_back(get_config()->server_name);
+  }
+
+  nva.push_back(nullptr);
+
+  rv = spdylay_submit_response(session_, downstream->get_stream_id(),
+                               nva.data(), data_prd_ptr);
+  if (rv < SPDYLAY_ERR_FATAL) {
+    ULOG(FATAL, this) << "spdylay_submit_response() failed: "
+                      << spdylay_strerror(rv);
+    return -1;
+  }
+
+  auto buf = downstream->get_response_buf();
+
+  buf->append(body, bodylen);
+
+  downstream->set_response_state(Downstream::MSG_COMPLETE);
+
+  return 0;
+}
+
 int SpdyUpstream::error_reply(Downstream *downstream,
                               unsigned int status_code) {
   int rv;

src/shrpx_spdy_upstream.h

@@ -74,6 +74,9 @@ public:
   virtual void on_handler_delete();
   virtual int on_downstream_reset(bool no_retry);
 
+  virtual int send_reply(Downstream *downstream, const uint8_t *body,
+                         size_t bodylen);
+
   bool get_flow_control() const;
 
   int consume(int32_t stream_id, size_t len);

src/shrpx_upstream.h

@@ -62,6 +62,8 @@ public:
   virtual void pause_read(IOCtrlReason reason) = 0;
   virtual int resume_read(IOCtrlReason reason, Downstream *downstream,
                           size_t consumed) = 0;
+  virtual int send_reply(Downstream *downstream, const uint8_t *body,
+                         size_t bodylen) = 0;
 };
 
 } // namespace shrpx