fix sendLocalResponse in wasm

source/extensions/common/wasm/context.cc

@@ -1638,6 +1638,12 @@ void Context::failStream(WasmStreamType stream_type) {
 WasmResult Context::sendLocalResponse(uint32_t response_code, std::string_view body_text,
                                       Pairs additional_headers, uint32_t grpc_status,
                                       std::string_view details) {
+  // This flag is used to avoid calling sendLocalReply() twice, even if wasm code has this
+  // logic. We can't reuse "local_reply_sent_" here because it can't avoid calling nested
+  // sendLocalReply() during encodeHeaders().
+  if (local_reply_hold_) {
+    return WasmResult::BadArgument;
+  }
   // "additional_headers" is a collection of string_views. These will no longer
   // be valid when "modify_headers" is finally called below, so we must
   // make copies of all the headers.
@@ -1662,10 +1668,17 @@ WasmResult Context::sendLocalResponse(uint32_t response_code, std::string_view b
                                   modify_headers = std::move(modify_headers), grpc_status,
                                   details = StringUtil::replaceAllEmptySpace(
                                       absl::string_view(details.data(), details.size()))] {
+      // When the wasm vm fails, failStream() is called if the plugin is fail-closed, we need
+      // this flag to avoid calling sendLocalReply() twice.
+      if (local_reply_sent_) {
+        return;
+      }
       decoder_callbacks_->sendLocalReply(static_cast<Envoy::Http::Code>(response_code), body_text,
                                          modify_headers, grpc_status, details);
+      local_reply_sent_ = true;
     });
   }
+  local_reply_hold_ = true;
   return WasmResult::Ok;
 }
 

source/extensions/common/wasm/context.h

@@ -437,6 +437,7 @@ class Context : public proxy_wasm::ContextBase,
   bool buffering_response_body_ = false;
   bool end_of_stream_ = false;
   bool local_reply_sent_ = false;
+  bool local_reply_hold_ = false;
   ProtobufWkt::Struct temporary_metadata_;
 
   // MB: must be a node-type map as we take persistent references to the entries.

test/extensions/common/wasm/test_data/test_context_cpp.cc

@@ -79,4 +79,42 @@ void TestRootContext::onQueueReady(uint32_t) {
   }
 }
 
+class DupReplyContext : public Context {
+public:
+  explicit DupReplyContext(uint32_t id, RootContext* root) : Context(id, root) {}
+  FilterDataStatus onRequestBody(size_t body_buffer_length, bool end_of_stream) override;
+
+private:
+  EnvoyRootContext* root() { return static_cast<EnvoyRootContext*>(Context::root()); }
+};
+
+FilterDataStatus DupReplyContext::onRequestBody(size_t, bool) {
+  sendLocalResponse(200, "ok", "body", {});
+  sendLocalResponse(200, "not send", "body", {});
+  return FilterDataStatus::Continue;
+}
+
+class PanicReplyContext : public Context {
+public:
+  explicit PanicReplyContext(uint32_t id, RootContext* root) : Context(id, root) {}
+  FilterDataStatus onRequestBody(size_t body_buffer_length, bool end_of_stream) override;
+
+private:
+  EnvoyRootContext* root() { return static_cast<EnvoyRootContext*>(Context::root()); }
+};
+
+FilterDataStatus PanicReplyContext::onRequestBody(size_t, bool) {
+  sendLocalResponse(200, "not send", "body", {});
+  int* badptr = nullptr;
+  *badptr = 0;
+  return FilterDataStatus::Continue;
+}
+
+static RegisterContextFactory register_DupReplyContext(CONTEXT_FACTORY(DupReplyContext),
+                                                       ROOT_FACTORY(EnvoyRootContext),
+                                                       "send local reply twice");
+static RegisterContextFactory register_PanicReplyContext(CONTEXT_FACTORY(PanicReplyContext),
+                                                         ROOT_FACTORY(EnvoyRootContext),
+                                                         "panic after sending local reply");
+
 END_WASM_PLUGIN

test/extensions/common/wasm/wasm_test.cc

@@ -1298,7 +1298,7 @@ TEST_P(WasmCommonTest, ThreadLocalCopyRetainsEnforcement) {
   thread_local_wasm->start(plugin);
 }
 
-class WasmCommonContextTest : public Common::Wasm::WasmTestBase<
+class WasmCommonContextTest : public Common::Wasm::WasmHttpFilterTestBase<
                                   testing::TestWithParam<std::tuple<std::string, std::string>>> {
 public:
   WasmCommonContextTest() = default;
@@ -1312,16 +1312,10 @@ class WasmCommonContextTest : public Common::Wasm::WasmTestBase<
               });
   }
 
-  void setupContext() {
-    context_ =
-        std::make_unique<TestContext>(wasm_->wasm().get(), root_context_->id(), plugin_handle_);
-    context_->onCreate();
-  }
+  void setupContext() { setupFilterBase<TestContext>(); }
 
   TestContext& rootContext() { return *static_cast<TestContext*>(root_context_); }
-  TestContext& context() { return *context_; }
-
-  std::unique_ptr<TestContext> context_;
+  TestContext& context() { return *static_cast<TestContext*>(context_.get()); }
 };
 
 INSTANTIATE_TEST_SUITE_P(Runtimes, WasmCommonContextTest,
@@ -1399,6 +1393,56 @@ TEST_P(WasmCommonContextTest, EmptyContext) {
   root_context_->validateConfiguration("", plugin_);
 }
 
+// test that we don't send the local reply twice, even though it's specified in the wasm code
+TEST_P(WasmCommonContextTest, DuplicateLocalReply) {
+  std::string code;
+  if (std::get<0>(GetParam()) != "null") {
+    code = TestEnvironment::readFileToStringForTest(TestEnvironment::substitute(absl::StrCat(
+        "{{ test_rundir }}/test/extensions/common/wasm/test_data/test_context_cpp.wasm")));
+  } else {
+    // The name of the Null VM plugin.
+    code = "CommonWasmTestContextCpp";
+  }
+  EXPECT_FALSE(code.empty());
+
+  setup(code, "context", "send local reply twice");
+  setupContext();
+  EXPECT_CALL(decoder_callbacks_, encodeHeaders_(_, _))
+      .WillOnce([this](Http::ResponseHeaderMap&, bool) { context().onResponseHeaders(0, false); });
+  EXPECT_CALL(decoder_callbacks_,
+              sendLocalReply(Envoy::Http::Code::OK, testing::Eq("body"), _, _, testing::Eq("ok")));
+
+  // Create in-VM context.
+  context().onCreate();
+  EXPECT_EQ(proxy_wasm::FilterDataStatus::StopIterationNoBuffer, context().onRequestBody(0, false));
+}
+
+// test that we don't send the local reply twice when the wasm code panics
+TEST_P(WasmCommonContextTest, LocalReplyWhenPanic) {
+  std::string code;
+  if (std::get<0>(GetParam()) != "null") {
+    code = TestEnvironment::readFileToStringForTest(TestEnvironment::substitute(absl::StrCat(
+        "{{ test_rundir }}/test/extensions/common/wasm/test_data/test_context_cpp.wasm")));
+  } else {
+    // no need test the Null VM plugin.
+    return;
+  }
+  EXPECT_FALSE(code.empty());
+
+  setup(code, "context", "panic after sending local reply");
+  setupContext();
+  // In the case of VM failure, failStream is called, so we need to make sure that we don't send the
+  // local reply twice.
+  EXPECT_CALL(decoder_callbacks_,
+              sendLocalReply(Envoy::Http::Code::ServiceUnavailable, testing::Eq(""), _,
+                             testing::Eq(Grpc::Status::WellKnownGrpcStatus::Unavailable),
+                             testing::Eq("wasm_fail_stream")));
+
+  // Create in-VM context.
+  context().onCreate();
+  EXPECT_EQ(proxy_wasm::FilterDataStatus::StopIterationNoBuffer, context().onRequestBody(0, false));
+}
+
 } // namespace Wasm
 } // namespace Common
 } // namespace Extensions