Implement wasi.environ_{get,sizes_get}. (proxy-wasm#126)

Signed-off-by: Takeshi Yoneda <[email protected]>

Author: mathetake

bazel/wasm.bzl

@@ -19,6 +19,11 @@ def _wasm_rust_transition_impl(settings, attr):
         "//command_line_option:platforms": "@rules_rust//rust/platform:wasm",
     }
 
+def _wasi_rust_transition_impl(settings, attr):
+    return {
+        "//command_line_option:platforms": "@rules_rust//rust/platform:wasi",
+    }
+
 wasm_rust_transition = transition(
     implementation = _wasm_rust_transition_impl,
     inputs = [],
@@ -27,6 +32,14 @@ wasm_rust_transition = transition(
     ],
 )
 
+wasi_rust_transition = transition(
+    implementation = _wasi_rust_transition_impl,
+    inputs = [],
+    outputs = [
+        "//command_line_option:platforms",
+    ],
+)
+
 def _wasm_binary_impl(ctx):
     out = ctx.actions.declare_file(ctx.label.name)
     ctx.actions.run(
@@ -49,7 +62,12 @@ wasm_rust_binary_rule = rule(
     attrs = _wasm_attrs(wasm_rust_transition),
 )
 
-def wasm_rust_binary(name, tags = [], **kwargs):
+wasi_rust_binary_rule = rule(
+    implementation = _wasm_binary_impl,
+    attrs = _wasm_attrs(wasi_rust_transition),
+)
+
+def wasm_rust_binary(name, tags = [], wasi = False, **kwargs):
     wasm_name = "_wasm_" + name.replace(".", "_")
     kwargs.setdefault("visibility", ["//visibility:public"])
 
@@ -62,7 +80,11 @@ def wasm_rust_binary(name, tags = [], **kwargs):
         **kwargs
     )
 
-    wasm_rust_binary_rule(
+    bin_rule = wasm_rust_binary_rule
+    if wasi:
+        bin_rule = wasi_rust_binary_rule
+
+    bin_rule(
         name = name,
         binary = ":" + wasm_name,
         tags = tags + ["manual"],

include/proxy-wasm/null_vm.h

@@ -43,6 +43,7 @@ struct NullVm : public WasmVm {
   bool setMemory(uint64_t pointer, uint64_t size, const void *data) override;
   bool setWord(uint64_t pointer, Word data) override;
   bool getWord(uint64_t pointer, Word *data) override;
+  size_t getWordSize() override;
   std::string_view getCustomSection(std::string_view name) override;
   std::string_view getPrecompiledSectionName() override;
 

include/proxy-wasm/wasm.h

@@ -53,6 +53,7 @@ class WasmBase : public std::enable_shared_from_this<WasmBase> {
 public:
   WasmBase(std::unique_ptr<WasmVm> wasm_vm, std::string_view vm_id,
            std::string_view vm_configuration, std::string_view vm_key,
+           std::unordered_map<std::string, std::string> envs,
            AllowedCapabilitiesMap allowed_capabilities);
   WasmBase(const std::shared_ptr<WasmHandleBase> &other, WasmVmFactory factory);
   virtual ~WasmBase();
@@ -136,6 +137,8 @@ class WasmBase : public std::enable_shared_from_this<WasmBase> {
 
   AbiVersion abiVersion() { return abi_version_; }
 
+  const std::unordered_map<std::string, std::string> &envs() { return envs_; }
+
   // Called to raise the flag which indicates that the context should stop iteration regardless of
   // returned filter status from Proxy-Wasm extensions. For example, we ignore
   // FilterHeadersStatus::Continue after a local reponse is sent by the host.
@@ -190,6 +193,8 @@ class WasmBase : public std::enable_shared_from_this<WasmBase> {
   std::unordered_map<uint32_t, ContextBase *> contexts_;                 // Contains all contexts.
   std::unordered_map<uint32_t, std::chrono::milliseconds> timer_period_; // per root_id.
   std::unique_ptr<ShutdownHandle> shutdown_handle_;
+  std::unordered_map<std::string, std::string>
+      envs_; // environment variables passed through wasi.environ_get
 
   WasmCallVoid<0> _initialize_; /* Emscripten v1.39.17+ */
   WasmCallVoid<0> _start_;      /* Emscripten v1.39.0+ */

include/proxy-wasm/wasm_vm.h

@@ -243,6 +243,11 @@ class WasmVm {
    */
   virtual bool setWord(uint64_t pointer, Word data) = 0;
 
+  /**
+   * @return the Word size in this VM.
+   */
+  virtual size_t getWordSize() = 0;
+
   /**
    * Get the contents of the custom section with the given name or "" if it does not exist.
    * @param name the name of the custom section to get.

src/exports.cc

@@ -767,18 +767,46 @@ Word wasi_unstable_fd_fdstat_get(void *raw_context, Word fd, Word statOut) {
 }
 
 // __wasi_errno_t __wasi_environ_get(char **environ, char *environ_buf);
-Word wasi_unstable_environ_get(void *, Word, Word) {
+Word wasi_unstable_environ_get(void *raw_context, Word environ_array_ptr, Word environ_buf) {
+  auto context = WASM_CONTEXT(raw_context);
+  auto word_size = context->wasmVm()->getWordSize();
+  auto &envs = context->wasm()->envs();
+  for (auto e : envs) {
+    if (!context->wasmVm()->setWord(environ_array_ptr, environ_buf)) {
+      return 21; // __WASI_EFAULT
+    }
+
+    std::string data;
+    data.reserve(e.first.size() + e.second.size() + 2);
+    data.append(e.first);
+    data.append("=");
+    data.append(e.second);
+    data.append("\x0");
+    if (!context->wasmVm()->setMemory(environ_buf, data.size(), data.c_str())) {
+      return 21; // __WASI_EFAULT
+    }
+    environ_buf = environ_buf.u64_ + data.size();
+    environ_array_ptr = environ_array_ptr.u64_ + word_size;
+  }
+
   return 0; // __WASI_ESUCCESS
 }
 
 // __wasi_errno_t __wasi_environ_sizes_get(size_t *environ_count, size_t
 // *environ_buf_size);
 Word wasi_unstable_environ_sizes_get(void *raw_context, Word count_ptr, Word buf_size_ptr) {
   auto context = WASM_CONTEXT(raw_context);
-  if (!context->wasmVm()->setWord(count_ptr, Word(0))) {
+  auto &envs = context->wasm()->envs();
+  if (!context->wasmVm()->setWord(count_ptr, Word(envs.size()))) {
     return 21; // __WASI_EFAULT
   }
-  if (!context->wasmVm()->setWord(buf_size_ptr, Word(0))) {
+
+  size_t size = 0;
+  for (auto e : envs) {
+    // len(key) + len(value) + 1('=') + 1(null terminator)
+    size += e.first.size() + e.second.size() + 2;
+  }
+  if (!context->wasmVm()->setWord(buf_size_ptr, Word(size))) {
     return 21; // __WASI_EFAULT
   }
   return 0; // __WASI_ESUCCESS

src/null/null_vm.cc

@@ -102,6 +102,8 @@ bool NullVm::getWord(uint64_t pointer, Word *data) {
   return true;
 }
 
+size_t NullVm::getWordSize() { return sizeof(uint64_t); }
+
 std::string_view NullVm::getCustomSection(std::string_view /* name */) {
   // Return nothing: there is no WASM file.
   return {};

src/v8/v8.cc

@@ -67,6 +67,7 @@ class V8 : public WasmVm {
   bool setMemory(uint64_t pointer, uint64_t size, const void *data) override;
   bool getWord(uint64_t pointer, Word *word) override;
   bool setWord(uint64_t pointer, Word word) override;
+  size_t getWordSize() override { return sizeof(uint32_t); };
 
 #define _REGISTER_HOST_FUNCTION(T)                                                                 \
   void registerCallback(std::string_view module_name, std::string_view function_name, T,           \

src/wasm.cc

@@ -193,6 +193,7 @@ WasmBase::WasmBase(const std::shared_ptr<WasmHandleBase> &base_wasm_handle, Wasm
     : std::enable_shared_from_this<WasmBase>(*base_wasm_handle->wasm()),
       vm_id_(base_wasm_handle->wasm()->vm_id_), vm_key_(base_wasm_handle->wasm()->vm_key_),
       started_from_(base_wasm_handle->wasm()->wasm_vm()->cloneable()),
+      envs_(base_wasm_handle->wasm()->envs()),
       allowed_capabilities_(base_wasm_handle->wasm()->allowed_capabilities_),
       base_wasm_handle_(base_wasm_handle) {
   if (started_from_ != Cloneable::NotCloneable) {
@@ -209,9 +210,10 @@ WasmBase::WasmBase(const std::shared_ptr<WasmHandleBase> &base_wasm_handle, Wasm
 
 WasmBase::WasmBase(std::unique_ptr<WasmVm> wasm_vm, std::string_view vm_id,
                    std::string_view vm_configuration, std::string_view vm_key,
+                   std::unordered_map<std::string, std::string> envs,
                    AllowedCapabilitiesMap allowed_capabilities)
     : vm_id_(std::string(vm_id)), vm_key_(std::string(vm_key)), wasm_vm_(std::move(wasm_vm)),
-      allowed_capabilities_(std::move(allowed_capabilities)),
+      envs_(envs), allowed_capabilities_(std::move(allowed_capabilities)),
       vm_configuration_(std::string(vm_configuration)), vm_id_handle_(getVmIdHandle(vm_id)) {
   if (!wasm_vm_) {
     failed_ = FailState::UnableToCreateVM;

src/wasmtime/wasmtime.cc

@@ -65,6 +65,7 @@ class Wasmtime : public WasmVm {
   bool setMemory(uint64_t pointer, uint64_t size, const void *data) override;
   bool getWord(uint64_t pointer, Word *word) override;
   bool setWord(uint64_t pointer, Word word) override;
+  size_t getWordSize() override { return sizeof(uint32_t); };
 
 #define _REGISTER_HOST_FUNCTION(T)                                                                 \
   void registerCallback(std::string_view module_name, std::string_view function_name, T,           \

src/wavm/wavm.cc

@@ -229,6 +229,7 @@ struct Wavm : public WasmVm {
   bool setMemory(uint64_t pointer, uint64_t size, const void *data) override;
   bool getWord(uint64_t pointer, Word *data) override;
   bool setWord(uint64_t pointer, Word data) override;
+  size_t getWordSize() override { return sizeof(uint32_t); };
   std::string_view getCustomSection(std::string_view name) override;
   std::string_view getPrecompiledSectionName() override;
   AbiVersion getAbiVersion() override;

test/BUILD

@@ -1,4 +1,4 @@
-load("@rules_cc//cc:defs.bzl", "cc_test")
+load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test")
 load("@proxy_wasm_cpp_host//bazel:variables.bzl", "COPTS", "LINKOPTS")
 
 cc_test(
@@ -23,6 +23,23 @@ cc_test(
     ],
     linkopts = LINKOPTS,
     deps = [
+        ":utility_lib",
+        "//:lib",
+        "@com_google_googletest//:gtest",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "exports_test",
+    srcs = ["exports_test.cc"],
+    copts = COPTS,
+    data = [
+        "//test/test_data:env.wasm",
+    ],
+    linkopts = LINKOPTS,
+    deps = [
+        ":utility_lib",
         "//:lib",
         "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",
@@ -72,3 +89,17 @@ cc_test(
         "@com_google_googletest//:gtest_main",
     ],
 )
+
+cc_library(
+    name = "utility_lib",
+    srcs = [
+        "utility.cc",
+        "utility.h",
+    ],
+    hdrs = ["utility.h"],
+    copts = COPTS,
+    deps = [
+        "//:lib",
+        "@com_google_googletest//:gtest",
+    ],
+)

test/exports_test.cc

@@ -0,0 +1,95 @@
+// Copyright 2021 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 <fstream>
+#include <iostream>
+#include <memory>
+#include <sstream>
+#include <string>
+#include <vector>
+
+#include "include/proxy-wasm/context.h"
+#include "include/proxy-wasm/exports.h"
+#include "include/proxy-wasm/wasm.h"
+
+#include "test/utility.h"
+
+namespace proxy_wasm {
+namespace {
+
+auto test_values = testing::ValuesIn(getRuntimes());
+
+INSTANTIATE_TEST_SUITE_P(Runtimes, TestVM, test_values);
+
+class TestContext : public ContextBase {
+public:
+  TestContext(WasmBase *base) : ContextBase(base){};
+  WasmResult log(uint32_t, std::string_view msg) override {
+    log_ += std::string(msg) + "\n";
+    return WasmResult::Ok;
+  }
+  std::string &log_msg() { return log_; }
+
+private:
+  std::string log_;
+};
+
+TEST_P(TestVM, Environment) {
+  std::unordered_map<std::string, std::string> envs = {{"KEY1", "VALUE1"}, {"KEY2", "VALUE2"}};
+  initialize("env.wasm");
+
+  auto wasm_base = WasmBase(std::move(vm_), "vm_id", "", "", envs, {});
+  ASSERT_TRUE(wasm_base.wasm_vm()->load(source_, false));
+
+  TestContext context(&wasm_base);
+  current_context_ = &context;
+
+  wasm_base.registerCallbacks();
+
+  ASSERT_TRUE(wasm_base.wasm_vm()->link(""));
+
+  WasmCallVoid<0> run;
+  wasm_base.wasm_vm()->getFunction("run", &run);
+
+  run(current_context_);
+
+  auto msg = context.log_msg();
+  EXPECT_NE(std::string::npos, msg.find("KEY1: VALUE1")) << msg;
+  EXPECT_NE(std::string::npos, msg.find("KEY2: VALUE2")) << msg;
+}
+
+TEST_P(TestVM, WithoutEnvironment) {
+  initialize("env.wasm");
+  auto wasm_base = WasmBase(std::move(vm_), "vm_id", "", "", {}, {});
+  ASSERT_TRUE(wasm_base.wasm_vm()->load(source_, false));
+
+  TestContext context(&wasm_base);
+  current_context_ = &context;
+
+  wasm_base.registerCallbacks();
+
+  ASSERT_TRUE(wasm_base.wasm_vm()->link(""));
+
+  WasmCallVoid<0> run;
+  wasm_base.wasm_vm()->getFunction("run", &run);
+
+  run(current_context_);
+
+  EXPECT_EQ(context.log_msg(), "");
+}
+
+} // namespace
+} // namespace proxy_wasm