nghttpx: Add clear_headers and add_header methods to Request and Response

src/shrpx_mruby_module_request.cc

@@ -184,7 +184,7 @@ mrb_value request_get_headers(mrb_state *mrb, mrb_value self) {
 } // namespace
 
 namespace {
-mrb_value request_set_header(mrb_state *mrb, mrb_value self) {
+mrb_value request_mod_header(mrb_state *mrb, mrb_value self, bool repl) {
   auto data = static_cast<MRubyAssocData *>(mrb->ud);
   auto downstream = data->downstream;
 
@@ -197,6 +197,7 @@ mrb_value request_set_header(mrb_state *mrb, mrb_value self) {
 
   key = mrb_funcall(mrb, key, "downcase", 0);
 
+  if (repl) {
     // making name empty will effectively delete header fields
     for (auto &hd : downstream->get_request_headers()) {
       if (util::streq(std::begin(hd.name), hd.name.size(), RSTRING_PTR(key),
@@ -204,6 +205,7 @@ mrb_value request_set_header(mrb_state *mrb, mrb_value self) {
         hd.name = "";
       }
     }
+  }
 
   if (mrb_obj_is_instance_of(mrb, values, mrb->array_class)) {
     auto n = mrb_ary_len(mrb, values);
@@ -225,6 +227,29 @@ mrb_value request_set_header(mrb_state *mrb, mrb_value self) {
 }
 } // namespace
 
+namespace {
+mrb_value request_set_header(mrb_state *mrb, mrb_value self) {
+  return request_mod_header(mrb, self, true);
+}
+} // namespace
+
+namespace {
+mrb_value request_add_header(mrb_state *mrb, mrb_value self) {
+  return request_mod_header(mrb, self, false);
+}
+} // namespace
+
+namespace {
+mrb_value request_clear_headers(mrb_state *mrb, mrb_value self) {
+  auto data = static_cast<MRubyAssocData *>(mrb->ud);
+  auto downstream = data->downstream;
+
+  downstream->clear_request_headers();
+
+  return mrb_nil_value();
+}
+} // namespace
+
 namespace {
 mrb_value request_get_remote_addr(mrb_state *mrb, mrb_value self) {
   auto data = static_cast<MRubyAssocData *>(mrb->ud);
@@ -266,8 +291,12 @@ void init_request_class(mrb_state *mrb, RClass *module) {
                     MRB_ARGS_REQ(1));
   mrb_define_method(mrb, request_class, "headers", request_get_headers,
                     MRB_ARGS_NONE());
+  mrb_define_method(mrb, request_class, "add_header", request_add_header,
+                    MRB_ARGS_REQ(2));
   mrb_define_method(mrb, request_class, "set_header", request_set_header,
                     MRB_ARGS_REQ(2));
+  mrb_define_method(mrb, request_class, "clear_headers", request_clear_headers,
+                    MRB_ARGS_NONE());
   mrb_define_method(mrb, request_class, "remote_addr", request_get_remote_addr,
                     MRB_ARGS_NONE());
 }

src/shrpx_mruby_module_response.cc

@@ -98,7 +98,7 @@ mrb_value response_get_headers(mrb_state *mrb, mrb_value self) {
 } // namespace
 
 namespace {
-mrb_value response_set_header(mrb_state *mrb, mrb_value self) {
+mrb_value response_mod_header(mrb_state *mrb, mrb_value self, bool repl) {
   auto data = static_cast<MRubyAssocData *>(mrb->ud);
   auto downstream = data->downstream;
 
@@ -111,6 +111,7 @@ mrb_value response_set_header(mrb_state *mrb, mrb_value self) {
 
   key = mrb_funcall(mrb, key, "downcase", 0);
 
+  if (repl) {
     // making name empty will effectively delete header fields
     for (auto &hd : downstream->get_response_headers()) {
       if (util::streq(std::begin(hd.name), hd.name.size(), RSTRING_PTR(key),
@@ -118,6 +119,7 @@ mrb_value response_set_header(mrb_state *mrb, mrb_value self) {
         hd.name = "";
       }
     }
+  }
 
   if (mrb_obj_is_instance_of(mrb, values, mrb->array_class)) {
     auto n = mrb_ary_len(mrb, values);
@@ -139,6 +141,29 @@ mrb_value response_set_header(mrb_state *mrb, mrb_value self) {
 }
 } // namespace
 
+namespace {
+mrb_value response_set_header(mrb_state *mrb, mrb_value self) {
+  return response_mod_header(mrb, self, true);
+}
+} // namespace
+
+namespace {
+mrb_value response_add_header(mrb_state *mrb, mrb_value self) {
+  return response_mod_header(mrb, self, false);
+}
+} // namespace
+
+namespace {
+mrb_value response_clear_headers(mrb_state *mrb, mrb_value self) {
+  auto data = static_cast<MRubyAssocData *>(mrb->ud);
+  auto downstream = data->downstream;
+
+  downstream->clear_response_headers();
+
+  return mrb_nil_value();
+}
+} // namespace
+
 namespace {
 mrb_value response_return(mrb_state *mrb, mrb_value self) {
   auto data = static_cast<MRubyAssocData *>(mrb->ud);
@@ -209,8 +234,12 @@ void init_response_class(mrb_state *mrb, RClass *module) {
                     MRB_ARGS_REQ(1));
   mrb_define_method(mrb, response_class, "headers", response_get_headers,
                     MRB_ARGS_NONE());
+  mrb_define_method(mrb, response_class, "add_header", response_add_header,
+                    MRB_ARGS_REQ(2));
   mrb_define_method(mrb, response_class, "set_header", response_set_header,
                     MRB_ARGS_REQ(2));
+  mrb_define_method(mrb, response_class, "clear_headers",
+                    response_clear_headers, MRB_ARGS_NONE());
   mrb_define_method(mrb, response_class, "return", response_return,
                     MRB_ARGS_OPT(1));
 }