restrict WASI

Signed-off-by: Ryan Apilado <[email protected]>
proxy-wasm · Nov 19, 2020 · 18a6367 · 18a6367
1 parent 8d24cf5
commit 18a6367
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 22 deletions.
diff --git a/include/proxy-wasm/exports.h b/include/proxy-wasm/exports.h
@@ -172,6 +172,11 @@ ::proxy_wasm::ContextBase *ContextOrEffectiveContext(::proxy_wasm::ContextBase *
   _f(get_configuration) _f(continue_request) _f(continue_response) _f(clear_route_cache)           \
       _f(continue_stream) _f(close_stream) _f(get_log_level)
 
+#define FOR_ALL_WASI_CAPABILITIES(_f)                                                              \
+  _f(fd_write) _f(fd_read) _f(fd_seek) _f(fd_close) _f(fd_fdstat_get) _f(environ_get)              \
+      _f(environ_sizes_get) _f(args_get) _f(args_sizes_get) _f(clock_time_get) _f(random_get)      \
+          _f(proc_exit)
+
 // Helpers to generate a stub to pass to VM, in place of a restricted export.
 #define _CREATE_EXPORT_STUB(_fn)                                                                   \
   template <typename F> struct _fn##Stub;                                                          \
@@ -182,9 +187,17 @@ ::proxy_wasm::ContextBase *ContextOrEffectiveContext(::proxy_wasm::ContextBase *
       context->wasmVm()->error("Attempted call to restricted capability: " #_fn);                  \
       return WasmResult::InternalFailure;                                                          \
     }                                                                                              \
+  };                                                                                               \
+  template <typename... Args> struct _fn##Stub<void(void *, Args...)> {                            \
+    static void stub(void *raw_context, Args...) {                                                 \
+      auto context = exports::ContextOrEffectiveContext(                                           \
+          static_cast<ContextBase *>((void)raw_context, current_context_));                        \
+      context->wasmVm()->error("Attempted call to restricted capability: " #_fn);                  \
+    }                                                                                              \
   };
-FOR_ALL_HOST_IMPLEMENTED_ABI_FUNCTIONS(_CREATE_EXPORT_STUB)
-FOR_ALL_HOST_IMPLEMENTED_ABI_FUNCTIONS_ABI_SPECIFIC(_CREATE_EXPORT_STUB)
+FOR_ALL_HOST_IMPLEMENTED_CAPABILITIES(_CREATE_EXPORT_STUB)
+FOR_ALL_HOST_IMPLEMENTED_CAPABILITIES_ABI_SPECIFIC(_CREATE_EXPORT_STUB)
+FOR_ALL_WASI_CAPABILITIES(_CREATE_EXPORT_STUB)
 #undef _CREATE_EXPORT_STUB
 
 } // namespace exports

diff --git a/src/wasm.cc b/src/wasm.cc
@@ -99,26 +99,28 @@ void WasmBase::registerCallbacks() {
 #undef _REGISTER
 
 #define _REGISTER_WASI(_fn)                                                                        \
-  wasm_vm_->registerCallback(                                                                      \
-      "wasi_unstable", #_fn, &exports::wasi_unstable_##_fn,                                        \
-      &ConvertFunctionWordToUint32<decltype(exports::wasi_unstable_##_fn),                         \
-                                   exports::wasi_unstable_##_fn>::convertFunctionWordToUint32);    \
-  wasm_vm_->registerCallback(                                                                      \
-      "wasi_snapshot_preview1", #_fn, &exports::wasi_unstable_##_fn,                               \
-      &ConvertFunctionWordToUint32<decltype(exports::wasi_unstable_##_fn),                         \
-                                   exports::wasi_unstable_##_fn>::convertFunctionWordToUint32)
-  _REGISTER_WASI(fd_write);
-  _REGISTER_WASI(fd_read);
-  _REGISTER_WASI(fd_seek);
-  _REGISTER_WASI(fd_close);
-  _REGISTER_WASI(fd_fdstat_get);
-  _REGISTER_WASI(environ_get);
-  _REGISTER_WASI(environ_sizes_get);
-  _REGISTER_WASI(args_get);
-  _REGISTER_WASI(args_sizes_get);
-  _REGISTER_WASI(clock_time_get);
-  _REGISTER_WASI(random_get);
-  _REGISTER_WASI(proc_exit);
+  if (capabilityAllowed(#_fn)) {                                                                   \
+    wasm_vm_->registerCallback(                                                                    \
+        "wasi_unstable", #_fn, &exports::wasi_unstable_##_fn,                                      \
+        &ConvertFunctionWordToUint32<decltype(exports::wasi_unstable_##_fn),                       \
+                                     exports::wasi_unstable_##_fn>::convertFunctionWordToUint32);  \
+    wasm_vm_->registerCallback(                                                                    \
+        "wasi_snapshot_preview1", #_fn, &exports::wasi_unstable_##_fn,                             \
+        &ConvertFunctionWordToUint32<decltype(exports::wasi_unstable_##_fn),                       \
+                                     exports::wasi_unstable_##_fn>::convertFunctionWordToUint32);  \
+  } else {                                                                                         \
+    typedef decltype(exports::wasi_unstable_##_fn) export_type;                                    \
+    constexpr export_type *stub = &exports::_fn##Stub<export_type>::stub;                          \
+    wasm_vm_->registerCallback(                                                                    \
+        "wasi_unstable", #_fn, stub,                                                               \
+        &ConvertFunctionWordToUint32<export_type, stub>::convertFunctionWordToUint32);             \
+    wasm_vm_->registerCallback(                                                                    \
+        "wasi_snapshot_preview1", #_fn, stub,                                                      \
+        &ConvertFunctionWordToUint32<export_type, stub>::convertFunctionWordToUint32);             \
+  }
+
+  FOR_ALL_WASI_CAPABILITIES(_REGISTER_WASI);
+
 #undef _REGISTER_WASI
 
   // Register the capability with the VM if it has been allowed, otherwise register a stub.