v8: set upper limit for maximum Wasm memory size. (#281)

Signed-off-by: Piotr Sikora <piotrsikora@google.com>

Author: PiotrSikora

BUILD

@@ -23,6 +23,7 @@ filegroup(
 cc_library(
     name = "wasm_vm_headers",
     hdrs = [
+        "include/proxy-wasm/limits.h",
         "include/proxy-wasm/wasm_vm.h",
         "include/proxy-wasm/word.h",
     ],

include/proxy-wasm/limits.h

@@ -0,0 +1,25 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+// Wasm memory page is always 64 KiB.
+#define PROXY_WASM_HOST_WASM_MEMORY_PAGE_SIZE_BYTES (64 * 1024)
+
+// Maximum allowed Wasm memory size.
+#ifndef PROXY_WASM_HOST_MAX_WASM_MEMORY_SIZE_BYTES
+#define PROXY_WASM_HOST_MAX_WASM_MEMORY_SIZE_BYTES (1024 * 1024 * 1024)
+#endif

src/v8/v8.cc

@@ -26,11 +26,17 @@
 #include <utility>
 #include <vector>
 
+#include "include/proxy-wasm/limits.h"
+
 #include "include/v8-version.h"
 #include "include/v8.h"
 #include "src/wasm/c-api.h"
 #include "wasm-api/wasm.hh"
 
+namespace v8::internal {
+extern unsigned int FLAG_wasm_max_mem_pages;
+} // namespace v8::internal
+
 namespace proxy_wasm {
 namespace v8 {
 
@@ -39,6 +45,8 @@ wasm::Engine *engine() {
   static wasm::own<wasm::Engine> engine;
 
   std::call_once(init, []() {
+    ::v8::internal::FLAG_wasm_max_mem_pages =
+        PROXY_WASM_HOST_MAX_WASM_MEMORY_SIZE_BYTES / PROXY_WASM_HOST_WASM_MEMORY_PAGE_SIZE_BYTES;
     ::v8::V8::EnableWebAssemblyTrapHandler(true);
     engine = wasm::Engine::make();
   });

test/BUILD

@@ -56,7 +56,7 @@ cc_test(
     data = [
         "//test/test_data:callback.wasm",
         "//test/test_data:clock.wasm",
-        "//test/test_data:infinite_loop.wasm",
+        "//test/test_data:resource_limits.wasm",
         "//test/test_data:trap.wasm",
     ],
     linkstatic = 1,

test/runtime_test.cc

@@ -19,6 +19,7 @@
 #include <thread>
 
 #include "include/proxy-wasm/context.h"
+#include "include/proxy-wasm/limits.h"
 #include "include/proxy-wasm/wasm.h"
 
 #include "test/utility.h"
@@ -86,7 +87,7 @@ TEST_P(TestVm, TerminateExecution) {
   if (engine_ != "v8") {
     return;
   }
-  auto source = readTestWasmFile("infinite_loop.wasm");
+  auto source = readTestWasmFile("resource_limits.wasm");
   ASSERT_FALSE(source.empty());
   auto wasm = TestWasm(std::move(vm_));
   ASSERT_TRUE(wasm.load(source, false));
@@ -112,6 +113,35 @@ TEST_P(TestVm, TerminateExecution) {
   }
 }
 
+TEST_P(TestVm, WasmMemoryLimit) {
+  // TODO(PiotrSikora): enforce memory limits in other engines.
+  if (engine_ != "v8") {
+    return;
+  }
+  auto source = readTestWasmFile("resource_limits.wasm");
+  ASSERT_FALSE(source.empty());
+  auto wasm = TestWasm(std::move(vm_));
+  ASSERT_TRUE(wasm.load(source, false));
+  ASSERT_TRUE(wasm.initialize());
+
+  WasmCallVoid<0> infinite_memory;
+  wasm.wasm_vm()->getFunction("infinite_memory", &infinite_memory);
+  ASSERT_TRUE(infinite_memory != nullptr);
+  infinite_memory(wasm.vm_context());
+
+  EXPECT_GE(wasm.wasm_vm()->getMemorySize(), PROXY_WASM_HOST_MAX_WASM_MEMORY_SIZE_BYTES * 0.95);
+  EXPECT_LE(wasm.wasm_vm()->getMemorySize(), PROXY_WASM_HOST_MAX_WASM_MEMORY_SIZE_BYTES);
+
+  // Check integration logs.
+  auto *host = dynamic_cast<TestIntegration *>(wasm.wasm_vm()->integration().get());
+  EXPECT_TRUE(host->isErrorLogged("Function: infinite_memory failed"));
+  if (engine_ == "v8") {
+    EXPECT_TRUE(host->isErrorLogged("Uncaught RuntimeError: unreachable"));
+    EXPECT_TRUE(host->isErrorLogged("Proxy-Wasm plugin in-VM backtrace:"));
+    EXPECT_TRUE(host->isErrorLogged(" - rust_oom"));
+  }
+}
+
 TEST_P(TestVm, Trap) {
   auto source = readTestWasmFile("trap.wasm");
   ASSERT_FALSE(source.empty());

test/test_data/BUILD

@@ -38,8 +38,8 @@ wasm_rust_binary(
 )
 
 wasm_rust_binary(
-    name = "infinite_loop.wasm",
-    srcs = ["infinite_loop.rs"],
+    name = "resource_limits.wasm",
+    srcs = ["resource_limits.rs"],
 )
 
 wasm_rust_binary(

test/test_data/infinite_loop.rs test/test_data/resource_limits.rstest/test_data/infinite_loop.rs renamed to test/test_data/resource_limits.rs

@@ -27,3 +27,11 @@ pub extern "C" fn infinite_loop() {
         _count += 1;
     }
 }
+
+#[no_mangle]
+pub extern "C" fn infinite_memory() {
+    let mut vec = Vec::new();
+    loop {
+       vec.push(Vec::<u32>::with_capacity(16384));
+    }
+}