wasm: support precompiled modules in V8-based runtime. (envoyproxy#9691)

Signed-off-by: Piotr Sikora <[email protected]>

bazel/external/wee8.BUILD

@@ -8,9 +8,13 @@ cc_library(
         "libwee8.a",
     ],
     hdrs = [
+        "wee8/include/v8-version.h",
         "wee8/third_party/wasm-api/wasm.hh",
     ],
-    includes = ["wee8/third_party"],
+    includes = [
+        "wee8/include",
+        "wee8/third_party",
+    ],
     visibility = ["//visibility:public"],
 )
 

source/extensions/common/wasm/null/null_vm.cc

@@ -82,6 +82,11 @@ absl::string_view NullVm::getCustomSection(absl::string_view /* name */) {
   return {};
 }
 
+absl::string_view NullVm::getPrecompiledSectionName() {
+  // Return nothing: there is no WASM file.
+  return {};
+}
+
 } // namespace Null
 } // namespace Wasm
 } // namespace Common

source/extensions/common/wasm/null/null_vm.h

@@ -38,6 +38,7 @@ struct NullVm : public WasmVmBase {
   bool setWord(uint64_t pointer, Word data) override;
   bool getWord(uint64_t pointer, Word* data) override;
   absl::string_view getCustomSection(absl::string_view name) override;
+  absl::string_view getPrecompiledSectionName() override;
 
 #define _FORWARD_GET_FUNCTION(_T)                                                                  \
   void getFunction(absl::string_view function_name, _T* f) override {                              \

source/extensions/common/wasm/v8/v8.cc

@@ -11,6 +11,7 @@
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/strings/match.h"
+#include "v8-version.h"
 #include "wasm-api/wasm.hh"
 
 namespace Envoy {
@@ -43,6 +44,7 @@ class V8 : public WasmVmBase {
 
   bool load(const std::string& code, bool allow_precompiled) override;
   absl::string_view getCustomSection(absl::string_view name) override;
+  absl::string_view getPrecompiledSectionName() override;
   void link(absl::string_view debug_name) override;
 
   Cloneable cloneable() override { return Cloneable::CompiledBytecode; }
@@ -70,6 +72,8 @@ class V8 : public WasmVmBase {
 #undef _GET_MODULE_FUNCTION
 
 private:
+  wasm::vec<byte_t> getStrippedSource();
+
   template <typename... Args>
   void registerHostFunctionImpl(absl::string_view module_name, absl::string_view function_name,
                                 void (*function)(void*, Args...));
@@ -239,14 +243,46 @@ template <typename T, typename U> constexpr T convertValTypesToArgsTuple(const U
 
 // V8 implementation.
 
-bool V8::load(const std::string& code, bool /* allow_precompiled */) {
+bool V8::load(const std::string& code, bool allow_precompiled) {
   ENVOY_LOG(trace, "load()");
   store_ = wasm::Store::make(engine());
 
+  // Wasm file header is 8 bytes (magic number + version).
+  static const uint8_t magic_number[4] = {0x00, 0x61, 0x73, 0x6d};
+  if (code.size() < 8 || ::memcmp(code.data(), magic_number, 4) != 0) {
+    return false;
+  }
+
   source_ = wasm::vec<byte_t>::make_uninitialized(code.size());
   ::memcpy(source_.get(), code.data(), code.size());
 
-  module_ = wasm::Module::make(store_.get(), source_);
+  if (allow_precompiled) {
+    const auto section_name = getPrecompiledSectionName();
+    if (!section_name.empty()) {
+      const auto precompiled = getCustomSection(section_name);
+      if (!precompiled.empty()) {
+        auto vec = wasm::vec<byte_t>::make_uninitialized(precompiled.size());
+        ::memcpy(vec.get(), precompiled.data(), precompiled.size());
+
+        // TODO(PiotrSikora): fuzz loading of precompiled Wasm modules.
+        // See: https://github.com/envoyproxy/envoy/issues/9731
+        module_ = wasm::Module::deserialize(store_.get(), vec);
+        if (!module_) {
+          // Precompiled module that cannot be loaded is considered a hard error,
+          // so don't fallback to compiling the bytecode.
+          return false;
+        }
+      }
+    }
+  }
+
+  if (!module_) {
+    // TODO(PiotrSikora): fuzz loading of Wasm modules.
+    // See: https://github.com/envoyproxy/envoy/issues/9731
+    const auto stripped_source = getStrippedSource();
+    module_ = wasm::Module::make(store_.get(), stripped_source ? stripped_source : source_);
+  }
+
   if (module_) {
     shared_module_ = module_->share();
   }
@@ -266,39 +302,95 @@ WasmVmPtr V8::clone() {
   return clone;
 }
 
+// Get Wasm module without Custom Sections to save some memory in workers.
+wasm::vec<byte_t> V8::getStrippedSource() {
+  ENVOY_LOG(trace, "getStrippedSource()");
+  ASSERT(source_.get() != nullptr);
+
+  std::vector<byte_t> stripped;
+
+  const byte_t* pos = source_.get() + 8 /* Wasm header */;
+  const byte_t* end = source_.get() + source_.size();
+  while (pos < end) {
+    const auto section_start = pos;
+    if (pos + 1 > end) {
+      return wasm::vec<byte_t>::invalid();
+    }
+    const auto section_type = *pos++;
+    const auto section_len = parseVarint(pos, end);
+    if (section_len == static_cast<uint32_t>(-1) || pos + section_len > end) {
+      return wasm::vec<byte_t>::invalid();
+    }
+    pos += section_len;
+    if (section_type == 0 /* custom section */) {
+      if (stripped.empty()) {
+        const byte_t* start = source_.get();
+        stripped.insert(stripped.end(), start, section_start);
+      }
+    } else if (!stripped.empty()) {
+      stripped.insert(stripped.end(), section_start, pos /* section end */);
+    }
+  }
+
+  // No custom sections found, use the original source.
+  if (stripped.empty()) {
+    return wasm::vec<byte_t>::invalid();
+  }
+
+  // Return stripped source, without custom sections.
+  return wasm::vec<byte_t>::make(stripped.size(), stripped.data());
+}
+
 absl::string_view V8::getCustomSection(absl::string_view name) {
   ENVOY_LOG(trace, "getCustomSection(\"{}\")", name);
   ASSERT(source_.get() != nullptr);
 
+  const byte_t* pos = source_.get() + 8 /* Wasm header */;
   const byte_t* end = source_.get() + source_.size();
-  const byte_t* pos = source_.get() + 8; // skip header
   while (pos < end) {
     if (pos + 1 > end) {
       throw WasmVmException("Failed to parse corrupted WASM module");
     }
-    auto type = *pos++;
-    auto rest = parseVarint(pos, end);
-    if (pos + rest > end) {
+    const auto section_type = *pos++;
+    const auto section_len = parseVarint(pos, end);
+    if (section_len == static_cast<uint32_t>(-1) || pos + section_len > end) {
       throw WasmVmException("Failed to parse corrupted WASM module");
     }
-    if (type == 0 /* custom section */) {
-      auto start = pos;
-      auto len = parseVarint(pos, end);
-      if (pos + len > end) {
+    if (section_type == 0 /* custom section */) {
+      const auto section_data_start = pos;
+      const auto section_name_len = parseVarint(pos, end);
+      if (section_name_len == static_cast<uint32_t>(-1) || pos + section_name_len > end) {
         throw WasmVmException("Failed to parse corrupted WASM module");
       }
-      pos += len;
-      rest -= (pos - start);
-      if (len == name.size() && ::memcmp(pos - len, name.data(), len) == 0) {
-        ENVOY_LOG(trace, "getCustomSection(\"{}\") found, size: {}", name, rest);
-        return {pos, rest};
+      if (section_name_len == name.size() && ::memcmp(pos, name.data(), section_name_len) == 0) {
+        pos += section_name_len;
+        ENVOY_LOG(trace, "getCustomSection(\"{}\") found, size: {}", name,
+                  section_data_start + section_len - pos);
+        return {pos, static_cast<size_t>(section_data_start + section_len - pos)};
       }
+      pos = section_data_start + section_len;
+    } else {
+      pos += section_len;
     }
-    pos += rest;
   }
   return "";
 }
 
+#if defined(__linux__) && defined(__x86_64__)
+#define WEE8_WASM_PRECOMPILE_PLATFORM "linux_x86_64"
+#endif
+
+absl::string_view V8::getPrecompiledSectionName() {
+#ifndef WEE8_WASM_PRECOMPILE_PLATFORM
+  return "";
+#else
+  static const auto name =
+      absl::StrCat("precompiled_v8_v", V8_MAJOR_VERSION, ".", V8_MINOR_VERSION, ".",
+                   V8_BUILD_NUMBER, ".", V8_PATCH_LEVEL, "_", WEE8_WASM_PRECOMPILE_PLATFORM);
+  return name;
+#endif
+}
+
 void V8::link(absl::string_view debug_name) {
   ENVOY_LOG(trace, "link(\"{}\")", debug_name);
   ASSERT(module_ != nullptr);

source/extensions/common/wasm/wasm_vm.h

@@ -235,6 +235,12 @@ class WasmVm : public Logger::Loggable<Logger::Id::wasm> {
    */
   virtual absl::string_view getCustomSection(absl::string_view name) PURE;
 
+  /**
+   * Get the name of the custom section that contains precompiled module.
+   * @return the name of the custom section that contains precompiled module.
+   */
+  virtual absl::string_view getPrecompiledSectionName() PURE;
+
   /**
    * Get typed function exported by the WASM module.
    */

test/extensions/common/wasm/test_data/test_rust.rs

@@ -1,5 +1,9 @@
+// TODO(PiotrSikora): build test data with Bazel rules.
+// See: https://github.com/envoyproxy/envoy/issues/9733
+//
 // Build using:
 // $ rustc -C lto -C opt-level=3 -C panic=abort -C link-arg=-S -C link-arg=-zstack-size=32768 --crate-type cdylib --target wasm32-unknown-unknown test_rust.rs
+// $ ../../../../../bazel-bin/test/tools/wee8_compile/wee8_compile_tool test_rust.wasm test_rust.wasm
 
 // Import functions exported from the host environment.
 extern "C" {

test/extensions/common/wasm/wasm_vm_test.cc

@@ -125,42 +125,54 @@ Word bad_pong2(void*, Word) { return 2; }
 // pong() with wrong argument type.
 double bad_pong3(void*, double) { return 3; }
 
-class WasmVmTest : public BaseVmTest {
+class WasmVmTest : public testing::TestWithParam<bool> {
 public:
+  WasmVmTest() : scope_(Stats::ScopeSharedPtr(stats_store.createScope("wasm."))) {}
+
   void SetUp() override { g_host_functions = new MockHostFunctions(); }
   void TearDown() override { delete g_host_functions; }
+
+protected:
+  Stats::IsolatedStoreImpl stats_store;
+  Stats::ScopeSharedPtr scope_;
 };
 
-TEST_F(WasmVmTest, V8BadCode) {
+INSTANTIATE_TEST_SUITE_P(AllowPrecompiled, WasmVmTest, testing::Values(false, true));
+
+TEST_P(WasmVmTest, V8BadCode) {
   auto wasm_vm = createWasmVm("envoy.wasm.runtime.v8", scope_);
   ASSERT_TRUE(wasm_vm != nullptr);
 
-  EXPECT_FALSE(wasm_vm->load("bad code", false));
+  EXPECT_FALSE(wasm_vm->load("bad code", GetParam()));
 }
 
-TEST_F(WasmVmTest, V8Code) {
+TEST_P(WasmVmTest, V8Code) {
   auto wasm_vm = createWasmVm("envoy.wasm.runtime.v8", scope_);
   ASSERT_TRUE(wasm_vm != nullptr);
   EXPECT_TRUE(wasm_vm->runtime() == "envoy.wasm.runtime.v8");
 
   auto code = TestEnvironment::readFileToStringForTest(TestEnvironment::substitute(
       "{{ test_rundir }}/test/extensions/common/wasm/test_data/test_rust.wasm"));
-  EXPECT_TRUE(wasm_vm->load(code, false));
+  EXPECT_TRUE(wasm_vm->load(code, GetParam()));
 
+  // Sanity checks for the expected test file.
+  if (!wasm_vm->getPrecompiledSectionName().empty()) {
+    EXPECT_TRUE(!wasm_vm->getCustomSection(wasm_vm->getPrecompiledSectionName()).empty());
+  }
   EXPECT_THAT(wasm_vm->getCustomSection("producers"), HasSubstr("rustc"));
   EXPECT_TRUE(wasm_vm->getCustomSection("emscripten_metadata").empty());
 
   EXPECT_TRUE(wasm_vm->cloneable() == Cloneable::CompiledBytecode);
   EXPECT_TRUE(wasm_vm->clone() != nullptr);
 }
 
-TEST_F(WasmVmTest, V8BadHostFunctions) {
+TEST_P(WasmVmTest, V8BadHostFunctions) {
   auto wasm_vm = createWasmVm("envoy.wasm.runtime.v8", scope_);
   ASSERT_TRUE(wasm_vm != nullptr);
 
   auto code = TestEnvironment::readFileToStringForTest(TestEnvironment::substitute(
       "{{ test_rundir }}/test/extensions/common/wasm/test_data/test_rust.wasm"));
-  EXPECT_TRUE(wasm_vm->load(code, false));
+  EXPECT_TRUE(wasm_vm->load(code, GetParam()));
 
   wasm_vm->registerCallback("env", "random", &random, CONVERT_FUNCTION_WORD_TO_UINT32(random));
   EXPECT_THROW_WITH_MESSAGE(wasm_vm->link("test"), WasmVmException,
@@ -182,13 +194,13 @@ TEST_F(WasmVmTest, V8BadHostFunctions) {
                             "want: i32 -> void, but host exports: f64 -> f64");
 }
 
-TEST_F(WasmVmTest, V8BadModuleFunctions) {
+TEST_P(WasmVmTest, V8BadModuleFunctions) {
   auto wasm_vm = createWasmVm("envoy.wasm.runtime.v8", scope_);
   ASSERT_TRUE(wasm_vm != nullptr);
 
   auto code = TestEnvironment::readFileToStringForTest(TestEnvironment::substitute(
       "{{ test_rundir }}/test/extensions/common/wasm/test_data/test_rust.wasm"));
-  EXPECT_TRUE(wasm_vm->load(code, false));
+  EXPECT_TRUE(wasm_vm->load(code, GetParam()));
 
   wasm_vm->registerCallback("env", "pong", &pong, CONVERT_FUNCTION_WORD_TO_UINT32(pong));
   wasm_vm->registerCallback("env", "random", &random, CONVERT_FUNCTION_WORD_TO_UINT32(random));
@@ -210,13 +222,13 @@ TEST_F(WasmVmTest, V8BadModuleFunctions) {
                             "Bad function signature for: sum");
 }
 
-TEST_F(WasmVmTest, V8FunctionCalls) {
+TEST_P(WasmVmTest, V8FunctionCalls) {
   auto wasm_vm = createWasmVm("envoy.wasm.runtime.v8", scope_);
   ASSERT_TRUE(wasm_vm != nullptr);
 
   auto code = TestEnvironment::readFileToStringForTest(TestEnvironment::substitute(
       "{{ test_rundir }}/test/extensions/common/wasm/test_data/test_rust.wasm"));
-  EXPECT_TRUE(wasm_vm->load(code, false));
+  EXPECT_TRUE(wasm_vm->load(code, GetParam()));
 
   wasm_vm->registerCallback("env", "pong", &pong, CONVERT_FUNCTION_WORD_TO_UINT32(pong));
   wasm_vm->registerCallback("env", "random", &random, CONVERT_FUNCTION_WORD_TO_UINT32(random));
@@ -248,13 +260,13 @@ TEST_F(WasmVmTest, V8FunctionCalls) {
                             "Function: abort failed: Uncaught RuntimeError: unreachable");
 }
 
-TEST_F(WasmVmTest, V8Memory) {
+TEST_P(WasmVmTest, V8Memory) {
   auto wasm_vm = createWasmVm("envoy.wasm.runtime.v8", scope_);
   ASSERT_TRUE(wasm_vm != nullptr);
 
   auto code = TestEnvironment::readFileToStringForTest(TestEnvironment::substitute(
       "{{ test_rundir }}/test/extensions/common/wasm/test_data/test_rust.wasm"));
-  EXPECT_TRUE(wasm_vm->load(code, false));
+  EXPECT_TRUE(wasm_vm->load(code, GetParam()));
 
   wasm_vm->registerCallback("env", "pong", &pong, CONVERT_FUNCTION_WORD_TO_UINT32(pong));
   wasm_vm->registerCallback("env", "random", &random, CONVERT_FUNCTION_WORD_TO_UINT32(random));

test/tools/wee8_compile/BUILD

@@ -0,0 +1,21 @@
+licenses(["notice"])  # Apache 2
+
+load(
+    "//bazel:envoy_build_system.bzl",
+    "envoy_cc_test_binary",
+    "envoy_cc_test_library",
+    "envoy_package",
+)
+
+envoy_package()
+
+envoy_cc_test_binary(
+    name = "wee8_compile_tool",
+    deps = [":wee8_compile_lib"],
+)
+
+envoy_cc_test_library(
+    name = "wee8_compile_lib",
+    srcs = ["wee8_compile.cc"],
+    external_deps = ["wee8"],
+)