diff --git a/plugins/samples/body_chunking/BUILD b/plugins/samples/body_chunking/BUILD new file mode 100644 index 00000000..61ef08aa --- /dev/null +++ b/plugins/samples/body_chunking/BUILD @@ -0,0 +1,25 @@ +load("//:plugins.bzl", "proxy_wasm_plugin_cpp", "proxy_wasm_plugin_rust", "proxy_wasm_tests") + +licenses(["notice"]) # Apache 2 + + +proxy_wasm_plugin_cpp( + name = "plugin_cpp.wasm", + srcs = ["plugin.cc"], + deps = [ + "@proxy_wasm_cpp_sdk//contrib:contrib_lib", + ], +) + +proxy_wasm_tests( + name = "tests", + data = [ + ":request_body.data", + ":response_body.data", + ":expected_request_body.data", + ], + plugins = [ + ":plugin_cpp.wasm", + ], + tests = ":tests.textpb", +) diff --git a/plugins/samples/body_chunking/expected_request_body.data b/plugins/samples/body_chunking/expected_request_body.data new file mode 100644 index 00000000..faecae41 --- /dev/null +++ b/plugins/samples/body_chunking/expected_request_body.data @@ -0,0 +1 @@ +12foo34foo56foo78foo90foo \ No newline at end of file diff --git a/plugins/samples/body_chunking/plugin.cc b/plugins/samples/body_chunking/plugin.cc new file mode 100644 index 00000000..a287f8ee --- /dev/null +++ b/plugins/samples/body_chunking/plugin.cc @@ -0,0 +1,39 @@ +// Copyright 2024 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. + +// [START serviceextensions_plugin_body_chunking] +#include "proxy_wasm_intrinsics.h" + +class MyHttpContext : public Context { + public: + explicit MyHttpContext(uint32_t id, RootContext* root) : Context(id, root) {} + + // Add foo onto the end of each request body chunk + FilterDataStatus onRequestBody(size_t chunk_len, + bool end_of_stream) override { + setBuffer(WasmBufferType::HttpRequestBody, chunk_len, 0, "foo"); + return FilterDataStatus::Continue; + } + + // Add bar onto the end of each response body chunk + FilterDataStatus onResponseBody(size_t chunk_len, + bool end_of_stream) override { + setBuffer(WasmBufferType::HttpResponseBody, chunk_len, 0, "bar"); + return FilterDataStatus::Continue; + } +}; + +static RegisterContextFactory register_StaticContext( + CONTEXT_FACTORY(MyHttpContext), ROOT_FACTORY(RootContext)); +// [END serviceextensions_plugin_body_chunking] diff --git a/plugins/samples/body_chunking/request_body.data b/plugins/samples/body_chunking/request_body.data new file mode 100644 index 00000000..6a537b5b --- /dev/null +++ b/plugins/samples/body_chunking/request_body.data @@ -0,0 +1 @@ +1234567890 \ No newline at end of file diff --git a/plugins/samples/body_chunking/response_body.data b/plugins/samples/body_chunking/response_body.data new file mode 100644 index 00000000..1712d202 --- /dev/null +++ b/plugins/samples/body_chunking/response_body.data @@ -0,0 +1 @@ +0987654321 \ No newline at end of file diff --git a/plugins/samples/body_chunking/tests.textpb b/plugins/samples/body_chunking/tests.textpb new file mode 100644 index 00000000..ed1beaa2 --- /dev/null +++ b/plugins/samples/body_chunking/tests.textpb @@ -0,0 +1,78 @@ +test { + name: "NumChunkWithFileInput" + num_chunks: 5 + request_body { + input { file: "request_body.data" } + result { body { exact: "12foo34foo56foo78foo90foo" } } + } + response_body { + input { file: "response_body.data" } + result { body { exact: "09bar87bar65bar43bar21bar" } } + } +} +test { + name: "NumChunkWithFileInputFileOutput" + num_chunks: 5 + request_body { + input { file: "request_body.data" } + result { body { file: "expected_request_body.data" } } + } +} +test { + name: "NumChunkWithContentInput" + num_chunks: 5 + request_body { + input { content: "1234567890" } + result { body { exact: "12foo34foo56foo78foo90foo" } } + } +} +test { + name: "ChunkSizeWithFileInput" + chunk_size: 2 + request_body { + input { file: "request_body.data" } + result { body { exact: "12foo34foo56foo78foo90foo" } } + } +} +test { + name: "ChunkSizeWithContentInput" + chunk_size: 2 + request_body { + input { content: "1234567890" } + result { body { exact: "12foo34foo56foo78foo90foo" } } + } +} + +test { + name: "NoChunking" + request_body{ + input { + content: "12" + } + result { + body { + exact: "12foo" + } + } + } + request_body{ + input { + content: "34" + } + result { + body { + exact: "34foo" + } + } + } + request_body{ + input { + content: "56" + } + result { + body { + exact: "56foo" + } + } + } +} diff --git a/plugins/test/dynamic_test.cc b/plugins/test/dynamic_test.cc index 240e2932..3d43e90e 100644 --- a/plugins/test/dynamic_test.cc +++ b/plugins/test/dynamic_test.cc @@ -33,16 +33,38 @@ #include "test/runner.pb.h" namespace service_extensions_samples { +namespace { +// Helper to read data from a path which may be relative to the test config. +absl::StatusOr ReadContent(const std::string& path, + const pb::Env& env) { + boost::filesystem::path in(path); + boost::filesystem::path cfg(env.test_path()); + // If relative path, resolve relative to test config file. + if (!in.is_absolute()) { + in = cfg.parent_path() / in; + } + return ReadDataFile(in.string()); +} // Helper class to support string/regex positive/negative matching. class StringMatcher { public: // Factory method. Creates a matcher or returns invalid argument status. - static absl::StatusOr Create(const pb::StringMatcher& expect) { + static absl::StatusOr Create(const pb::StringMatcher& expect, + const pb::Env& env) { StringMatcher sm; sm.invert_ = expect.invert(); if (expect.has_exact()) { sm.exact_ = expect.exact(); + } else if (expect.has_file()) { + absl::StatusOr file_content = + ReadContent(expect.file(), env); + if (file_content.ok()) { + sm.exact_ = *file_content; + } else { + return absl::InvalidArgumentError( + absl::StrCat("Bad file: ", file_content.status())); + } } else if (expect.has_regex()) { sm.re_ = std::make_unique(expect.regex(), RE2::Quiet); if (!sm.re_->ok()) { @@ -65,6 +87,8 @@ class StringMatcher { return invert_; }; + std::string getExact() { return exact_; } + private: StringMatcher() = default; @@ -73,6 +97,29 @@ class StringMatcher { std::unique_ptr re_ = nullptr; }; +// Helper to add some logging at construction and teardown. +class LogTestBounds { + public: + LogTestBounds(ContextOptions& options) : options_(options) { + if (options_.log_file) { + const auto* test = testing::UnitTest::GetInstance()->current_test_info(); + options_.log_file << "--- Starting test: " << test->name() << " ---" + << std::endl; + } + } + ~LogTestBounds() { + if (options_.log_file) { + const auto* test = testing::UnitTest::GetInstance()->current_test_info(); + options_.log_file << "--- Finished test: " << test->name() << " ---" + << std::endl; + } + } + + private: + ContextOptions& options_; +}; +} // namespace + absl::StatusOr> DynamicTest::LoadWasm(bool benchmark) { // Set log level. Default to INFO. Disable in benchmarks. @@ -119,30 +166,6 @@ DynamicTest::LoadWasm(bool benchmark) { absl::StrJoin(context.phase_logs(), "\n")); \ } -namespace { -// Helper to add some logging at construction and teardown. -class LogTestBounds { - public: - LogTestBounds(ContextOptions& options) : options_(options) { - if (options_.log_file) { - const auto* test = testing::UnitTest::GetInstance()->current_test_info(); - options_.log_file << "--- Starting test: " << test->name() << " ---" - << std::endl; - } - } - ~LogTestBounds() { - if (options_.log_file) { - const auto* test = testing::UnitTest::GetInstance()->current_test_info(); - options_.log_file << "--- Finished test: " << test->name() << " ---" - << std::endl; - } - } - - private: - ContextOptions& options_; -}; -} // namespace - void DynamicTest::TestBody() { // Initialize VM. auto load_wasm = LoadWasm(/*benchmark=*/false); @@ -177,14 +200,47 @@ void DynamicTest::TestBody() { ASSERT_VM_HEALTH("request_headers", handle, stream); CheckPhaseResults("request_headers", invoke.result(), stream, res); } - if (cfg_.request_body_size() > 0) { - for (auto& invoke : *cfg_.mutable_request_body()) { - auto res = stream.SendRequestBody(std::move( - *absl::WrapUnique(invoke.mutable_input()->release_content()))); - ASSERT_VM_HEALTH("request_body", handle, stream); - CheckPhaseResults("request_body", invoke.result(), stream, res); + auto run_body_test = [&handle, &stream, this]( + std::string phase, + google::protobuf::RepeatedPtrField + invocations, + auto invoke_wasm) { + if (invocations.size() == 0) return; + auto body_chunking_plan = cfg_.body_chunking_plan_case(); + if (invocations.size() != 1 && + body_chunking_plan != + pb::Test::BodyChunkingPlanCase::BODY_CHUNKING_PLAN_NOT_SET) { + FAIL() + << "Cannot specify body_chunking_plan with multiple body invocations"; + return; } - } + for (const pb::Invocation& invocation : invocations) { + TestHttpContext::Result body_result = TestHttpContext::Result{}; + absl::StatusOr complete_input_body = ParseBodyInput(invocation.input()); + if (!complete_input_body.ok()) { + FAIL() << complete_input_body.status(); + } + + std::vector chunks; + if (body_chunking_plan != + pb::Test::BodyChunkingPlanCase::BODY_CHUNKING_PLAN_NOT_SET) { + chunks = ChunkBody(*complete_input_body, cfg_); + } else { + chunks = {*complete_input_body}; + } + for (std::string& body_chunk : chunks) { + TestHttpContext::Result res = invoke_wasm(std::move(body_chunk)); + ASSERT_VM_HEALTH(phase, handle, stream); + body_result.body = body_result.body.append(res.body); + } + CheckPhaseResults(phase, invocation.result(), stream, body_result); + } + }; + + ASSERT_NO_FATAL_FAILURE(run_body_test( + "request_body", cfg_.request_body(), + [&stream](std::string chunk) { return stream.SendRequestBody(chunk); })); + if (cfg_.has_response_headers()) { const auto& invoke = cfg_.response_headers(); auto headers = ParseHeaders(invoke.input(), /*is_request=*/false); @@ -193,14 +249,10 @@ void DynamicTest::TestBody() { ASSERT_VM_HEALTH("response_headers", handle, stream); CheckPhaseResults("response_headers", invoke.result(), stream, res); } - if (cfg_.response_body_size() > 0) { - for (auto& invoke : *cfg_.mutable_response_body()) { - auto res = stream.SendResponseBody(std::move( - *absl::WrapUnique(invoke.mutable_input()->release_content()))); - ASSERT_VM_HEALTH("response_body", handle, stream); - CheckPhaseResults("response_body", invoke.result(), stream, res); - } - } + + ASSERT_NO_FATAL_FAILURE(run_body_test( + "response_body", cfg_.response_body(), + [&stream](std::string chunk) { return stream.SendResponseBody(chunk); })); // Tear down HTTP context. stream.TearDown(); @@ -286,14 +338,13 @@ void DynamicTest::BenchHttpHandlers(benchmark::State& state) { BM_RETURN_IF_ERROR(headers.status()); response_headers = *headers; } - std::vector request_body_chunks; - for (const auto& request_body : cfg_.request_body()) { - request_body_chunks.emplace_back(request_body.input().content()); - } - std::vector response_body_chunks; - for (const auto& response_body : cfg_.response_body()) { - response_body_chunks.emplace_back(response_body.input().content()); - } + absl::StatusOr> request_body_chunks = + PrepBodyCallbackBenchmark(cfg_, cfg_.request_body()); + BM_RETURN_IF_ERROR(request_body_chunks.status()); + absl::StatusOr> response_body_chunks = + PrepBodyCallbackBenchmark(cfg_, cfg_.response_body()); + BM_RETURN_IF_ERROR(response_body_chunks.status()); + std::optional stream; for (auto _ : state) { // Pausing timing here is not recommended. One way we could avoid it: @@ -301,8 +352,9 @@ void DynamicTest::BenchHttpHandlers(benchmark::State& state) { // - don't hand ownership of body chunks to stream context state.PauseTiming(); stream.emplace(handle); // create/destroy TestHttpContext - std::vector request_body_chunks_copies = request_body_chunks; - std::vector response_body_chunks_copies = response_body_chunks; + std::vector request_body_chunks_copies = *request_body_chunks; + std::vector response_body_chunks_copies = + *response_body_chunks; state.ResumeTiming(); if (request_headers) { @@ -423,7 +475,7 @@ void DynamicTest::CheckPhaseResults(const std::string& phase, void DynamicTest::FindString(const std::string& phase, const std::string& type, const pb::StringMatcher& expect, const std::vector& contents) { - auto matcher = StringMatcher::Create(expect); + auto matcher = StringMatcher::Create(expect, env_); if (!matcher.ok()) { ADD_FAILURE() << absl::Substitute("[$0] $1", phase, matcher.status().ToString()); @@ -434,7 +486,7 @@ void DynamicTest::FindString(const std::string& phase, const std::string& type, "[$0] expected $1 of $2 $3: '$4', actual: \n$5", phase, expect.invert() ? "absence" : "presence", expect.has_regex() ? "regex" : "exact", type, - expect.has_regex() ? expect.regex() : expect.exact(), + expect.has_regex() ? expect.regex() : matcher->getExact(), absl::StrJoin(contents, "\n")); } } @@ -505,7 +557,7 @@ absl::StatusOr DynamicTest::ParseHeaders( if (!input.file().empty()) { // Handle file input. - auto content = ReadContent(input.file()); + auto content = ReadContent(input.file(), env_); if (!content.ok()) return content.status(); auto parse = ParseHTTP1Headers(*content, is_request, hdrs); if (!parse.ok()) return parse; @@ -521,15 +573,56 @@ absl::StatusOr DynamicTest::ParseHeaders( } return hdrs; } +absl::StatusOr DynamicTest::ParseBodyInput( + const pb::Input& input) { + if (!input.file().empty()) { + // Handle file input. + return ReadContent(input.file(), env_); + } else { + return input.content(); + } +} -absl::StatusOr DynamicTest::ReadContent(const std::string& path) { - boost::filesystem::path in(path); - boost::filesystem::path cfg(env_.test_path()); - // If relative path, resolve relative to test config file. - if (!in.is_absolute()) { - in = cfg.parent_path() / in; +std::vector DynamicTest::ChunkBody( + const std::string& complete_body, const pb::Test& test) { + if (test.has_num_chunks()) { + int chunk_length = complete_body.size() / test.num_chunks(); + std::vector body_chunks = absl::StrSplit( + complete_body, + absl::MaxSplits(absl::ByLength(chunk_length), test.num_chunks() - 1)); + return body_chunks; + } else { + std::vector body_chunks = + absl::StrSplit(complete_body, absl::ByLength(test.chunk_size())); + return body_chunks; } - return ReadDataFile(in.string()); +} + +absl::StatusOr> DynamicTest::PrepBodyCallbackBenchmark( + const pb::Test& test, + google::protobuf::RepeatedPtrField invocations) { + if (invocations.size() == 0) return std::vector {}; + auto body_chunking_plan = test.body_chunking_plan_case(); + if (invocations.size() != 1 && + body_chunking_plan != + pb::Test::BodyChunkingPlanCase::BODY_CHUNKING_PLAN_NOT_SET) { + return absl::InvalidArgumentError( + "Cannot specify body_chunking_plan with multiple body invocations"); + } + std::vector chunks; + for (const pb::Invocation& invocation : invocations) { + absl::StatusOr complete_input_body = ParseBodyInput(invocation.input()); + if (!complete_input_body.ok()) { + return complete_input_body.status(); + } + if (body_chunking_plan != + pb::Test::BodyChunkingPlanCase::BODY_CHUNKING_PLAN_NOT_SET) { + chunks = ChunkBody(*complete_input_body, cfg_); + } else { + chunks.emplace_back(*complete_input_body); + } + } + return chunks; } } // namespace service_extensions_samples diff --git a/plugins/test/dynamic_test.h b/plugins/test/dynamic_test.h index e155d30e..8863da90 100644 --- a/plugins/test/dynamic_test.h +++ b/plugins/test/dynamic_test.h @@ -80,8 +80,18 @@ class DynamicTest : public DynamicFixture { // Helper to generate Headers struct from proto, string, or file. absl::StatusOr ParseHeaders(const pb::Input& input, bool is_request); - // Helper to read data from a path which may be relative to the test config. - absl::StatusOr ReadContent(const std::string& path); + + // Helper to generate body string from test config + absl::StatusOr ParseBodyInput(const pb::Input& input); + + // Helper to break body down into chunks + std::vector ChunkBody(const std::string& complete_body, + const pb::Test& test); + + // Helper to prep body callbacks for benchmarking + absl::StatusOr> PrepBodyCallbackBenchmark( + const pb::Test& test, + google::protobuf::RepeatedPtrField invocations); std::string engine_; pb::Env env_; diff --git a/plugins/test/runner.proto b/plugins/test/runner.proto index 5e8f0882..8889e885 100644 --- a/plugins/test/runner.proto +++ b/plugins/test/runner.proto @@ -27,6 +27,8 @@ message StringMatcher { bytes exact = 2; // Regex expectation (RE2 full-match). string regex = 3; + // File exact match expectation + string file = 4; } } @@ -49,7 +51,9 @@ message Input { bytes content = 2; // Serialized HTTP content, but provided via a separate file. - // File path can be absolute or relative to the test config file. + // For body input, don't include any content other than the body you would + // like to test. File path can be absolute or relative to the test + // config file. string file = 3; } @@ -92,14 +96,25 @@ message Test { string name = 1; // Whether to run benchmarks, in addition to verifying any expectations. bool benchmark = 2; - + // If a chunking_plan is set, body invocations will be treated as the complete + // body and will broken into chunks according to specified plan. Expectations + // will be evaluated against the reconstructed complete body after processing + // and not as individual chunks. + oneof body_chunking_plan { + int32 num_chunks = 10; + int64 chunk_size = 11; + } // Wasm invocations of various HTTP phases. If multiple phases are provided, // they are executed in sequence on the same HttpContext. Each phase specifies // its own expectations. If any phase results in an immediate response, // further phases are not executed, and later expectations will fail. Invocation request_headers = 3; + // Only add one request_body invocation if using a chunking_plan. If multiple + // request_body invocations are present, tests will fail. repeated Invocation request_body = 4; Invocation response_headers = 5; + // Only add one request_body invocation if using a chunking_plan. If multiple + // response_body invocations are present, tests will fail. repeated Invocation response_body = 6; // Expectations for plugin and stream lifecycle. Useful for testing side