feat: add knob to customise onRequestHeaders StopIteration behavior

Signed-off-by: Michael Warres <mpw@google.com>

Author: mpwarres

test/BUILD

@@ -132,6 +132,21 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "stop_iteration_test",
+    srcs = ["stop_iteration_test.cc"],
+    data = [
+        "//test/test_data:stop_iteration.wasm",
+    ],
+    linkstatic = 1,
+    deps = [
+        ":utility_lib",
+        "//:lib",
+        "@com_google_googletest//:gtest",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_test(
     name = "security_test",
     srcs = ["security_test.cc"],

test/stop_iteration_test.cc

@@ -0,0 +1,83 @@
+// Copyright 2025 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.
+
+#include "gtest/gtest.h"
+#include "include/proxy-wasm/wasm.h"
+#include "test/utility.h"
+
+namespace proxy_wasm {
+
+INSTANTIATE_TEST_SUITE_P(WasmEngines, TestVm, testing::ValuesIn(getWasmEngines()),
+                         [](const testing::TestParamInfo<std::string> &info) {
+                           return info.param;
+                         });
+
+// TestVm is parameterized for each engine and creates a VM on construction.
+TEST_P(TestVm, AllowOnRequestHeadersStopIteration) {
+  // Read the wasm source.
+  auto source = readTestWasmFile("stop_iteration.wasm");
+  ASSERT_FALSE(source.empty());
+
+  // Create a WasmBase and load the plugin.
+  auto wasm = std::make_shared<TestWasm>(std::move(vm_));
+  ASSERT_TRUE(wasm->load(source, /*allow_precompiled=*/false));
+  ASSERT_TRUE(wasm->initialize());
+
+  // Create a plugin.
+  const auto plugin = std::make_shared<PluginBase>(
+      /*name=*/"test", /*root_id=*/"", /*vm_id=*/"",
+      /*engine=*/wasm->wasm_vm()->getEngineName(), /*plugin_config=*/"",
+      /*fail_open=*/false, /*key=*/"");
+
+  // Create root context, call onStart().
+  ContextBase *root_context = wasm->start(plugin);
+  ASSERT_TRUE(root_context != nullptr);
+
+  // On the root context, call onConfigure().
+  ASSERT_TRUE(wasm->configure(root_context, plugin));
+
+  // By default, stream context onRequestHeaders translates
+  // FilterHeadersStatus::StopIteration to
+  // FilterHeadersStatus::StopAllIterationAndWatermark.
+  {
+    auto wasm_handle = std::make_shared<WasmHandleBase>(wasm);
+    auto plugin_handle = std::make_shared<PluginHandleBase>(wasm_handle, plugin);
+    auto stream_context = TestContext(wasm.get(), root_context->id(), plugin_handle);
+    stream_context.set_allow_on_request_headers_stop_iteration(true);
+    stream_context.onCreate();
+    EXPECT_EQ(stream_context.onRequestHeaders(/*headers=*/0, /*end_of_stream=*/false),
+              FilterHeadersStatus::StopIteration);
+    stream_context.onResponseHeaders(/*headers=*/0, /*end_of_stream=*/false);
+    stream_context.onDone();
+    stream_context.onDelete();
+  }
+  ASSERT_FALSE(wasm->isFailed());
+
+  // Create a stream context that propagates FilterHeadersStatus::StopIteration.
+  {
+    auto wasm_handle = std::make_shared<WasmHandleBase>(wasm);
+    auto plugin_handle = std::make_shared<PluginHandleBase>(wasm_handle, plugin);
+    auto stream_context = TestContext(wasm.get(), root_context->id(), plugin_handle);
+    stream_context.set_allow_on_request_headers_stop_iteration(true);
+    stream_context.onCreate();
+    EXPECT_EQ(stream_context.onRequestHeaders(/*headers=*/0, /*end_of_stream=*/false),
+              FilterHeadersStatus::StopIteration);
+    stream_context.onResponseHeaders(/*headers=*/0, /*end_of_stream=*/false);
+    stream_context.onDone();
+    stream_context.onDelete();
+  }
+  ASSERT_FALSE(wasm->isFailed());
+}
+
+} // namespace proxy_wasm

test/test_data/BUILD

@@ -89,3 +89,8 @@ proxy_wasm_cc_binary(
     name = "http_logging.wasm",
     srcs = ["http_logging.cc"],
 )
+
+proxy_wasm_cc_binary(
+    name = "stop_iteration.wasm",
+    srcs = ["stop_iteration.cc"],
+)

test/test_data/stop_iteration.cc

@@ -0,0 +1,27 @@
+// Copyright 2025 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.
+
+#include "proxy_wasm_intrinsics.h"
+
+class StopIterationContext : public Context {
+public:
+  explicit StopIterationContext(uint32_t id, RootContext *root) : Context(id, root) {}
+
+  FilterHeadersStatus onRequestHeaders(uint32_t headers, bool end_of_stream) override {
+    return FilterHeadersStatus::StopIteration;
+  }
+};
+
+static RegisterContextFactory register_StaticContext(CONTEXT_FACTORY(StopIterationContext),
+                                                     ROOT_FACTORY(RootContext));

test/utility.h

@@ -133,6 +133,10 @@ class TestContext : public ContextBase {
         .count();
   }
 
+  void set_allow_on_request_headers_stop_iteration(bool allow) {
+    allow_on_request_headers_stop_iteration_ = allow;
+  }
+
 private:
   std::string log_;
   static std::string global_log_;