tmp

Signed-off-by: mathetake <[email protected]>

Author: mathetake

include/proxy-wasm/context.h

@@ -170,6 +170,11 @@ class ContextBase : public RootInterface,
   // Called before deleting the context.
   virtual void destroy();
 
+  // Called to raise the flag which indicates that the context should stop iteration regardless of
+  // returned filter status from WASM extensions. For example, we ignore
+  // FilterHeadersStatus::Continue after a local reponse is sent by the host.
+  void stopIteration() { stop_iteration_ = true; };
+
   /**
    * Calls into the VM.
    * These are implemented by the proxy-independent host code. They are virtual to support some
@@ -388,6 +393,7 @@ class ContextBase : public RootInterface,
   std::shared_ptr<PluginBase> temp_plugin_; // Remove once ABI v0.1.0 is gone.
   bool in_vm_context_created_ = false;
   bool destroyed_ = false;
+  bool stop_iteration_ = false;
 };
 
 class DeferAfterCallActions {

src/context.cc

@@ -13,16 +13,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "include/proxy-wasm/context.h"
+#include "include/proxy-wasm/wasm.h"
+#include <cassert>
 #include <deque>
 #include <map>
 #include <memory>
 #include <mutex>
 #include <unordered_map>
 #include <unordered_set>
 
-#include "include/proxy-wasm/context.h"
-#include "include/proxy-wasm/wasm.h"
-
 #define CHECK_FAIL(_call, _stream_type, _return_open, _return_closed)                              \
   if (isFailed()) {                                                                                \
     if (plugin_->fail_open_) {                                                                     \
@@ -469,44 +469,62 @@ template <typename P> static uint32_t headerSize(const P &p) { return p ? p->siz
 FilterHeadersStatus ContextBase::onRequestHeaders(uint32_t headers, bool end_of_stream) {
   CHECK_HTTP2(on_request_headers_abi_01_, on_request_headers_abi_02_, FilterHeadersStatus::Continue,
               FilterHeadersStatus::StopIteration);
-  DeferAfterCallActions actions(this);
-  auto result = wasm_->on_request_headers_abi_01_
-                    ? wasm_->on_request_headers_abi_01_(this, id_, headers).u64_
-                    : wasm_
-                          ->on_request_headers_abi_02_(this, id_, headers,
-                                                       static_cast<uint32_t>(end_of_stream))
-                          .u64_;
-  if (result > static_cast<uint64_t>(FilterHeadersStatus::StopAllIterationAndWatermark))
+  uint64_t result;
+  {
+    DeferAfterCallActions actions(this);
+    result = wasm_->on_request_headers_abi_01_
+                 ? wasm_->on_request_headers_abi_01_(this, id_, headers).u64_
+                 : wasm_
+                       ->on_request_headers_abi_02_(this, id_, headers,
+                                                    static_cast<uint32_t>(end_of_stream))
+                       .u64_;
+  }
+
+  if (stop_iteration_) {
+    return FilterHeadersStatus::StopIteration;
+  } else if (result > static_cast<uint64_t>(FilterHeadersStatus::StopAllIterationAndWatermark)) {
     return FilterHeadersStatus::StopAllIterationAndWatermark;
+  }
   return static_cast<FilterHeadersStatus>(result);
 }
 
 FilterDataStatus ContextBase::onRequestBody(uint32_t data_length, bool end_of_stream) {
   CHECK_HTTP(on_request_body_, FilterDataStatus::Continue, FilterDataStatus::StopIterationNoBuffer);
-  DeferAfterCallActions actions(this);
-  auto result =
-      wasm_->on_request_body_(this, id_, data_length, static_cast<uint32_t>(end_of_stream)).u64_;
-  if (result > static_cast<uint64_t>(FilterDataStatus::StopIterationNoBuffer))
+  uint64_t result;
+  {
+    DeferAfterCallActions actions(this);
+    result =
+        wasm_->on_request_body_(this, id_, data_length, static_cast<uint32_t>(end_of_stream)).u64_;
+  }
+  if (stop_iteration_ || result > static_cast<uint64_t>(FilterDataStatus::StopIterationNoBuffer))
     return FilterDataStatus::StopIterationNoBuffer;
   return static_cast<FilterDataStatus>(result);
 }
 
 FilterTrailersStatus ContextBase::onRequestTrailers(uint32_t trailers) {
   CHECK_HTTP(on_request_trailers_, FilterTrailersStatus::Continue,
              FilterTrailersStatus::StopIteration);
-  DeferAfterCallActions actions(this);
-  if (static_cast<FilterTrailersStatus>(wasm_->on_request_trailers_(this, id_, trailers).u64_) ==
-      FilterTrailersStatus::Continue) {
+
+  uint64_t result;
+  {
+    DeferAfterCallActions actions(this);
+    result = wasm_->on_request_trailers_(this, id_, trailers).u64_;
+  }
+  if (!stop_iteration_ &&
+      static_cast<FilterTrailersStatus>(result) == FilterTrailersStatus::Continue) {
     return FilterTrailersStatus::Continue;
   }
   return FilterTrailersStatus::StopIteration;
 }
 
 FilterMetadataStatus ContextBase::onRequestMetadata(uint32_t elements) {
   CHECK_HTTP(on_request_metadata_, FilterMetadataStatus::Continue, FilterMetadataStatus::Continue);
-  DeferAfterCallActions actions(this);
-  if (static_cast<FilterMetadataStatus>(wasm_->on_request_metadata_(this, id_, elements).u64_) ==
-      FilterMetadataStatus::Continue) {
+  uint64_t result;
+  {
+    DeferAfterCallActions actions(this);
+    result = wasm_->on_request_metadata_(this, id_, elements).u64_;
+  }
+  if (static_cast<FilterMetadataStatus>(result) == FilterMetadataStatus::Continue) {
     return FilterMetadataStatus::Continue;
   }
   return FilterMetadataStatus::Continue; // This is currently the only return code.
@@ -515,35 +533,53 @@ FilterMetadataStatus ContextBase::onRequestMetadata(uint32_t elements) {
 FilterHeadersStatus ContextBase::onResponseHeaders(uint32_t headers, bool end_of_stream) {
   CHECK_HTTP2(on_response_headers_abi_01_, on_response_headers_abi_02_,
               FilterHeadersStatus::Continue, FilterHeadersStatus::StopIteration);
-  DeferAfterCallActions actions(this);
-  auto result = wasm_->on_response_headers_abi_01_
-                    ? wasm_->on_response_headers_abi_01_(this, id_, headers).u64_
-                    : wasm_
-                          ->on_response_headers_abi_02_(this, id_, headers,
-                                                        static_cast<uint32_t>(end_of_stream))
-                          .u64_;
-  if (result > static_cast<uint64_t>(FilterHeadersStatus::StopAllIterationAndWatermark))
+
+  uint64_t result;
+  {
+    DeferAfterCallActions actions(this);
+    result = wasm_->on_response_headers_abi_01_
+                 ? wasm_->on_response_headers_abi_01_(this, id_, headers).u64_
+                 : wasm_
+                       ->on_response_headers_abi_02_(this, id_, headers,
+                                                     static_cast<uint32_t>(end_of_stream))
+                       .u64_;
+  }
+
+  if (stop_iteration_) {
+    return FilterHeadersStatus::StopIteration;
+  } else if (result > static_cast<uint64_t>(FilterHeadersStatus::StopAllIterationAndWatermark)) {
     return FilterHeadersStatus::StopAllIterationAndWatermark;
+  }
   return static_cast<FilterHeadersStatus>(result);
 }
 
 FilterDataStatus ContextBase::onResponseBody(uint32_t body_length, bool end_of_stream) {
   CHECK_HTTP(on_response_body_, FilterDataStatus::Continue,
              FilterDataStatus::StopIterationNoBuffer);
-  DeferAfterCallActions actions(this);
-  auto result =
-      wasm_->on_response_body_(this, id_, body_length, static_cast<uint32_t>(end_of_stream)).u64_;
-  if (result > static_cast<uint64_t>(FilterDataStatus::StopIterationNoBuffer))
+
+  uint64_t result;
+  {
+
+    DeferAfterCallActions actions(this);
+    result =
+        wasm_->on_response_body_(this, id_, body_length, static_cast<uint32_t>(end_of_stream)).u64_;
+  }
+
+  if (stop_iteration_ || result > static_cast<uint64_t>(FilterDataStatus::StopIterationNoBuffer))
     return FilterDataStatus::StopIterationNoBuffer;
   return static_cast<FilterDataStatus>(result);
 }
 
 FilterTrailersStatus ContextBase::onResponseTrailers(uint32_t trailers) {
   CHECK_HTTP(on_response_trailers_, FilterTrailersStatus::Continue,
              FilterTrailersStatus::StopIteration);
-  DeferAfterCallActions actions(this);
-  if (static_cast<FilterTrailersStatus>(wasm_->on_response_trailers_(this, id_, trailers).u64_) ==
-      FilterTrailersStatus::Continue) {
+  uint64_t result;
+  {
+    DeferAfterCallActions actions(this);
+    result = wasm_->on_response_trailers_(this, id_, trailers).u64_;
+  }
+  if (!stop_iteration_ &&
+      static_cast<FilterTrailersStatus>(result) == FilterTrailersStatus::Continue) {
     return FilterTrailersStatus::Continue;
   }
   return FilterTrailersStatus::StopIteration;

src/exports.cc

@@ -186,6 +186,7 @@ Word send_local_response(void *raw_context, Word response_code, Word response_co
   auto additional_headers = toPairs(additional_response_header_pairs.value());
   context->sendLocalResponse(response_code, body.value(), std::move(additional_headers), grpc_code,
                              details.value());
+  context->wasm()->addAfterVmCallAction([context] { context->stopIteration(); });
   return WasmResult::Ok;
 }