Don't pause processing when send_local_response fails

src/exports.cc

@@ -153,8 +153,25 @@ Word send_local_response(Word response_code, Word response_code_details_ptr,
     return WasmResult::InvalidMemoryAccess;
   }
   auto additional_headers = PairsUtil::toPairs(additional_response_header_pairs.value());
-  context->sendLocalResponse(response_code, body.value(), std::move(additional_headers),
-                             grpc_status, details.value());
+  auto status = context->sendLocalResponse(
+      response_code, body.value(), std::move(additional_headers), grpc_status, details.value());
+  // Only stop processing if we actually triggered local response.
+  //
+  // For context, Envoy sends local replies through the filter chain,
+  // so wasm filter can be called to handle a local reply that the
+  // filter itself triggered.
+  //
+  // Normally that is not an issue, unless wasm filter calls
+  // proxy_send_local_response again (which they probably shouldn't).
+  // In this case, no new local response will be generated and
+  // sendLocalResponse will fail.
+  //
+  // If at this point we stop processing, we end up in a situation when
+  // no response was sent, even though we tried twice, and the connection
+  // is stuck, because processing is stopped.
+  if (status != WasmResult::Ok) {
+    return status;
+  }
   context->wasm()->stopNextIteration(true);
   return WasmResult::Ok;
 }

test/BUILD

@@ -89,6 +89,7 @@ cc_test(
     data = [
         "//test/test_data:clock.wasm",
         "//test/test_data:env.wasm",
+        "//test/test_data:local_response.wasm",
         "//test/test_data:random.wasm",
     ],
     linkstatic = 1,

test/exports_test.cc

@@ -157,5 +157,39 @@ TEST_P(TestVm, RandomTooLarge) {
   EXPECT_TRUE(context->isLogged("random_get(66560) failed."));
 }
 
+TEST_P(TestVm, SendLocalResponse) {
+  auto source = readTestWasmFile("local_response.wasm");
+  ASSERT_FALSE(source.empty());
+  auto wasm = TestWasm(std::move(vm_));
+  ASSERT_TRUE(wasm.load(source, false));
+  ASSERT_TRUE(wasm.initialize());
+
+  auto *context = dynamic_cast<TestContext *>(wasm.vm_context());
+
+  // We first try the negative case - proxy_send_local_response fails
+  WasmCallVoid<0> run_fail;
+  wasm.wasm_vm()->getFunction("run_fail", &run_fail);
+  ASSERT_TRUE(run_fail != nullptr);
+  run_fail(context);
+
+  // We expect application to log whatever status
+  // proxy_send_local_response returns.
+  EXPECT_TRUE(context->isLogged(stringify("proxy_send_local_response returned ",
+                                          static_cast<uint32_t>(WasmResult::Unimplemented))));
+  // When we fail to send local response we don't pause processing.
+  EXPECT_FALSE(context->wasm()->isNextIterationStopped());
+
+  // Then we try the positive case - proxy_send_local_response succeeds
+  WasmCallVoid<0> run_success;
+  wasm.wasm_vm()->getFunction("run_success", &run_success);
+  ASSERT_TRUE(run_success != nullptr);
+  run_success(context);
+
+  EXPECT_TRUE(context->isLogged(
+      stringify("proxy_send_local_response returned ", static_cast<uint32_t>(WasmResult::Ok))));
+  // When we succeed to send local response we stop processing.
+  EXPECT_TRUE(context->wasm()->isNextIterationStopped());
+}
+
 } // namespace
 } // namespace proxy_wasm

test/test_data/BUILD

@@ -80,6 +80,12 @@ wasm_rust_binary(
     wasi = True,
 )
 
+wasm_rust_binary(
+    name = "local_response.wasm",
+    srcs = ["local_response.rs"],
+    wasi = True,
+)
+
 proxy_wasm_cc_binary(
     name = "canary_check.wasm",
     srcs = ["canary_check.cc"],

test/test_data/local_response.rs

@@ -0,0 +1,62 @@
+// Copyright 2021 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#[no_mangle]
+pub extern "C" fn proxy_abi_version_0_2_0() {}
+
+#[no_mangle]
+pub extern "C" fn proxy_on_memory_allocate(_: usize) -> *mut u8 {
+    std::ptr::null_mut()
+}
+
+fn send_http_response(status_code: u32) -> u32 {
+    let headers = 0u32.to_le_bytes().to_vec();
+    unsafe {
+        proxy_send_local_response(
+            status_code,
+            std::ptr::null(),
+            0,
+            std::ptr::null(),
+            0,
+            headers.as_ptr(),
+            headers.len(),
+            -1)
+    }
+}
+
+#[no_mangle]
+pub extern "C" fn run_fail() {
+    println!(
+        "proxy_send_local_response returned {}",
+        send_http_response(404));
+}
+
+#[no_mangle]
+pub extern "C" fn run_success() {
+    println!(
+        "proxy_send_local_response returned {}",
+        send_http_response(200));
+}
+
+extern "C" {
+    fn proxy_send_local_response(
+        status_code: u32,
+        status_code_details_data: *const u8,
+        status_code_details_size: usize,
+        body_data: *const u8,
+        body_size: usize,
+        headers_data: *const u8,
+        headers_size: usize,
+        grpc_status: i32,
+    ) -> u32;
+}

test/utility.h

@@ -45,6 +45,31 @@ namespace proxy_wasm {
 std::vector<std::string> getWasmEngines();
 std::string readTestWasmFile(const std::string &filename);
 
+namespace internal {
+
+template <typename... Args> struct Stringify {
+  static void convert(std::ostream &out) {}
+};
+
+template <typename... Args> void stringify_impl(std::ostream &out, Args... args) {
+  Stringify<Args...>::convert(out, std::forward<Args>(args)...);
+}
+
+template <typename A, typename... Args> struct Stringify<A, Args...> {
+  static void convert(std::ostream &out, A arg, Args... args) {
+    out << arg;
+    stringify_impl(out, std::forward<Args>(args)...);
+  }
+};
+
+} // namespace internal
+
+template <typename... Args> std::string stringify(Args... args) {
+  std::ostringstream out(std::ostringstream::ate);
+  internal::stringify_impl(out, std::forward<Args>(args)...);
+  return out.str();
+}
+
 class TestIntegration : public WasmVmIntegration {
 public:
   ~TestIntegration() override = default;
@@ -133,6 +158,15 @@ class TestContext : public ContextBase {
         .count();
   }
 
+  WasmResult sendLocalResponse(uint32_t response_code, std::string_view /*unused*/,
+                               Pairs /*unused*/, GrpcStatusCode /*unused*/,
+                               std::string_view /*unused*/) override {
+    if (response_code >= 200 && response_code < 300) {
+      return WasmResult::Ok;
+    }
+    return WasmResult::Unimplemented;
+  }
+
 private:
   std::string log_;
   static std::string global_log_;