wasm: add proxy_set_buffer_bytes. (envoyproxy#442) (envoyproxy#184)

This function can be used to: 1. Prepend data to the buffer using proxy_set_buffer(0, 0, data); 2. Append data to the buffer using proxy_set_buffer(MAX, 0, data); 3. Replace data in the buffer using proxy_set_buffer(0, MAX, data); 4. Remove data from the buffer using proxy_set_buffer(0, MAX, ""); While there, allow using MAX in proxy_get_buffer_bytes(). Signed-off-by: Piotr Sikora <[email protected]>
gargnupur · Mar 4, 2020 · c90c744 · c90c744
1 parent 80d0be0
commit c90c744
Show file tree

Hide file tree

Showing 12 changed files with 229 additions and 26 deletions.
diff --git a/api/wasm/cpp/proxy_wasm_api.h b/api/wasm/cpp/proxy_wasm_api.h
@@ -728,6 +728,10 @@ inline WasmResult getBufferStatus(BufferType type, size_t* size, uint32_t* flags
   return proxy_get_buffer_status(type, size, flags);
 }
 
+inline WasmResult setBuffer(BufferType type, size_t start, size_t length, StringView data) {
+  return proxy_set_buffer_bytes(type, start, length, data.data(), data.size());
+}
+
 // HTTP
 
 inline void MakeHeaderStringPairsBuffer(const HeaderStringPairs& headers, void** buffer_ptr,

diff --git a/api/wasm/cpp/proxy_wasm_externs.h b/api/wasm/cpp/proxy_wasm_externs.h
@@ -94,6 +94,8 @@ extern "C" WasmResult proxy_get_buffer_bytes(BufferType type, uint32_t start, ui
                                              const char** ptr, size_t* size);
 extern "C" WasmResult proxy_get_buffer_status(BufferType type, size_t* length_ptr,
                                               uint32_t* flags_ptr);
+extern "C" WasmResult proxy_set_buffer_bytes(BufferType type, uint32_t start, uint32_t length,
+                                             const char* ptr, size_t size);
 
 // HTTP
 extern "C" WasmResult proxy_http_call(const char* uri_ptr, size_t uri_size, void* header_pairs_ptr,

diff --git a/api/wasm/cpp/proxy_wasm_intrinsics.js b/api/wasm/cpp/proxy_wasm_intrinsics.js
@@ -23,6 +23,7 @@ mergeInto(LibraryManager.library, {
     proxy_remove_header_map_value: function () {},
     proxy_get_buffer_bytes: function () {},
     proxy_get_buffer_status: function () {},
+    proxy_set_buffer_bytes: function () {},
     proxy_http_call: function () {},
     proxy_define_metric: function () {},
     proxy_increment_metric: function () {},

diff --git a/source/extensions/common/wasm/context.cc b/source/extensions/common/wasm/context.cc
@@ -817,6 +817,33 @@ const Buffer::Instance* Context::getBuffer(BufferType type) {
   return nullptr;
 }
 
+WasmResult Context::setBuffer(BufferType type, std::function<void(Buffer::Instance&)> callback) {
+  switch (type) {
+  case BufferType::HttpRequestBody:
+    if (buffering_request_body_) {
+      decoder_callbacks_->modifyDecodingBuffer(callback);
+    } else {
+      callback(*request_body_buffer_);
+    }
+    return WasmResult::Ok;
+  case BufferType::HttpResponseBody:
+    if (buffering_response_body_) {
+      encoder_callbacks_->modifyEncodingBuffer(callback);
+    } else {
+      callback(*response_body_buffer_);
+    }
+    return WasmResult::Ok;
+  case BufferType::NetworkDownstreamData:
+    callback(*network_downstream_data_buffer_);
+    return WasmResult::Ok;
+  case BufferType::NetworkUpstreamData:
+    callback(*network_upstream_data_buffer_);
+    return WasmResult::Ok;
+  default:
+    return WasmResult::BadArgument;
+  }
+}
+
 // Async call via HTTP
 WasmResult Context::httpCall(absl::string_view cluster, const Pairs& request_headers,
                              absl::string_view request_body, const Pairs& request_trailers,
@@ -1208,9 +1235,9 @@ Http::FilterDataStatus Context::onRequestBody(bool end_of_stream) {
   }
   DeferAfterCallActions actions(this);
   const auto buffer = getBuffer(BufferType::HttpRequestBody);
-  const auto body_len = (buffer == nullptr) ? 0 : buffer->length();
+  const auto buffer_length = (buffer == nullptr) ? 0 : buffer->length();
   switch (wasm_
-              ->on_request_body_(this, id_, static_cast<uint32_t>(body_len),
+              ->on_request_body_(this, id_, static_cast<uint32_t>(buffer_length),
                                  static_cast<uint32_t>(end_of_stream))
               .u64_) {
   case 0:
@@ -1275,9 +1302,9 @@ Http::FilterDataStatus Context::onResponseBody(bool end_of_stream) {
   }
   DeferAfterCallActions actions(this);
   const auto buffer = getBuffer(BufferType::HttpResponseBody);
-  const auto body_len = (buffer == nullptr) ? 0 : buffer->length();
+  const auto buffer_length = (buffer == nullptr) ? 0 : buffer->length();
   switch (wasm_
-              ->on_response_body_(this, id_, static_cast<uint32_t>(body_len),
+              ->on_response_body_(this, id_, static_cast<uint32_t>(buffer_length),
                                   static_cast<uint32_t>(end_of_stream))
               .u64_) {
   case 0:

diff --git a/source/extensions/common/wasm/context.h b/source/extensions/common/wasm/context.h
@@ -266,6 +266,7 @@ class Context : public Logger::Loggable<Logger::Id::wasm>,
 
   // Buffer
   virtual const Buffer::Instance* getBuffer(BufferType type);
+  virtual WasmResult setBuffer(BufferType type, std::function<void(Buffer::Instance&)> callback);
   bool end_of_stream() { return end_of_stream_; }
 
   // HTTP

diff --git a/source/extensions/common/wasm/exports.cc b/source/extensions/common/wasm/exports.cc
@@ -457,22 +457,23 @@ Word get_buffer_bytes(void* raw_context, Word type, Word start, Word length, Wor
     return wasmResultToWord(WasmResult::NotFound);
   }
   // NB: check for overflow.
-  if (buffer->length() < start.u64_ + length.u64_ || start.u64_ > start.u64_ + length.u64_) {
+  if (start.u64_ > start.u64_ + length.u64_) {
     return wasmResultToWord(WasmResult::BadArgument);
   }
   uint64_t pointer = 0;
   void* p = nullptr;
-  if (length.u64_ > 0) {
-    p = context->wasm()->allocMemory(length.u64_, &pointer);
+  uint64_t written = std::min(buffer->length(), length.u64_);
+  if (written > 0) {
+    p = context->wasm()->allocMemory(written, &pointer);
     if (!p) {
       return wasmResultToWord(WasmResult::InvalidMemoryAccess);
     }
-    buffer->copyOut(start.u64_, length.u64_, p);
+    buffer->copyOut(start.u64_, written, p);
   }
   if (!context->wasmVm()->setWord(ptr_ptr.u64_, Word(pointer))) {
     return wasmResultToWord(WasmResult::InvalidMemoryAccess);
   }
-  if (!context->wasmVm()->setWord(size_ptr.u64_, Word(length.u64_))) {
+  if (!context->wasmVm()->setWord(size_ptr.u64_, Word(written))) {
     return wasmResultToWord(WasmResult::InvalidMemoryAccess);
   }
   return wasmResultToWord(WasmResult::Ok);
@@ -498,6 +499,46 @@ Word get_buffer_status(void* raw_context, Word type, Word length_ptr, Word flags
   return wasmResultToWord(WasmResult::Ok);
 }
 
+Word set_buffer_bytes(void* raw_context, Word type, Word start, Word length, Word data_ptr,
+                      Word data_size) {
+  if (type.u64_ > static_cast<uint64_t>(BufferType::MAX)) {
+    return wasmResultToWord(WasmResult::BadArgument);
+  }
+  auto context = WASM_CONTEXT(raw_context);
+  auto buffer = context->getBuffer(static_cast<BufferType>(type.u64_));
+  if (!buffer) {
+    return wasmResultToWord(WasmResult::NotFound);
+  }
+  auto data = context->wasmVm()->getMemory(data_ptr.u64_, data_size.u64_);
+  if (!data) {
+    return wasmResultToWord(WasmResult::InvalidMemoryAccess);
+  }
+  WasmResult result = WasmResult::InternalFailure;
+  if (context->setBuffer(static_cast<BufferType>(type.u64_),
+                         [start = start.u64_, length = length.u64_, data, &result](auto& buffer) {
+                           if (start == 0) {
+                             if (length == 0) {
+                               buffer.prepend(data.value());
+                               result = WasmResult::Ok;
+                             } else if (length >= buffer.length()) {
+                               buffer.drain(buffer.length());
+                               buffer.add(data.value());
+                               result = WasmResult::Ok;
+                             } else {
+                               result = WasmResult::BadArgument;
+                             }
+                           } else if (start >= buffer.length()) {
+                             buffer.add(data.value());
+                             result = WasmResult::Ok;
+                           } else {
+                             result = WasmResult::BadArgument;
+                           }
+                         }) != WasmResult::Ok) {
+    return wasmResultToWord(WasmResult::BadArgument);
+  }
+  return wasmResultToWord(result);
+}
+
 Word http_call(void* raw_context, Word uri_ptr, Word uri_size, Word header_pairs_ptr,
                Word header_pairs_size, Word body_ptr, Word body_size, Word trailer_pairs_ptr,
                Word trailer_pairs_size, Word timeout_milliseconds, Word token_ptr) {

diff --git a/source/extensions/common/wasm/exports.h b/source/extensions/common/wasm/exports.h
@@ -38,6 +38,8 @@ Word enqueue_shared_queue(void* raw_context, Word token, Word data_ptr, Word dat
 Word get_buffer_bytes(void* raw_context, Word type, Word start, Word length, Word ptr_ptr,
                       Word size_ptr);
 Word get_buffer_status(void* raw_context, Word type, Word length_ptr, Word flags_ptr);
+Word set_buffer_bytes(void* raw_context, Word type, Word start, Word length, Word data_ptr,
+                      Word data_size);
 Word add_header_map_value(void* raw_context, Word type, Word key_ptr, Word key_size, Word value_ptr,
                           Word value_size);
 Word get_header_map_value(void* raw_context, Word type, Word key_ptr, Word key_size,

diff --git a/source/extensions/common/wasm/null/wasm_api_impl.h b/source/extensions/common/wasm/null/wasm_api_impl.h
@@ -135,6 +135,12 @@ inline WasmResult proxy_get_buffer_status(BufferType type, size_t* length_ptr,
       Exports::get_buffer_status(current_context_, WS(type), WR(length_ptr), WR(flags_ptr)));
 }
 
+inline WasmResult proxy_set_buffer_bytes(BufferType type, uint64_t start, uint64_t length,
+                                         const char* data_ptr, size_t data_size) {
+  return wordToWasmResult(Exports::set_buffer_bytes(current_context_, WS(type), WS(start),
+                                                    WS(length), WR(data_ptr), WS(data_size)));
+}
+
 // Headers/Trailers/Metadata Maps
 inline WasmResult proxy_add_header_map_value(HeaderMapType type, const char* key_ptr,
                                              size_t key_size, const char* value_ptr,

diff --git a/source/extensions/common/wasm/wasm.cc b/source/extensions/common/wasm/wasm.cc
@@ -186,6 +186,7 @@ void Wasm::registerCallbacks() {
 
   _REGISTER_PROXY(get_buffer_status);
   _REGISTER_PROXY(get_buffer_bytes);
+  _REGISTER_PROXY(set_buffer_bytes);
 
   _REGISTER_PROXY(http_call);
 

diff --git a/test/extensions/filters/http/wasm/test_data/body_cpp.cc b/test/extensions/filters/http/wasm/test_data/body_cpp.cc
@@ -16,8 +16,8 @@ class ExampleContext : public Context {
   FilterDataStatus onResponseBody(size_t body_buffer_length, bool end_of_stream) override;
 
 private:
-  FilterDataStatus onBody(BufferType bt, size_t bufLen, bool end);
-  static void logBody(BufferType bt);
+  FilterDataStatus onBody(BufferType type, size_t buffer_length, bool end);
+  static void logBody(BufferType type);
 
   std::string test_op_;
   int num_chunks_ = 0;
@@ -34,27 +34,61 @@ FilterHeadersStatus ExampleContext::onResponseHeaders(uint32_t) {
   return FilterHeadersStatus::Continue;
 }
 
-void ExampleContext::logBody(BufferType bt) {
-  size_t bufferedSize;
+void ExampleContext::logBody(BufferType type) {
+  size_t buffered_size;
   uint32_t flags;
-  getBufferStatus(bt, &bufferedSize, &flags);
-  auto body = getBufferBytes(bt, 0, bufferedSize);
+  getBufferStatus(type, &buffered_size, &flags);
+  auto body = getBufferBytes(type, 0, buffered_size);
   logError(std::string("onRequestBody ") + std::string(body->view()));
 }
 
-FilterDataStatus ExampleContext::onBody(BufferType bt, size_t bufLen, bool end) {
+FilterDataStatus ExampleContext::onBody(BufferType type, size_t buffer_length, bool end_of_stream) {
   if (test_op_ == "ReadBody") {
-    auto body = getBufferBytes(bt, 0, bufLen);
+    auto body = getBufferBytes(type, 0, buffer_length);
     logError("onRequestBody " + std::string(body->view()));
 
+  } else if (test_op_ == "PrependAndAppendToBody") {
+    setBuffer(BufferType::HttpRequestBody, 0, 0, "prepend.");
+    setBuffer(BufferType::HttpRequestBody, 0xFFFFFFFF, 0, ".append");
+    auto updated = getBufferBytes(BufferType::HttpRequestBody, 0, 0xFFFFFFFF);
+    logError("onRequestBody " + std::string(updated->view()));
+
+  } else if (test_op_ == "ReplaceBody") {
+    setBuffer(BufferType::HttpRequestBody, 0, 0xFFFFFFFF, "replace");
+    auto replaced = getBufferBytes(BufferType::HttpRequestBody, 0, 0xFFFFFFFF);
+    logError("onRequestBody " + std::string(replaced->view()));
+
+  } else if (test_op_ == "RemoveBody") {
+    setBuffer(BufferType::HttpRequestBody, 0, 0xFFFFFFFF, "");
+    auto erased = getBufferBytes(BufferType::HttpRequestBody, 0, 0xFFFFFFFF);
+    logError("onRequestBody " + std::string(erased->view()));  
+
   } else if (test_op_ == "BufferBody") {
-    logBody(bt);
-    return end ? FilterDataStatus::Continue : FilterDataStatus::StopIterationAndBuffer;
+    logBody(type);
+    return end_of_stream ? FilterDataStatus::Continue : FilterDataStatus::StopIterationAndBuffer;
+
+  } else if (test_op_ == "PrependAndAppendToBufferedBody") {
+    setBuffer(BufferType::HttpRequestBody, 0, 0, "prepend.");
+    setBuffer(BufferType::HttpRequestBody, 0xFFFFFFFF, 0, ".append");
+    logBody(type);
+    return end_of_stream ? FilterDataStatus::Continue : FilterDataStatus::StopIterationAndBuffer;
+
+  } else if (test_op_ == "ReplaceBufferedBody") {
+    setBuffer(BufferType::HttpRequestBody, 0, 0xFFFFFFFF, "replace");
+    auto replaced = getBufferBytes(BufferType::HttpRequestBody, 0, 0xFFFFFFFF);
+    logBody(type);
+    return end_of_stream ? FilterDataStatus::Continue : FilterDataStatus::StopIterationAndBuffer;
+
+  } else if (test_op_ == "RemoveBufferedBody") {
+    setBuffer(BufferType::HttpRequestBody, 0, 0xFFFFFFFF, "");
+    auto erased = getBufferBytes(BufferType::HttpRequestBody, 0, 0xFFFFFFFF);
+    logBody(type);
+    return end_of_stream ? FilterDataStatus::Continue : FilterDataStatus::StopIterationAndBuffer;
 
   } else if (test_op_ == "BufferTwoBodies") {
-    logBody(bt);
+    logBody(type);
     num_chunks_++;
-    if (end || num_chunks_ > 2) {
+    if (end_of_stream || num_chunks_ > 2) {
       return FilterDataStatus::Continue;
     }
     return FilterDataStatus::StopIterationAndBuffer;
@@ -73,4 +107,4 @@ FilterDataStatus ExampleContext::onRequestBody(size_t body_buffer_length, bool e
 
 FilterDataStatus ExampleContext::onResponseBody(size_t body_buffer_length, bool end_of_stream) {
   return onBody(BufferType::HttpResponseBody, body_buffer_length, end_of_stream);
-}
+}
diff --git a/test/extensions/filters/http/wasm/test_data/body_cpp.wasm b/test/extensions/filters/http/wasm/test_data/body_cpp.wasm