v8: add restricted callbacks.

include/proxy-wasm/wasm.h

@@ -392,7 +392,15 @@ inline void *WasmBase::allocMemory(uint64_t size, uint64_t *address) {
   if (!malloc_) {
     return nullptr;
   }
+  wasm_vm_->setRestrictedCallback(
+      true, {// logging (Proxy-Wasm)
+             "env.proxy_log",
+             // logging (stdout/stderr)
+             "wasi_unstable.fd_write", "wasi_snapshot_preview1.fd_write",
+             // time
+             "wasi_unstable.clock_time_get", "wasi_snapshot_preview1.clock_time_get"});
   Word a = malloc_(vm_context(), size);
+  wasm_vm_->setRestrictedCallback(false);
   if (!a.u64_) {
     return nullptr;
   }

include/proxy-wasm/wasm_vm.h

@@ -20,6 +20,7 @@
 #include <optional>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 #include "include/proxy-wasm/word.h"
@@ -314,6 +315,16 @@ class WasmVm {
     fail_callbacks_.push_back(fail_callback);
   }
 
+  bool isHostFunctionAllowed(const std::string &name) {
+    return !restricted_callback_ || allowed_hostcalls_.find(name) != allowed_hostcalls_.end();
+  }
+
+  void setRestrictedCallback(bool restricted,
+                             std::unordered_set<std::string> allowed_hostcalls = {}) {
+    restricted_callback_ = restricted;
+    allowed_hostcalls_ = std::move(allowed_hostcalls);
+  }
+
   // Integrator operations.
   std::unique_ptr<WasmVmIntegration> &integration() { return integration_; }
   bool cmpLogLevel(proxy_wasm::LogLevel level) { return integration_->getLogLevel() <= level; }
@@ -322,6 +333,10 @@ class WasmVm {
   std::unique_ptr<WasmVmIntegration> integration_;
   FailState failed_ = FailState::Ok;
   std::vector<std::function<void(FailState)>> fail_callbacks_;
+
+private:
+  bool restricted_callback_{false};
+  std::unordered_set<std::string> allowed_hostcalls_{};
 };
 
 // Thread local state set during a call into a WASM VM so that calls coming out of the

src/v8/v8.cc

@@ -107,6 +107,8 @@ class V8 : public WasmVm {
   void terminate() override;
 
 private:
+  wasm::own<wasm::Trap> trap(std::string message);
+
   std::string getFailMessage(std::string_view function_name, wasm::own<wasm::Trap> trap);
 
   template <typename... Args>
@@ -503,6 +505,10 @@ bool V8::setWord(uint64_t pointer, Word word) {
   return true;
 }
 
+wasm::own<wasm::Trap> V8::trap(std::string message) {
+  return wasm::Trap::make(store_.get(), wasm::Message::make(std::move(message)));
+}
+
 template <typename... Args>
 void V8::registerHostFunctionImpl(std::string_view module_name, std::string_view function_name,
                                   void (*function)(Args...)) {
@@ -519,6 +525,9 @@ void V8::registerHostFunctionImpl(std::string_view module_name, std::string_view
           func_data->vm_->integration()->trace("[vm->host] " + func_data->name_ + "(" +
                                                printValues(params, sizeof...(Args)) + ")");
         }
+        if (!func_data->vm_->isHostFunctionAllowed(func_data->name_)) {
+          return dynamic_cast<V8 *>(func_data->vm_)->trap("restricted_callback");
+        }
         auto args = convertValTypesToArgsTuple<std::tuple<Args...>>(params);
         auto function = reinterpret_cast<void (*)(Args...)>(func_data->raw_func_);
         std::apply(function, args);
@@ -552,6 +561,9 @@ void V8::registerHostFunctionImpl(std::string_view module_name, std::string_view
           func_data->vm_->integration()->trace("[vm->host] " + func_data->name_ + "(" +
                                                printValues(params, sizeof...(Args)) + ")");
         }
+        if (!func_data->vm_->isHostFunctionAllowed(func_data->name_)) {
+          return dynamic_cast<V8 *>(func_data->vm_)->trap("restricted_callback");
+        }
         auto args = convertValTypesToArgsTuple<std::tuple<Args...>>(params);
         auto function = reinterpret_cast<R (*)(Args...)>(func_data->raw_func_);
         R rvalue = std::apply(function, args);

src/wasm.cc

@@ -355,6 +355,25 @@ ContextBase *WasmBase::getRootContext(const std::shared_ptr<PluginBase> &plugin,
 }
 
 void WasmBase::startVm(ContextBase *root_context) {
+  // wasi_snapshot_preview1.clock_time_get
+  wasm_vm_->setRestrictedCallback(
+      true, {// logging (Proxy-Wasm)
+             "env.proxy_log",
+             // logging (stdout/stderr)
+             "wasi_unstable.fd_write", "wasi_snapshot_preview1.fd_write",
+             // args
+             "wasi_unstable.args_sizes_get", "wasi_snapshot_preview1.args_sizes_get",
+             "wasi_unstable.args_get", "wasi_snapshot_preview1.args_get",
+             // environment variables
+             "wasi_unstable.environ_sizes_get", "wasi_snapshot_preview1.environ_sizes_get",
+             "wasi_unstable.environ_get", "wasi_snapshot_preview1.environ_get",
+             // preopened files/directories
+             "wasi_unstable.fd_prestat_get", "wasi_snapshot_preview1.fd_prestat_get",
+             "wasi_unstable.fd_prestat_dir_name", "wasi_snapshot_preview1.fd_prestat_dir_name",
+             // time
+             "wasi_unstable.clock_time_get", "wasi_snapshot_preview1.clock_time_get",
+             // random
+             "wasi_unstable.random_get", "wasi_snapshot_preview1.random_get"});
   if (_initialize_) {
     // WASI reactor.
     _initialize_(root_context);
@@ -370,6 +389,7 @@ void WasmBase::startVm(ContextBase *root_context) {
     // WASI command.
     _start_(root_context);
   }
+  wasm_vm_->setRestrictedCallback(false);
 }
 
 bool WasmBase::configure(ContextBase *root_context, std::shared_ptr<PluginBase> plugin) {

test/BUILD

@@ -102,6 +102,21 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "security_test",
+    srcs = ["security_test.cc"],
+    data = [
+        "//test/test_data:bad_malloc.wasm",
+    ],
+    linkstatic = 1,
+    deps = [
+        ":utility_lib",
+        "//:lib",
+        "@com_google_googletest//:gtest",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_test(
     name = "shared_data",
     srcs = ["shared_data_test.cc"],

test/security_test.cc

@@ -0,0 +1,106 @@
+// Copyright 2022 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.
+
+#include "gtest/gtest.h"
+
+#include <memory>
+#include <string>
+
+#include "include/proxy-wasm/context.h"
+#include "include/proxy-wasm/wasm.h"
+
+#include "test/utility.h"
+
+namespace proxy_wasm {
+
+INSTANTIATE_TEST_SUITE_P(WasmEngines, TestVm, testing::ValuesIn(getWasmEngines()),
+                         [](const testing::TestParamInfo<std::string> &info) {
+                           return info.param;
+                         });
+
+TEST_P(TestVm, MallocNoHostcalls) {
+  if (engine_ != "v8") {
+    return;
+  }
+  auto source = readTestWasmFile("bad_malloc.wasm");
+  ASSERT_FALSE(source.empty());
+  auto wasm = TestWasm(std::move(vm_));
+  ASSERT_TRUE(wasm.load(source, false));
+  ASSERT_TRUE(wasm.initialize());
+
+  uint64_t ptr = 0;
+  void *result = wasm.allocMemory(0x1000, &ptr);
+  EXPECT_NE(result, nullptr);
+  EXPECT_FALSE(wasm.isFailed());
+
+  // Check application logs.
+  auto *context = dynamic_cast<TestContext *>(wasm.vm_context());
+  EXPECT_TRUE(context->isLogEmpty());
+  // Check integration logs.
+  auto *integration = dynamic_cast<TestIntegration *>(wasm.wasm_vm()->integration().get());
+  EXPECT_FALSE(integration->isErrorLogged("Function: proxy_on_memory_allocate failed"));
+  EXPECT_FALSE(integration->isErrorLogged("restricted_callback"));
+}
+
+TEST_P(TestVm, MallocWithLog) {
+  if (engine_ != "v8") {
+    return;
+  }
+  auto source = readTestWasmFile("bad_malloc.wasm");
+  ASSERT_FALSE(source.empty());
+  auto wasm = TestWasm(std::move(vm_));
+  ASSERT_TRUE(wasm.load(source, false));
+  ASSERT_TRUE(wasm.initialize());
+
+  uint64_t ptr = 0;
+  // 0xAAAA => hostcall to proxy_log (allowed).
+  void *result = wasm.allocMemory(0xAAAA, &ptr);
+  EXPECT_NE(result, nullptr);
+  EXPECT_FALSE(wasm.isFailed());
+
+  // Check application logs.
+  auto *context = dynamic_cast<TestContext *>(wasm.vm_context());
+  EXPECT_TRUE(context->isLogged("this is fine"));
+  // Check integration logs.
+  auto *integration = dynamic_cast<TestIntegration *>(wasm.wasm_vm()->integration().get());
+  EXPECT_FALSE(integration->isErrorLogged("Function: proxy_on_memory_allocate failed"));
+  EXPECT_FALSE(integration->isErrorLogged("restricted_callback"));
+}
+
+TEST_P(TestVm, MallocWithHostcall) {
+  if (engine_ != "v8") {
+    return;
+  }
+  auto source = readTestWasmFile("bad_malloc.wasm");
+  ASSERT_FALSE(source.empty());
+  auto wasm = TestWasm(std::move(vm_));
+  ASSERT_TRUE(wasm.load(source, false));
+  ASSERT_TRUE(wasm.initialize());
+
+  uint64_t ptr = 0;
+  // 0xBBBB => hostcall to proxy_done (not allowed).
+  void *result = wasm.allocMemory(0xBBBB, &ptr);
+  EXPECT_EQ(result, nullptr);
+  EXPECT_TRUE(wasm.isFailed());
+
+  // Check application logs.
+  auto *context = dynamic_cast<TestContext *>(wasm.vm_context());
+  EXPECT_TRUE(context->isLogEmpty());
+  // Check integration logs.
+  auto *integration = dynamic_cast<TestIntegration *>(wasm.wasm_vm()->integration().get());
+  EXPECT_TRUE(integration->isErrorLogged("Function: proxy_on_memory_allocate failed"));
+  EXPECT_TRUE(integration->isErrorLogged("restricted_callback"));
+}
+
+} // namespace proxy_wasm

test/test_data/BUILD

@@ -28,6 +28,11 @@ wasm_rust_binary(
     ],
 )
 
+wasm_rust_binary(
+    name = "bad_malloc.wasm",
+    srcs = ["bad_malloc.rs"],
+)
+
 wasm_rust_binary(
     name = "callback.wasm",
     srcs = ["callback.rs"],

test/test_data/bad_malloc.rs

@@ -0,0 +1,47 @@
+// Copyright 2022 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 std::mem::MaybeUninit;
+
+extern "C" {
+    fn proxy_log(level: u32, message_data: *const u8, message_size: usize) -> u32;
+    fn proxy_done() -> u32;
+}
+
+#[no_mangle]
+pub extern "C" fn proxy_abi_version_0_2_0() {}
+
+#[no_mangle]
+pub extern "C" fn proxy_on_memory_allocate(size: usize) -> *mut u8 {
+    let mut vec: Vec<MaybeUninit<u8>> = Vec::with_capacity(size);
+    unsafe {
+        vec.set_len(size);
+    }
+    match size {
+        0xAAAA => {
+            let message = "this is fine";
+            unsafe {
+                proxy_log(0, message.as_ptr(), message.len());
+            }
+        }
+        0xBBBB => {
+            unsafe {
+                proxy_done();
+            }
+        }
+        _ => {}
+    }
+    let slice = vec.into_boxed_slice();
+    Box::into_raw(slice) as *mut u8
+}