From 45e0d42c51bf847be676357d4a817847b285c639 Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Sat, 5 Sep 2015 01:07:29 +0900
Subject: [PATCH] nghttpx: Check allowed operations per phase

---
 src/shrpx_mruby.cc                | 15 +++++++++++----
 src/shrpx_mruby.h                 | 13 ++++++++++++-
 src/shrpx_mruby_module.cc         |  4 ++++
 src/shrpx_mruby_module_request.cc | 12 ++++++++++++
 4 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/src/shrpx_mruby.cc b/src/shrpx_mruby.cc
index 1ca5c0f8..79ba148a 100644
--- a/src/shrpx_mruby.cc
+++ b/src/shrpx_mruby.cc
@@ -44,14 +44,15 @@ MRubyContext::MRubyContext(mrb_state *mrb, RProc *on_request_proc,
 
 MRubyContext::~MRubyContext() { mrb_close(mrb_); }
 
-int MRubyContext::run_request_proc(Downstream *downstream, RProc *proc) {
+int MRubyContext::run_request_proc(Downstream *downstream, RProc *proc,
+                                   int phase) {
   if (!proc || running_) {
     return 0;
   }
 
   running_ = true;
 
-  MRubyAssocData data{downstream};
+  MRubyAssocData data{downstream, phase};
 
   mrb_->ud = &data;
 
@@ -92,11 +93,11 @@ int MRubyContext::run_request_proc(Downstream *downstream, RProc *proc) {
 }
 
 int MRubyContext::run_on_request_proc(Downstream *downstream) {
-  return run_request_proc(downstream, on_request_proc_);
+  return run_request_proc(downstream, on_request_proc_, PHASE_REQUEST);
 }
 
 int MRubyContext::run_on_response_proc(Downstream *downstream) {
-  return run_request_proc(downstream, on_response_proc_);
+  return run_request_proc(downstream, on_response_proc_, PHASE_RESPONSE);
 }
 
 void MRubyContext::delete_downstream(Downstream *downstream) {
@@ -188,6 +189,12 @@ mrb_sym intern_ptr(mrb_state *mrb, void *ptr) {
   return mrb_intern(mrb, reinterpret_cast<const char *>(&p), sizeof(p));
 }
 
+void check_phase(mrb_state *mrb, int phase, int phase_mask) {
+  if ((phase & phase_mask) == 0) {
+    mrb_raise(mrb, E_RUNTIME_ERROR, "operation was not allowed in this phase");
+  }
+}
+
 } // namespace mruby
 
 } // namespace shrpx
diff --git a/src/shrpx_mruby.h b/src/shrpx_mruby.h
index 24a9e6cd..a75e4975 100644
--- a/src/shrpx_mruby.h
+++ b/src/shrpx_mruby.h
@@ -46,7 +46,7 @@ public:
   int run_on_request_proc(Downstream *downstream);
   int run_on_response_proc(Downstream *downstream);
 
-  int run_request_proc(Downstream *downstream, RProc *proc);
+  int run_request_proc(Downstream *downstream, RProc *proc, int phase);
 
   void delete_downstream(Downstream *downstream);
 
@@ -57,8 +57,15 @@ private:
   bool running_;
 };
 
+enum {
+  PHASE_NONE = 0,
+  PHASE_REQUEST = 1,
+  PHASE_RESPONSE = 1 << 1,
+};
+
 struct MRubyAssocData {
   Downstream *downstream;
+  int phase;
   bool request_headers_dirty;
   bool response_headers_dirty;
 };
@@ -70,6 +77,10 @@ std::unique_ptr<MRubyContext> create_mruby_context();
 // Return interned |ptr|.
 mrb_sym intern_ptr(mrb_state *mrb, void *ptr);
 
+// Checks that |phase| is set in |phase_mask|.  If not set, raise
+// exception.
+void check_phase(mrb_state *mrb, int phase, int phase_mask);
+
 } // namespace mruby
 
 } // namespace shrpx
diff --git a/src/shrpx_mruby_module.cc b/src/shrpx_mruby_module.cc
index 16fba07b..3beb2ca6 100644
--- a/src/shrpx_mruby_module.cc
+++ b/src/shrpx_mruby_module.cc
@@ -90,6 +90,10 @@ void init_module(mrb_state *mrb) {
 
   mrb_define_class_method(mrb, module, "run", run,
                           MRB_ARGS_REQ(1) | MRB_ARGS_BLOCK());
+  mrb_define_const(mrb, module, "REQUEST_PHASE",
+                   mrb_fixnum_value(PHASE_REQUEST));
+  mrb_define_const(mrb, module, "RESPONSE_PHASE",
+                   mrb_fixnum_value(PHASE_RESPONSE));
 
   init_env_class(mrb, module);
   init_request_class(mrb, module);
diff --git a/src/shrpx_mruby_module_request.cc b/src/shrpx_mruby_module_request.cc
index eea4f0ee..0b8ead0a 100644
--- a/src/shrpx_mruby_module_request.cc
+++ b/src/shrpx_mruby_module_request.cc
@@ -76,6 +76,8 @@ mrb_value request_set_method(mrb_state *mrb, mrb_value self) {
   auto data = static_cast<MRubyAssocData *>(mrb->ud);
   auto downstream = data->downstream;
 
+  check_phase(mrb, data->phase, PHASE_REQUEST);
+
   const char *method;
   mrb_int n;
   mrb_get_args(mrb, "s", &method, &n);
@@ -109,6 +111,8 @@ mrb_value request_set_authority(mrb_state *mrb, mrb_value self) {
   auto data = static_cast<MRubyAssocData *>(mrb->ud);
   auto downstream = data->downstream;
 
+  check_phase(mrb, data->phase, PHASE_REQUEST);
+
   const char *authority;
   mrb_int n;
   mrb_get_args(mrb, "s", &authority, &n);
@@ -137,6 +141,8 @@ mrb_value request_set_scheme(mrb_state *mrb, mrb_value self) {
   auto data = static_cast<MRubyAssocData *>(mrb->ud);
   auto downstream = data->downstream;
 
+  check_phase(mrb, data->phase, PHASE_REQUEST);
+
   const char *scheme;
   mrb_int n;
   mrb_get_args(mrb, "s", &scheme, &n);
@@ -165,6 +171,8 @@ mrb_value request_set_path(mrb_state *mrb, mrb_value self) {
   auto data = static_cast<MRubyAssocData *>(mrb->ud);
   auto downstream = data->downstream;
 
+  check_phase(mrb, data->phase, PHASE_REQUEST);
+
   const char *path;
   mrb_int pathlen;
   mrb_get_args(mrb, "s", &path, &pathlen);
@@ -188,6 +196,8 @@ mrb_value request_mod_header(mrb_state *mrb, mrb_value self, bool repl) {
   auto data = static_cast<MRubyAssocData *>(mrb->ud);
   auto downstream = data->downstream;
 
+  check_phase(mrb, data->phase, PHASE_REQUEST);
+
   mrb_value key, values;
   mrb_get_args(mrb, "oo", &key, &values);
 
@@ -244,6 +254,8 @@ mrb_value request_clear_headers(mrb_state *mrb, mrb_value self) {
   auto data = static_cast<MRubyAssocData *>(mrb->ud);
   auto downstream = data->downstream;
 
+  check_phase(mrb, data->phase, PHASE_REQUEST);
+
   downstream->clear_request_headers();
 
   return mrb_nil_value();