[pku][wasm] Refactor PKU usage in Wasm

RwxMemoryWriteScope becomes the bottleneck for both MAP_JIT and PKU
machinery.
Wasm and V8 code space will use the same memory protection key.

This is a next step towards adding PKU support for V8 code space.

Bug: v8:13023
Change-Id: I647f8c09bc41e5ef8a1d74b58a48a43e08454e0d
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3702213Reviewed-by: Michael Lippautz <mlippautz@chromium.org>
Commit-Queue: Wenqin Yang <wenqin.yang@intel.com>
Reviewed-by: Igor Sheludko <ishell@chromium.org>
Reviewed-by: Clemens Backes <clemensb@chromium.org>
Cr-Commit-Position: refs/heads/main@{#81947}

src/base/build_config.h

@@ -34,6 +34,12 @@
 #define V8_HAS_PTHREAD_JIT_WRITE_PROTECT 0
 #endif
 
+#if defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
+#define V8_HAS_PKU_JIT_WRITE_PROTECT 1
+#else
+#define V8_HAS_PKU_JIT_WRITE_PROTECT 0
+#endif
+
 #if defined(V8_TARGET_ARCH_IA32) || defined(V8_TARGET_ARCH_X64)
 #define V8_TARGET_ARCH_STORES_RETURN_ADDRESS_ON_STACK true
 #else

src/base/platform/memory-protection-key.cc

@@ -4,7 +4,7 @@
 
 #include "src/base/platform/memory-protection-key.h"
 
-#if defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
 #include <sys/mman.h>     // For {mprotect()} protection macros.
 #include <sys/utsname.h>  // For {uname()}.
 #undef MAP_TYPE  // Conflicts with MAP_TYPE in Torque-generated instance-types.h
@@ -34,7 +34,7 @@
 // TODO(dlehmann): Move this import and freestanding functions below to
 // base/platform/platform.h {OS} (lower-level functions) and
 // {base::PageAllocator} (exported API).
-#if defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
 #include <dlfcn.h>
 #endif
 
@@ -55,11 +55,11 @@ pkey_get_t pkey_get = nullptr;
 pkey_set_t pkey_set = nullptr;
 
 #ifdef DEBUG
-bool pkey_initialized = false;
+bool pkey_api_initialized = false;
 #endif
 
 int GetProtectionFromMemoryPermission(PageAllocator::Permission permission) {
-#if defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
   // Mappings for PKU are either RWX (for code), no access (for uncommitted
   // memory), or RW (for assembler buffers).
   switch (permission) {
@@ -80,9 +80,9 @@ int GetProtectionFromMemoryPermission(PageAllocator::Permission permission) {
 }  // namespace
 
 void MemoryProtectionKey::InitializeMemoryProtectionKeySupport() {
-  // Flip {pkey_initialized} (in debug mode) and check the new value.
-  DCHECK_EQ(true, pkey_initialized = !pkey_initialized);
-#if defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
+  // Flip {pkey_api_initialized} (in debug mode) and check the new value.
+  DCHECK_EQ(true, pkey_api_initialized = !pkey_api_initialized);
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
   // PKU was broken on Linux kernels before 5.13 (see
   // https://lore.kernel.org/all/20210623121456.399107624@linutronix.de/).
   // A fix is also included in the 5.4.182 and 5.10.103 versions ("x86/fpu:
@@ -136,7 +136,7 @@ void MemoryProtectionKey::InitializeMemoryProtectionKeySupport() {
 // static
 DISABLE_CFI_ICALL
 int MemoryProtectionKey::AllocateKey() {
-  DCHECK(pkey_initialized);
+  DCHECK(pkey_api_initialized);
   if (!pkey_alloc) return kNoMemoryProtectionKey;
 
   // If there is support in glibc, try to allocate a new key.
@@ -153,7 +153,7 @@ int MemoryProtectionKey::AllocateKey() {
 // static
 DISABLE_CFI_ICALL
 void MemoryProtectionKey::FreeKey(int key) {
-  DCHECK(pkey_initialized);
+  DCHECK(pkey_api_initialized);
   // Only free the key if one was allocated.
   if (key == kNoMemoryProtectionKey) return;
 
@@ -168,7 +168,7 @@ DISABLE_CFI_ICALL
 bool MemoryProtectionKey::SetPermissionsAndKey(
     v8::PageAllocator* page_allocator, base::AddressRegion region,
     v8::PageAllocator::Permission page_permissions, int key) {
-  DCHECK(pkey_initialized);
+  DCHECK(pkey_api_initialized);
 
   void* address = reinterpret_cast<void*>(region.begin());
   size_t size = region.size();
@@ -207,7 +207,7 @@ bool MemoryProtectionKey::SetPermissionsAndKey(
 DISABLE_CFI_ICALL
 void MemoryProtectionKey::SetPermissionsForKey(int key,
                                                Permission permissions) {
-  DCHECK(pkey_initialized);
+  DCHECK(pkey_api_initialized);
   DCHECK_NE(kNoMemoryProtectionKey, key);
 
   // If a valid key was allocated, {pkey_set()} must also be available.
@@ -219,7 +219,7 @@ void MemoryProtectionKey::SetPermissionsForKey(int key,
 // static
 DISABLE_CFI_ICALL
 MemoryProtectionKey::Permission MemoryProtectionKey::GetKeyPermission(int key) {
-  DCHECK(pkey_initialized);
+  DCHECK(pkey_api_initialized);
   DCHECK_NE(kNoMemoryProtectionKey, key);
 
   // If a valid key was allocated, {pkey_get()} must also be available.

src/base/platform/memory-protection-key.h

@@ -5,7 +5,7 @@
 #ifndef V8_BASE_PLATFORM_MEMORY_PROTECTION_KEY_H_
 #define V8_BASE_PLATFORM_MEMORY_PROTECTION_KEY_H_
 
-#if defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
 #include <sys/mman.h>  // For static_assert of permission values.
 #undef MAP_TYPE  // Conflicts with MAP_TYPE in Torque-generated instance-types.h
 #endif

src/common/code-memory-access-inl.h

@@ -7,6 +7,9 @@
 
 #include "src/common/code-memory-access.h"
 #include "src/flags/flags.h"
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
+#include "src/base/platform/memory-protection-key.h"
+#endif
 
 namespace v8 {
 namespace internal {
@@ -31,9 +34,7 @@ RwxMemoryWriteScope::~RwxMemoryWriteScope() {
 #pragma clang diagnostic ignored "-Wunguarded-availability-new"
 
 // static
-bool RwxMemoryWriteScope::IsAllowed() {
-  return pthread_jit_write_protect_supported_np();
-}
+bool RwxMemoryWriteScope::IsSupported() { return true; }
 
 // static
 void RwxMemoryWriteScope::SetWritable() {
@@ -52,10 +53,46 @@ void RwxMemoryWriteScope::SetExecutable() {
 }
 #pragma clang diagnostic pop
 
-#else  // !V8_HAS_PTHREAD_JIT_WRITE_PROTECT
+#elif V8_HAS_PKU_JIT_WRITE_PROTECT
+// static
+bool RwxMemoryWriteScope::IsSupported() {
+  static_assert(base::MemoryProtectionKey::kNoMemoryProtectionKey == -1);
+  DCHECK(pkey_initialized);
+  return memory_protection_key_ >= 0;
+}
+
+// static
+void RwxMemoryWriteScope::SetWritable() {
+  DCHECK(pkey_initialized);
+  if (!IsSupported()) return;
+  if (code_space_write_nesting_level_ == 0) {
+    DCHECK_NE(
+        base::MemoryProtectionKey::GetKeyPermission(memory_protection_key_),
+        base::MemoryProtectionKey::kNoRestrictions);
+    base::MemoryProtectionKey::SetPermissionsForKey(
+        memory_protection_key_, base::MemoryProtectionKey::kNoRestrictions);
+  }
+  code_space_write_nesting_level_++;
+}
+
+// static
+void RwxMemoryWriteScope::SetExecutable() {
+  DCHECK(pkey_initialized);
+  if (!IsSupported()) return;
+  code_space_write_nesting_level_--;
+  if (code_space_write_nesting_level_ == 0) {
+    DCHECK_EQ(
+        base::MemoryProtectionKey::GetKeyPermission(memory_protection_key_),
+        base::MemoryProtectionKey::kNoRestrictions);
+    base::MemoryProtectionKey::SetPermissionsForKey(
+        memory_protection_key_, base::MemoryProtectionKey::kDisableWrite);
+  }
+}
+
+#else  // !V8_HAS_PTHREAD_JIT_WRITE_PROTECT && !V8_TRY_USE_PKU_JIT_WRITE_PROTECT
 
 // static
-bool RwxMemoryWriteScope::IsAllowed() { return true; }
+bool RwxMemoryWriteScope::IsSupported() { return false; }
 
 // static
 void RwxMemoryWriteScope::SetWritable() {}

src/common/code-memory-access.cc

@@ -7,19 +7,58 @@
 namespace v8 {
 namespace internal {
 
-#if V8_HAS_PTHREAD_JIT_WRITE_PROTECT
-
+#if V8_HAS_PTHREAD_JIT_WRITE_PROTECT || V8_HAS_PKU_JIT_WRITE_PROTECT
 thread_local int RwxMemoryWriteScope::code_space_write_nesting_level_ = 0;
+#endif  // V8_HAS_PTHREAD_JIT_WRITE_PROTECT || V8_HAS_PKU_JIT_WRITE_PROTECT
+
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
+int RwxMemoryWriteScope::memory_protection_key_ =
+    base::MemoryProtectionKey::kNoMemoryProtectionKey;
+
+#if DEBUG
+bool RwxMemoryWriteScope::pkey_initialized = false;
+#endif
+
+void RwxMemoryWriteScope::InitializeMemoryProtectionKey() {
+  // Flip {pkey_initialized} (in debug mode) and check the new value.
+  DCHECK_EQ(true, pkey_initialized = !pkey_initialized);
+  memory_protection_key_ = base::MemoryProtectionKey::AllocateKey();
+  DCHECK(memory_protection_key_ > 0 ||
+         memory_protection_key_ ==
+             base::MemoryProtectionKey::kNoMemoryProtectionKey);
+}
+
+bool RwxMemoryWriteScope::IsPKUWritable() {
+  DCHECK(pkey_initialized);
+  return base::MemoryProtectionKey::GetKeyPermission(memory_protection_key_) ==
+         base::MemoryProtectionKey::kNoRestrictions;
+}
 
-RwxMemoryWriteScopeForTesting::RwxMemoryWriteScopeForTesting() {
-  RwxMemoryWriteScope::SetWritable();
+ResetPKUPermissionsForThreadSpawning::ResetPKUPermissionsForThreadSpawning() {
+  if (!RwxMemoryWriteScope::IsSupported()) return;
+  int pkey = RwxMemoryWriteScope::memory_protection_key();
+  was_writable_ = base::MemoryProtectionKey::GetKeyPermission(pkey) ==
+                  base::MemoryProtectionKey::kNoRestrictions;
+  if (was_writable_) {
+    base::MemoryProtectionKey::SetPermissionsForKey(
+        pkey, base::MemoryProtectionKey::kDisableWrite);
+  }
 }
 
-RwxMemoryWriteScopeForTesting::~RwxMemoryWriteScopeForTesting() {
-  RwxMemoryWriteScope::SetExecutable();
+ResetPKUPermissionsForThreadSpawning::~ResetPKUPermissionsForThreadSpawning() {
+  if (!RwxMemoryWriteScope::IsSupported()) return;
+  int pkey = RwxMemoryWriteScope::memory_protection_key();
+  if (was_writable_) {
+    base::MemoryProtectionKey::SetPermissionsForKey(
+        pkey, base::MemoryProtectionKey::kNoRestrictions);
+  }
 }
+#endif  // V8_HAS_PKU_JIT_WRITE_PROTECT
+
+RwxMemoryWriteScopeForTesting::RwxMemoryWriteScopeForTesting()
+    : RwxMemoryWriteScope("For Testing") {}
 
-#endif  // V8_HAS_PTHREAD_JIT_WRITE_PROTECT
+RwxMemoryWriteScopeForTesting::~RwxMemoryWriteScopeForTesting() {}
 
 }  // namespace internal
 }  // namespace v8

src/common/code-memory-access.h

@@ -20,19 +20,28 @@ class CodeSpaceWriteScope;
 }
 
 // This scope is a wrapper for APRR/MAP_JIT machinery on MacOS on ARM64
-// ("Apple M1"/Apple Silicon) with respective semantics.
-// See pthread_jit_write_protect_np() for details.
-// The scope must not be used if the process does not have the
-// "com.apple.security.cs.allow-jit" entitlement (see IsAllowed()).
-// On other platforms the scope is a no-op and thus it's allowed to be used.
+// ("Apple M1"/Apple Silicon) or Intel PKU (aka. memory protection keys)
+// with respective low-level semantics.
 //
-// The semantics is the following: the scope switches permissions between
-// writable and executable for all the pages allocated with RWX permissions.
-// Only current thread is affected. This achieves "real" W^X and it's fast.
+// The semantics on MacOS on ARM64 is the following:
+// The scope switches permissions between writable and executable for all the
+// pages allocated with RWX permissions. Only current thread is affected.
+// This achieves "real" W^X and it's fast (see pthread_jit_write_protect_np()
+// for details).
 // By default it is assumed that the state is executable.
+// It's also assumed that the process has the "com.apple.security.cs.allow-jit"
+// entitlement.
+//
+// The semantics on Intel with PKU support is the following:
+// When Intel PKU is available, the scope switches the protection key's
+// permission between writable and not writable. The executable permission
+// cannot be retracted with PKU. That is, this "only" achieves write
+// protection, but is similarly thread-local and fast.
+//
+// On other platforms the scope is a no-op and thus it's allowed to be used.
 //
 // The scope is reentrant and thread safe.
-class V8_NODISCARD RwxMemoryWriteScope final {
+class V8_NODISCARD RwxMemoryWriteScope {
  public:
   // The comment argument is used only for ensuring that explanation about why
   // the scope is needed is given at particular use case.
@@ -44,11 +53,17 @@ class V8_NODISCARD RwxMemoryWriteScope final {
   RwxMemoryWriteScope(const RwxMemoryWriteScope&) = delete;
   RwxMemoryWriteScope& operator=(const RwxMemoryWriteScope&) = delete;
 
-  // Returns true if the configuration of the binary allows using of MAP_JIT
-  // machinery.
-  // This method is intended to be used for checking that the state of --jitless
-  // flag does not contradict the allowance of the MAP_JIT feature.
-  V8_INLINE static bool IsAllowed();
+  // Returns true if current configuration supports fast write-protection of
+  // executable pages.
+  V8_INLINE static bool IsSupported();
+
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
+  static int memory_protection_key() { return memory_protection_key_; }
+
+  static void InitializeMemoryProtectionKey();
+
+  static bool IsPKUWritable();
+#endif  // V8_HAS_PKU_JIT_WRITE_PROTECT
 
  private:
   friend class CodePageCollectionMemoryModificationScope;
@@ -63,10 +78,17 @@ class V8_NODISCARD RwxMemoryWriteScope final {
   V8_INLINE static void SetWritable();
   V8_INLINE static void SetExecutable();
 
-#if V8_HAS_PTHREAD_JIT_WRITE_PROTECT
+#if V8_HAS_PTHREAD_JIT_WRITE_PROTECT || V8_HAS_PKU_JIT_WRITE_PROTECT
   // This counter is used for supporting scope reentrance.
   static thread_local int code_space_write_nesting_level_;
-#endif  // V8_HAS_PTHREAD_JIT_WRITE_PROTECT
+#endif  // V8_HAS_PTHREAD_JIT_WRITE_PROTECT || V8_HAS_PKU_JIT_WRITE_PROTECT
+
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
+  static int memory_protection_key_;
+#if DEBUG
+  static bool pkey_initialized;
+#endif
+#endif  // V8_HAS_PKU_JIT_WRITE_PROTECT
 };
 
 // This class is a no-op version of the RwxMemoryWriteScope class above.
@@ -79,19 +101,35 @@ class V8_NODISCARD NopRwxMemoryWriteScope final {
   }
 };
 
+// Sometimes we need to call a function which will / might spawn a new thread,
+// like {JobHandle::NotifyConcurrencyIncrease}, while holding a
+// {RwxMemoryWriteScope}. This is problematic since the new thread will inherit
+// the parent thread's PKU permissions.
+// The {ResetPKUPermissionsForThreadSpawning} scope will thus reset the PKU
+// permissions as long as it is in scope, such that it is safe to spawn new
+// threads.
+class V8_NODISCARD ResetPKUPermissionsForThreadSpawning {
+ public:
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
+  ResetPKUPermissionsForThreadSpawning();
+  ~ResetPKUPermissionsForThreadSpawning();
+
+ private:
+  bool was_writable_;
+#else
+  // Define an empty constructor to avoid "unused variable" warnings.
+  ResetPKUPermissionsForThreadSpawning() {}
+#endif
+};
+
 // Same as the RwxMemoryWriteScope but without inlining the code.
 // This is a workaround for component build issue (crbug/1316800), when
 // a thread_local value can't be properly exported.
-class V8_NODISCARD RwxMemoryWriteScopeForTesting final {
+class V8_NODISCARD RwxMemoryWriteScopeForTesting final
+    : public RwxMemoryWriteScope {
  public:
-#if V8_HAS_PTHREAD_JIT_WRITE_PROTECT
   V8_EXPORT_PRIVATE RwxMemoryWriteScopeForTesting();
   V8_EXPORT_PRIVATE ~RwxMemoryWriteScopeForTesting();
-#else
-  V8_INLINE RwxMemoryWriteScopeForTesting() {
-    // Define a constructor to avoid unused variable warnings.
-  }
-#endif  // V8_HAS_PTHREAD_JIT_WRITE_PROTECT
 
   // Disable copy constructor and copy-assignment operator, since this manages
   // a resource and implicit copying of the scope can yield surprising errors.

src/heap/heap.cc

@@ -5675,11 +5675,6 @@ void Heap::SetUp(LocalHeap* main_thread_local_heap) {
       reinterpret_cast<uintptr_t>(v8::internal::GetRandomMmapAddr()) &
       ~kMmapRegionMask;
 
-  // Ensure that RwxMemoryWriteScope and other dependent scopes (in particular,
-  // CodePage*ModificationScope and CodeSpaceMemoryModificationScope)
-  // are allowed to be used when jitless mode is not enabled.
-  CHECK_IMPLIES(!FLAG_jitless, RwxMemoryWriteScope::IsAllowed());
-
   v8::PageAllocator* code_page_allocator;
   if (isolate_->RequiresCodeRange() || code_range_size_ != 0) {
     const size_t requested_size =

src/init/v8.cc

@@ -13,6 +13,7 @@
 #include "src/base/platform/platform.h"
 #include "src/codegen/cpu-features.h"
 #include "src/codegen/interface-descriptors.h"
+#include "src/common/code-memory-access.h"
 #include "src/debug/debug.h"
 #include "src/deoptimizer/deoptimizer.h"
 #include "src/execution/frames.h"
@@ -251,6 +252,12 @@ void V8::Initialize() {
   ElementsAccessor::InitializeOncePerProcess();
   Bootstrapper::InitializeOncePerProcess();
   CallDescriptors::InitializeOncePerProcess();
+
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
+  base::MemoryProtectionKey::InitializeMemoryProtectionKeySupport();
+  RwxMemoryWriteScope::InitializeMemoryProtectionKey();
+#endif
+
 #if V8_ENABLE_WEBASSEMBLY
   wasm::WasmEngine::InitializeOncePerProcess();
 #endif  // V8_ENABLE_WEBASSEMBLY

src/wasm/code-space-access.cc

@@ -71,9 +71,8 @@ bool CodeSpaceWriteScope::SwitchingPerNativeModule() { return false; }
 
 // static
 void CodeSpaceWriteScope::SetWritable() {
-  auto* code_manager = GetWasmCodeManager();
-  if (code_manager->MemoryProtectionKeysEnabled()) {
-    code_manager->SetThreadWritable(true);
+  if (WasmCodeManager::MemoryProtectionKeysEnabled()) {
+    RwxMemoryWriteScope::SetWritable();
   } else if (FLAG_wasm_write_protect_code_memory) {
     current_native_module_->AddWriter();
   }
@@ -81,10 +80,9 @@ void CodeSpaceWriteScope::SetWritable() {
 
 // static
 void CodeSpaceWriteScope::SetExecutable() {
-  auto* code_manager = GetWasmCodeManager();
-  if (code_manager->MemoryProtectionKeysEnabled()) {
+  if (WasmCodeManager::MemoryProtectionKeysEnabled()) {
     DCHECK(FLAG_wasm_memory_protection_keys);
-    code_manager->SetThreadWritable(false);
+    RwxMemoryWriteScope::SetExecutable();
   } else if (FLAG_wasm_write_protect_code_memory) {
     current_native_module_->RemoveWriter();
   }
@@ -92,21 +90,10 @@ void CodeSpaceWriteScope::SetExecutable() {
 
 // static
 bool CodeSpaceWriteScope::SwitchingPerNativeModule() {
-  return !GetWasmCodeManager()->MemoryProtectionKeysEnabled() &&
+  return !WasmCodeManager::MemoryProtectionKeysEnabled() &&
          FLAG_wasm_write_protect_code_memory;
 }
 
-ResetPKUPermissionsForThreadSpawning::ResetPKUPermissionsForThreadSpawning() {
-  auto* code_manager = GetWasmCodeManager();
-  was_writable_ = code_manager->MemoryProtectionKeysEnabled() &&
-                  code_manager->MemoryProtectionKeyWritable();
-  if (was_writable_) code_manager->SetThreadWritable(false);
-}
-
-ResetPKUPermissionsForThreadSpawning::~ResetPKUPermissionsForThreadSpawning() {
-  if (was_writable_) GetWasmCodeManager()->SetThreadWritable(true);
-}
-
 #endif  // !V8_HAS_PTHREAD_JIT_WRITE_PROTECT
 
 }  // namespace wasm

src/wasm/code-space-access.h

@@ -72,27 +72,6 @@ class V8_NODISCARD CodeSpaceWriteScope final {
   NativeModule* const previous_native_module_;
 };
 
-// Sometimes we need to call a function which will / might spawn a new thread,
-// like {JobHandle::NotifyConcurrencyIncrease}, while holding a
-// {CodeSpaceWriteScope}. This is problematic since the new thread will inherit
-// the parent thread's PKU permissions.
-// The {ResetPKUPermissionsForThreadSpawning} scope will thus reset the PKU
-// permissions as long as it is in scope, such that it is safe to spawn new
-// threads.
-class V8_NODISCARD ResetPKUPermissionsForThreadSpawning {
- public:
-#if !V8_HAS_PTHREAD_JIT_WRITE_PROTECT
-  ResetPKUPermissionsForThreadSpawning();
-  ~ResetPKUPermissionsForThreadSpawning();
-
- private:
-  bool was_writable_;
-#else
-  // Define an empty constructor to avoid "unused variable" warnings.
-  ResetPKUPermissionsForThreadSpawning() {}
-#endif
-};
-
 }  // namespace v8::internal::wasm
 
 #endif  // V8_WASM_CODE_SPACE_ACCESS_H_

src/wasm/module-compiler.cc

@@ -1577,7 +1577,7 @@ CompilationExecutionResult ExecuteCompilationUnits(
   // Also, open a CodeSpaceWriteScope now to have (thread-local) write access to
   // the assembler buffers.
   base::Optional<CodeSpaceWriteScope> write_scope_for_assembler_buffers;
-  if (GetWasmCodeManager()->MemoryProtectionKeysEnabled()) {
+  if (WasmCodeManager::MemoryProtectionKeysEnabled()) {
     optional_assembler_buffer_cache.emplace();
     assembler_buffer_cache = &*optional_assembler_buffer_cache;
     write_scope_for_assembler_buffers.emplace(nullptr);

src/wasm/wasm-code-manager.cc

@@ -12,7 +12,6 @@
 #include "src/base/build_config.h"
 #include "src/base/iterator.h"
 #include "src/base/macros.h"
-#include "src/base/platform/memory-protection-key.h"
 #include "src/base/platform/platform.h"
 #include "src/base/small-vector.h"
 #include "src/base/string-format.h"
@@ -20,6 +19,7 @@
 #include "src/codegen/assembler-inl.h"
 #include "src/codegen/macro-assembler-inl.h"
 #include "src/codegen/macro-assembler.h"
+#include "src/common/code-memory-access.h"
 #include "src/common/globals.h"
 #include "src/diagnostics/disassembler.h"
 #include "src/logging/counters.h"
@@ -517,10 +517,9 @@ int WasmCode::GetSourcePositionBefore(int offset) {
 constexpr size_t WasmCodeAllocator::kMaxCodeSpaceSize;
 
 WasmCodeAllocator::WasmCodeAllocator(std::shared_ptr<Counters> async_counters)
-    : protect_code_memory_(
-          !V8_HAS_PTHREAD_JIT_WRITE_PROTECT &&
-          FLAG_wasm_write_protect_code_memory &&
-          !GetWasmCodeManager()->MemoryProtectionKeysEnabled()),
+    : protect_code_memory_(!V8_HAS_PTHREAD_JIT_WRITE_PROTECT &&
+                           FLAG_wasm_write_protect_code_memory &&
+                           !WasmCodeManager::MemoryProtectionKeysEnabled()),
       async_counters_(std::move(async_counters)) {
   owned_code_space_.reserve(4);
 }
@@ -1883,18 +1882,11 @@ NativeModule::~NativeModule() {
 
 WasmCodeManager::WasmCodeManager()
     : max_committed_code_space_(FLAG_wasm_max_code_space * MB),
-      critical_committed_code_space_(max_committed_code_space_ / 2),
-      memory_protection_key_(base::MemoryProtectionKey::AllocateKey()) {
-  // Ensure that RwxMemoryWriteScope and other dependent scopes (in particular,
-  // wasm::CodeSpaceWriteScope) are allowed to be used.
-  CHECK(RwxMemoryWriteScope::IsAllowed());
-}
+      critical_committed_code_space_(max_committed_code_space_ / 2) {}
 
 WasmCodeManager::~WasmCodeManager() {
   // No more committed code space.
   DCHECK_EQ(0, total_committed_code_space_.load());
-
-  base::MemoryProtectionKey::FreeKey(memory_protection_key_);
 }
 
 #if defined(V8_OS_WIN64)
@@ -1944,14 +1936,19 @@ void WasmCodeManager::Commit(base::AddressRegion region) {
   // open when calling this method.
   PageAllocator::Permission permission = PageAllocator::kReadWriteExecute;
 
-  bool success;
+  bool success = false;
   if (MemoryProtectionKeysEnabled()) {
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
     TRACE_HEAP(
-        "Setting rwx permissions and memory protection key %d for 0x%" PRIxPTR
+        "Setting rwx permissions and memory protection key for 0x%" PRIxPTR
         ":0x%" PRIxPTR "\n",
-        memory_protection_key_, region.begin(), region.end());
+        region.begin(), region.end());
     success = base::MemoryProtectionKey::SetPermissionsAndKey(
-        GetPlatformPageAllocator(), region, permission, memory_protection_key_);
+        GetPlatformPageAllocator(), region, permission,
+        RwxMemoryWriteScope::memory_protection_key());
+#else
+    UNREACHABLE();
+#endif  // V8_HAS_PKU_JIT_WRITE_PROTECT
   } else {
     TRACE_HEAP("Setting rwx permissions for 0x%" PRIxPTR ":0x%" PRIxPTR "\n",
                region.begin(), region.end());
@@ -2148,52 +2145,37 @@ size_t WasmCodeManager::EstimateNativeModuleMetaDataSize(
          far_jump_table_size;
 }
 
-void WasmCodeManager::SetThreadWritable(bool writable) {
-  DCHECK(MemoryProtectionKeysEnabled());
-
-  auto permissions = writable ? base::MemoryProtectionKey::kNoRestrictions
-                              : base::MemoryProtectionKey::kDisableWrite;
-
-  // When switching to writable we should not already be writable. Otherwise
-  // this points at a problem with counting writers, or with wrong
-  // initialization (globally or per thread).
-  DCHECK_IMPLIES(writable, !MemoryProtectionKeyWritable());
-
-  TRACE_HEAP("Setting memory protection key %d to writable: %d.\n",
-             memory_protection_key_, writable);
-  base::MemoryProtectionKey::SetPermissionsForKey(memory_protection_key_,
-                                                  permissions);
-}
-
-bool WasmCodeManager::HasMemoryProtectionKeySupport() const {
-  return memory_protection_key_ !=
-         base::MemoryProtectionKey::kNoMemoryProtectionKey;
+// static
+bool WasmCodeManager::HasMemoryProtectionKeySupport() {
+  return RwxMemoryWriteScope::IsSupported();
 }
 
-bool WasmCodeManager::MemoryProtectionKeysEnabled() const {
+// static
+bool WasmCodeManager::MemoryProtectionKeysEnabled() {
   return HasMemoryProtectionKeySupport() && FLAG_wasm_memory_protection_keys;
 }
 
-bool WasmCodeManager::MemoryProtectionKeyWritable() const {
-  return base::MemoryProtectionKey::GetKeyPermission(memory_protection_key_) ==
-         base::MemoryProtectionKey::kNoRestrictions;
+// static
+bool WasmCodeManager::MemoryProtectionKeyWritable() {
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
+  return RwxMemoryWriteScope::IsPKUWritable();
+#else
+  return false;
+#endif  // V8_HAS_PKU_JIT_WRITE_PROTECT
 }
 
-void WasmCodeManager::InitializeMemoryProtectionKeyPermissionsIfSupported()
-    const {
+// static
+void WasmCodeManager::InitializeMemoryProtectionKeyPermissionsIfSupported() {
   if (!HasMemoryProtectionKeySupport()) return;
   // The default permission is {kDisableAccess}. Switch from that to
   // {kDisableWrite}. Leave other permissions untouched, as the thread did
   // already use the memory protection key in that case.
-  if (base::MemoryProtectionKey::GetKeyPermission(memory_protection_key_) ==
-      base::MemoryProtectionKey::kDisableAccess) {
-    base::MemoryProtectionKey::SetPermissionsForKey(
-        memory_protection_key_, base::MemoryProtectionKey::kDisableWrite);
-  }
+  RwxMemoryWriteScope initialize_permission_scope(
+      "For initialization if PKU is in kNoAccess permission case.");
 }
 
 base::AddressRegion WasmCodeManager::AllocateAssemblerBufferSpace(int size) {
-#if defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
   if (MemoryProtectionKeysEnabled()) {
     auto* page_allocator = GetPlatformPageAllocator();
     size_t page_size = page_allocator->AllocatePageSize();
@@ -2213,23 +2195,23 @@ base::AddressRegion WasmCodeManager::AllocateAssemblerBufferSpace(int size) {
         base::AddressRegionOf(reinterpret_cast<uint8_t*>(mapped), size);
     CHECK(base::MemoryProtectionKey::SetPermissionsAndKey(
         page_allocator, region, PageAllocator::kReadWrite,
-        memory_protection_key_));
+        RwxMemoryWriteScope::memory_protection_key()));
     return region;
   }
-#endif  // defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
+#endif  // V8_HAS_PKU_JIT_WRITE_PROTECT
   DCHECK(!MemoryProtectionKeysEnabled());
   return base::AddressRegionOf(new uint8_t[size], size);
 }
 
 void WasmCodeManager::FreeAssemblerBufferSpace(base::AddressRegion region) {
-#if defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
+#if V8_HAS_PKU_JIT_WRITE_PROTECT
   if (MemoryProtectionKeysEnabled()) {
     auto* page_allocator = GetPlatformPageAllocator();
     FreePages(page_allocator, reinterpret_cast<void*>(region.begin()),
               region.size());
     return;
   }
-#endif  // defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
+#endif  // V8_HAS_PKU_JIT_WRITE_PROTECT
   DCHECK(!MemoryProtectionKeysEnabled());
   delete[] reinterpret_cast<uint8_t*>(region.begin());
 }

src/wasm/wasm-code-manager.h

@@ -19,9 +19,9 @@
 #include "src/base/address-region.h"
 #include "src/base/bit-field.h"
 #include "src/base/macros.h"
-#include "src/base/platform/memory-protection-key.h"
 #include "src/base/vector.h"
 #include "src/builtins/builtins.h"
+#include "src/common/code-memory-access.h"
 #include "src/handles/handles.h"
 #include "src/tasks/operations-barrier.h"
 #include "src/trap-handler/trap-handler.h"
@@ -1063,29 +1063,22 @@ class V8_EXPORT_PRIVATE WasmCodeManager final {
   // generated code. This data still be stored on the C++ heap.
   static size_t EstimateNativeModuleMetaDataSize(const WasmModule* module);
 
-  // Set this thread's permission of all owned code space to read-write or
-  // read-only (if {writable} is false). Can only be called if
-  // {HasMemoryProtectionKeySupport()} is {true}.
-  // Since the permission is thread-local, there is no requirement to hold any
-  // lock when calling this method.
-  void SetThreadWritable(bool writable);
-
   // Returns true if there is hardware support for PKU. Use
   // {MemoryProtectionKeysEnabled} to also check if PKU usage is enabled via
   // flags.
-  bool HasMemoryProtectionKeySupport() const;
+  static bool HasMemoryProtectionKeySupport();
 
   // Returns true if PKU should be used.
-  bool MemoryProtectionKeysEnabled() const;
+  static bool MemoryProtectionKeysEnabled();
 
   // Returns {true} if the memory protection key is write-enabled for the
   // current thread.
   // Can only be called if {HasMemoryProtectionKeySupport()} is {true}.
-  bool MemoryProtectionKeyWritable() const;
+  static bool MemoryProtectionKeyWritable();
 
   // Initialize the current thread's permissions for the memory protection key,
   // if we have support.
-  void InitializeMemoryProtectionKeyPermissionsIfSupported() const;
+  static void InitializeMemoryProtectionKeyPermissionsIfSupported();
 
   // Allocate new memory for assembler buffers, potentially protected by PKU.
   base::AddressRegion AllocateAssemblerBufferSpace(int size);
@@ -1120,8 +1113,6 @@ class V8_EXPORT_PRIVATE WasmCodeManager final {
   // and updated after each GC.
   std::atomic<size_t> critical_committed_code_space_;
 
-  int memory_protection_key_;
-
   mutable base::Mutex native_modules_mutex_;
 
   //////////////////////////////////////////////////////////////////////////////

src/wasm/wasm-engine.cc

@@ -5,7 +5,6 @@
 #include "src/wasm/wasm-engine.h"
 
 #include "src/base/functional.h"
-#include "src/base/platform/memory-protection-key.h"
 #include "src/base/platform/time.h"
 #include "src/common/globals.h"
 #include "src/debug/debug.h"
@@ -1021,7 +1020,7 @@ void WasmEngine::AddIsolate(Isolate* isolate) {
   // In that case, the current thread might still have the default permissions
   // for the memory protection key (== no access). Thus initialize the
   // permissions now.
-  GetWasmCodeManager()->InitializeMemoryProtectionKeyPermissionsIfSupported();
+  WasmCodeManager::InitializeMemoryProtectionKeyPermissionsIfSupported();
 
   // Install sampling GC callback.
   // TODO(v8:7424): For now we sample module sizes in a GC callback. This will
@@ -1182,7 +1181,7 @@ std::shared_ptr<NativeModule> WasmEngine::NewNativeModule(
     isolate_info->pku_support_sampled = true;
     auto* histogram =
         isolate->counters()->wasm_memory_protection_keys_support();
-    bool has_mpk = GetWasmCodeManager()->HasMemoryProtectionKeySupport();
+    bool has_mpk = WasmCodeManager::HasMemoryProtectionKeySupport();
     histogram->AddSample(has_mpk ? 1 : 0);
   }
 
@@ -1620,7 +1619,6 @@ GlobalWasmState* global_wasm_state = nullptr;
 
 // static
 void WasmEngine::InitializeOncePerProcess() {
-  base::MemoryProtectionKey::InitializeMemoryProtectionKeySupport();
   DCHECK_NULL(global_wasm_state);
   global_wasm_state = new GlobalWasmState();
 }

test/unittests/wasm/memory-protection-unittest.cc

@@ -52,8 +52,8 @@ class MemoryProtectionTest : public TestWithNativeContext {
     bool enable_pku = mode == kPku || mode == kPkuWithMprotectFallback;
     FLAG_wasm_memory_protection_keys = enable_pku;
     // The key is initially write-protected.
-    CHECK_IMPLIES(GetWasmCodeManager()->HasMemoryProtectionKeySupport(),
-                  !GetWasmCodeManager()->MemoryProtectionKeyWritable());
+    CHECK_IMPLIES(WasmCodeManager::HasMemoryProtectionKeySupport(),
+                  !WasmCodeManager::MemoryProtectionKeyWritable());
 
     bool enable_mprotect =
         mode == kMprotect || mode == kPkuWithMprotectFallback;
@@ -108,8 +108,7 @@ class MemoryProtectionTest : public TestWithNativeContext {
     // M1 always uses MAP_JIT.
     if (V8_HAS_PTHREAD_JIT_WRITE_PROTECT) return false;
     bool param_has_pku = mode_ == kPku || mode_ == kPkuWithMprotectFallback;
-    return param_has_pku &&
-           GetWasmCodeManager()->HasMemoryProtectionKeySupport();
+    return param_has_pku && WasmCodeManager::HasMemoryProtectionKeySupport();
   }
 
  private: