From 9a85c5264aeba2c7156afab814c602bbe4ad86da Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Thu, 3 Nov 2016 21:29:07 +0900
Subject: [PATCH] nghttpx: Send 1xx non-final response using mruby script

---
 doc/nghttpx.h2r                    | 12 +++++
 src/shrpx_mruby_module_response.cc | 87 ++++++++++++++++++++++++++++++
 2 files changed, 99 insertions(+)

diff --git a/doc/nghttpx.h2r b/doc/nghttpx.h2r
index abf1fb7f..761735ca 100644
--- a/doc/nghttpx.h2r
+++ b/doc/nghttpx.h2r
@@ -482,6 +482,18 @@ respectively.
         existing header fields, and then add required header fields.
         It is an error to call this method twice for a given request.
 
+    .. rb:method:: send_info(status, headers)
+
+        Send non-final (informational) response to a client.  *status*
+        must be in the range [100, 199], inclusive.  *headers* is a
+        hash containing response header fields.  Its key must be a
+        string, and the associated value must be either string or
+        array of strings.  Since this is not a final response, even if
+        this method is invoked, request is still forwarded to a
+        backend unless :rb:meth:`Nghttpx::Response#return` is called.
+        This method can be called multiple times.  It cannot be called
+        after :rb:meth:`Nghttpx::Response#return` is called.
+
 MRUBY EXAMPLES
 ~~~~~~~~~~~~~~
 
diff --git a/src/shrpx_mruby_module_response.cc b/src/shrpx_mruby_module_response.cc
index 13b7192a..3088dbb6 100644
--- a/src/shrpx_mruby_module_response.cc
+++ b/src/shrpx_mruby_module_response.cc
@@ -274,6 +274,91 @@ mrb_value response_return(mrb_state *mrb, mrb_value self) {
 }
 } // namespace
 
+namespace {
+mrb_value response_send_info(mrb_state *mrb, mrb_value self) {
+  auto data = static_cast<MRubyAssocData *>(mrb->ud);
+  auto downstream = data->downstream;
+  auto &resp = downstream->response();
+  int rv;
+
+  if (downstream->get_response_state() == Downstream::MSG_COMPLETE) {
+    mrb_raise(mrb, E_RUNTIME_ERROR, "response has already been committed");
+  }
+
+  mrb_int http_status;
+  mrb_value hash;
+  mrb_get_args(mrb, "iH", &http_status, &hash);
+
+  if (http_status / 100 != 1) {
+    mrb_raise(mrb, E_RUNTIME_ERROR,
+              "status_code must be in range [100, 199], inclusive");
+  }
+
+  auto &balloc = downstream->get_block_allocator();
+
+  auto keys = mrb_hash_keys(mrb, hash);
+  auto keyslen = mrb_ary_len(mrb, keys);
+
+  for (int i = 0; i < keyslen; ++i) {
+    auto key = mrb_ary_ref(mrb, keys, i);
+    if (!mrb_string_p(key)) {
+      mrb_raise(mrb, E_RUNTIME_ERROR, "key must be string");
+    }
+
+    auto values = mrb_hash_get(mrb, hash, key);
+
+    auto ai = mrb_gc_arena_save(mrb);
+
+    key = mrb_funcall(mrb, key, "downcase", 0);
+
+    auto keyref = make_string_ref(
+        balloc,
+        StringRef{RSTRING_PTR(key), static_cast<size_t>(RSTRING_LEN(key))});
+
+    mrb_gc_arena_restore(mrb, ai);
+
+    auto token = http2::lookup_token(keyref.byte(), keyref.size());
+
+    if (mrb_array_p(values)) {
+      auto n = mrb_ary_len(mrb, values);
+      for (int i = 0; i < n; ++i) {
+        auto value = mrb_ary_ref(mrb, values, i);
+        if (!mrb_string_p(value)) {
+          mrb_raise(mrb, E_RUNTIME_ERROR, "value must be string");
+        }
+
+        resp.fs.add_header_token(
+            keyref,
+            make_string_ref(balloc,
+                            StringRef{RSTRING_PTR(value),
+                                      static_cast<size_t>(RSTRING_LEN(value))}),
+            false, token);
+      }
+    } else if (mrb_string_p(values)) {
+      resp.fs.add_header_token(
+          keyref,
+          make_string_ref(balloc,
+                          StringRef{RSTRING_PTR(values),
+                                    static_cast<size_t>(RSTRING_LEN(values))}),
+          false, token);
+    } else {
+      mrb_raise(mrb, E_RUNTIME_ERROR, "value must be string");
+    }
+  }
+
+  resp.http_status = http_status;
+
+  auto upstream = downstream->get_upstream();
+
+  rv = upstream->on_downstream_header_complete(downstream);
+  if (rv != 0) {
+    mrb_raise(mrb, E_RUNTIME_ERROR, "could not send non-final response");
+  }
+
+  return self;
+}
+} // namespace
+
 void init_response_class(mrb_state *mrb, RClass *module) {
   auto response_class =
       mrb_define_class_under(mrb, module, "Response", mrb->object_class);
@@ -298,6 +383,8 @@ void init_response_class(mrb_state *mrb, RClass *module) {
                     response_clear_headers, MRB_ARGS_NONE());
   mrb_define_method(mrb, response_class, "return", response_return,
                     MRB_ARGS_OPT(1));
+  mrb_define_method(mrb, response_class, "send_info", response_send_info,
+                    MRB_ARGS_REQ(2));
 }
 
 } // namespace mruby