Add support for setting HTTP bodies. (#2)

Signed-off-by: Gregory Brail <gregbrail@google.com>

Author: gbrail

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"]

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",
+    ],
+)

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<dyn HttpContext> { 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
+    }
+}

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,

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
     }