From 998ce82dca245c609bfc828c77c01573fe8bfd4b Mon Sep 17 00:00:00 2001 From: Greg Brail Date: Tue, 4 Aug 2020 20:47:23 -0700 Subject: [PATCH] Add support for setting HTTP bodies. (#2) Signed-off-by: Gregory Brail --- Cargo.toml | 5 +++ examples/BUILD | 12 ++++++++ examples/http_body.rs | 72 +++++++++++++++++++++++++++++++++++++++++++ src/hostcalls.rs | 24 +++++++++++++++ src/traits.rs | 8 +++++ 5 files changed, 121 insertions(+) create mode 100644 examples/http_body.rs diff --git a/Cargo.toml b/Cargo.toml index 1d6418a2..91d77997 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,3 +43,8 @@ crate-type = ["cdylib"] name = "http_headers" path = "examples/http_headers.rs" crate-type = ["cdylib"] + +[[example]] +name = "http_body" +path = "examples/http_body.rs" +crate-type = ["cdylib"] diff --git a/examples/BUILD b/examples/BUILD index dee0ce59..4bd527a9 100644 --- a/examples/BUILD +++ b/examples/BUILD @@ -36,3 +36,15 @@ rust_binary( "//cargo:log", ], ) + +rust_binary( + name = "http_body", + srcs = ["http_body.rs"], + crate_type = "cdylib", + edition = "2018", + out_binary = True, + deps = [ + "//:proxy_wasm", + "//cargo:log", + ], +) diff --git a/examples/http_body.rs b/examples/http_body.rs new file mode 100644 index 00000000..a4242f86 --- /dev/null +++ b/examples/http_body.rs @@ -0,0 +1,72 @@ +// Copyright 2020 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. + +use proxy_wasm::traits::*; +use proxy_wasm::types::*; + +#[no_mangle] +pub fn _start() { + proxy_wasm::set_log_level(LogLevel::Trace); + proxy_wasm::set_http_context(|_, _| -> Box { Box::new(HttpBody::new()) }); +} + +#[derive(Default)] +struct HttpBody { + total_body_size: usize, +} + +impl HttpBody { + fn new() -> HttpBody { + Default::default() + } +} + +impl Context for HttpBody {} + +impl HttpContext for HttpBody { + fn on_http_response_headers(&mut self, _: usize) -> Action { + // If there is a Content-Length header and we change the length of + // the body later, then clients will break. So remove it. + // We must do this here, because once we exit this function we + // can no longer modify the response headers. + self.set_http_response_header("content-length", None); + // Don't continue to the next callout in the chain because we might + // modify the body. + Action::Pause + } + + fn on_http_response_body(&mut self, body_size: usize, end_of_stream: bool) -> Action { + self.total_body_size += body_size; + if !end_of_stream { + // Wait -- we'll be called again when the complete body is buffered + // at the host side. + return Action::Pause; + } + + // Replace the message body if it contains the text "secret". + // Since we returned "Pause" previuously, this will return the whole body. + // However, we have to calculate the size ourselves. + if let Some(body_bytes) = self.get_http_response_body(0, self.total_body_size) { + let body_str = String::from_utf8(body_bytes).unwrap(); + if body_str.find("secret").is_some() { + let new_body = format!( + "Original message body ({} bytes) redacted.", + self.total_body_size + ); + self.set_http_response_body(0, self.total_body_size, &new_body.into_bytes()); + } + } + Action::Continue + } +} diff --git a/src/hostcalls.rs b/src/hostcalls.rs index 0b770a0f..1d7509d7 100644 --- a/src/hostcalls.rs +++ b/src/hostcalls.rs @@ -127,6 +127,30 @@ pub fn get_buffer( } } +extern "C" { + fn proxy_set_buffer_bytes( + buffer_type: BufferType, + start: usize, + size: usize, + buffer_data: *const u8, + buffer_size: usize, + ) -> Status; +} + +pub fn set_buffer( + buffer_type: BufferType, + start: usize, + size: usize, + value: &[u8], +) -> Result<(), Status> { + unsafe { + match proxy_set_buffer_bytes(buffer_type, start, size, value.as_ptr(), value.len()) { + Status::Ok => Ok(()), + status => panic!("unexpected status: {}", status as u32), + } + } +} + extern "C" { fn proxy_get_header_map_pairs( map_type: MapType, diff --git a/src/traits.rs b/src/traits.rs index ffad21a4..268ef39c 100644 --- a/src/traits.rs +++ b/src/traits.rs @@ -184,6 +184,10 @@ pub trait HttpContext: Context { hostcalls::get_buffer(BufferType::HttpRequestBody, start, max_size).unwrap() } + fn set_http_request_body(&self, start: usize, size: usize, value: &[u8]) { + hostcalls::set_buffer(BufferType::HttpRequestBody, start, size, value).unwrap() + } + fn on_http_request_trailers(&mut self, _num_trailers: usize) -> Action { Action::Continue } @@ -244,6 +248,10 @@ pub trait HttpContext: Context { hostcalls::get_buffer(BufferType::HttpResponseBody, start, max_size).unwrap() } + fn set_http_response_body(&self, start: usize, size: usize, value: &[u8]) { + hostcalls::set_buffer(BufferType::HttpResponseBody, start, size, value).unwrap() + } + fn on_http_response_trailers(&mut self, _num_trailers: usize) -> Action { Action::Continue }