diff --git a/plugins/samples/add_geo_query/BUILD b/plugins/samples/add_geo_query/BUILD new file mode 100644 index 000000000..77233b57f --- /dev/null +++ b/plugins/samples/add_geo_query/BUILD @@ -0,0 +1,20 @@ +load("//:plugins.bzl", "proxy_wasm_plugin_rust", "proxy_wasm_tests") + +licenses(["notice"]) # Apache 2 + +proxy_wasm_plugin_rust( + name = "plugin_rust.wasm", + srcs = ["plugin.rs"], + deps = [ + "//bazel/cargo/remote:log", + "//bazel/cargo/remote:proxy-wasm", + ], +) + +proxy_wasm_tests( + name = "tests", + plugins = [ + ":plugin_rust.wasm", + ], + tests = ":tests.textpb", +) diff --git a/plugins/samples/add_geo_query/plugin.rs b/plugins/samples/add_geo_query/plugin.rs new file mode 100644 index 000000000..70e2de0c2 --- /dev/null +++ b/plugins/samples/add_geo_query/plugin.rs @@ -0,0 +1,72 @@ +// Copyright 2026 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_country_query] +use proxy_wasm::traits::{Context, HttpContext}; +use proxy_wasm::types::{Action, LogLevel}; + +const CLIENT_REGION_PATH: &[&str] = &["request", "client_region"]; +const DEFAULT_COUNTRY: &str = "unknown"; + +proxy_wasm::main! {{ + proxy_wasm::set_log_level(LogLevel::Trace); + proxy_wasm::set_http_context(|_, _| -> Box { Box::new(MyHttpContext) }); +}} + +struct MyHttpContext; + +impl Context for MyHttpContext {} + +impl HttpContext for MyHttpContext { + fn on_http_request_headers(&mut self, _num_headers: usize, _end_of_stream: bool) -> Action { + let country_value = self.get_country_value(); + + log::info!("country: {}", country_value); + + let path = self.get_http_request_header(":path").unwrap_or_default(); + let new_path = self.add_country_parameter(&path, &country_value); + + self.set_http_request_header(":path", Some(&new_path)); + + Action::Continue + } +} + +impl MyHttpContext { + fn get_country_value(&self) -> String { + if let Some(bytes) = self.get_property(CLIENT_REGION_PATH.to_vec()) { + if let Ok(country) = String::from_utf8(bytes) { + if !country.is_empty() { + return country; + } + } + } + DEFAULT_COUNTRY.to_string() + } + + fn add_country_parameter(&self, path: &str, country: &str) -> String { + let mut new_path = String::with_capacity(path.len() + country.len() + 10); + new_path.push_str(path); + + if path.contains('?') { + new_path.push_str("&country="); + } else { + new_path.push_str("?country="); + } + new_path.push_str(country); + + new_path + } +} +// [END serviceextensions_plugin_country_query] diff --git a/plugins/samples/add_geo_query/tests.textpb b/plugins/samples/add_geo_query/tests.textpb new file mode 100644 index 000000000..e39e9e29f --- /dev/null +++ b/plugins/samples/add_geo_query/tests.textpb @@ -0,0 +1,139 @@ +# Test basic country parameter addition using request.client_region +test { + name: "AddCountryParameterWithClientRegion" + properties { + path: "request.client_region" + value: "BR" + } + request_headers { + input { + header { key: ":path" value: "/api/data" } + } + result { + has_header { key: ":path" value: "/api/data?country=BR" } + log { regex: ".*country: BR.*" } + } + } +} + +# Test fallback to unknown when client_region property is missing +test { + name: "FallbackToUnknownWhenNoClientRegion" + request_headers { + input { + header { key: ":path" value: "/api/data" } + } + result { + has_header { key: ":path" value: "/api/data?country=unknown" } + log { regex: ".*country: unknown.*" } + } + } +} + +# Test fallback to unknown when client_region is present but empty +test { + name: "FallbackToUnknownWhenClientRegionEmpty" + properties { + path: "request.client_region" + value: "" + } + request_headers { + input { + header { key: ":path" value: "/api/data" } + } + result { + has_header { key: ":path" value: "/api/data?country=unknown" } + log { regex: ".*country: unknown.*" } + } + } +} + +# Test appending to existing query string +test { + name: "AppendCountryToExistingQueryString" + properties { + path: "request.client_region" + value: "DE" + } + request_headers { + input { + header { key: ":path" value: "/api/data?page=1&limit=10" } + } + result { + headers { regex: ":path: /api/data\\?.*country=DE.*" } + log { regex: ".*country: DE.*" } + } + } +} + +# Test root path without query string +test { + name: "AddCountryToRootPath" + properties { + path: "request.client_region" + value: "MX" + } + request_headers { + input { + header { key: ":path" value: "/" } + } + result { + has_header { key: ":path" value: "/?country=MX" } + log { regex: ".*country: MX.*" } + } + } +} + +# Test special characters +test { + name: "HandleClientRegionCountryCodeUK" + properties { + path: "request.client_region" + value: "UK" + } + request_headers { + input { + header { key: ":path" value: "/api/data" } + } + result { + has_header { key: ":path" value: "/api/data?country=UK" } + log { regex: ".*country: UK.*" } + } + } +} + +# Test empty path scenario +test { + name: "HandleEmptyPath" + properties { + path: "request.client_region" + value: "AU" + } + request_headers { + input { + header { key: ":path" value: "" } + } + result { + has_header { key: ":path" value: "?country=AU" } + log { regex: ".*country: AU.*" } + } + } +} + +# Test anti-spoofing: replaces existing country parameter +test { + name: "RemoveExistingCountryParameter" + properties { + path: "request.client_region" + value: "US" + } + request_headers { + input { + header { key: ":path" value: "/api/data?country=SPOOFED&page=1" } + } + result { + headers { regex: ":path: /api/data\\?.*country=US.*" } + log { regex: ".*country: US.*" } + } + } +}