[wasm] Use PKUs for code space write protection

This is the second CL in a line of two to implement PKU-based
WebAssembly code space write protection. The first CL added two
low-level PKU functions; this CL uses them to grant/withdraw writable
permissions, local to each thread that wants to modify the code space.

In particular, when {--wasm-memory-protection-keys} is enabled, we first
associate a memory protection key with all code pages, which by
default does not allow any write access. Then, before each location that
needs to modify the code space, we open
{NativeModuleModificationScope}s (which are already present for
mprotect-based write protection). When the PKU flag is given, this then
first tries to set permissions of a memory protection key (which is
fast), and otherwise when {--wasm-write-protect-code-memory} is enabled,
falls back to mprotect-based write protection (which is much more
expensive and also not thread-local, but for the whole process).

R=clemensb@chromium.org

Bug: v8:11714
Change-Id: I3527906a8d9f776ed44c8d5db52539e78e1c52fd
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2882800
Commit-Queue: Daniel Lehmann <dlehmann@google.com>
Reviewed-by: Clemens Backes <clemensb@chromium.org>
Cr-Commit-Position: refs/heads/master@{#74501}

src/wasm/wasm-code-manager.cc

@@ -753,6 +753,26 @@ bool WasmCodeAllocator::SetWritable(bool writable) {
   return true;
 }
 
+bool WasmCodeAllocator::SetThreadWritable(bool writable) {
+  static thread_local int writable_nesting_level = 0;
+  if (writable) {
+    if (++writable_nesting_level > 1) return true;
+  } else {
+    DCHECK_GT(writable_nesting_level, 0);
+    if (--writable_nesting_level > 0) return true;
+  }
+  writable = writable_nesting_level > 0;
+
+  int key = code_manager_->memory_protection_key_;
+
+  MemoryProtectionKeyPermission permissions =
+      writable ? kNoRestrictions : kDisableWrite;
+
+  TRACE_HEAP("Setting memory protection key %d to writable: %d.\n", key,
+             writable);
+  return SetPermissionsForMemoryProtectionKey(key, permissions);
+}
+
 void WasmCodeAllocator::FreeCode(Vector<WasmCode* const> codes) {
   // Zap code area and collect freed code regions.
   DisjointAllocationPool freed_regions;
@@ -1746,12 +1766,22 @@ void WasmCodeManager::Commit(base::AddressRegion region) {
   // method.
   PageAllocator::Permission permission = PageAllocator::kReadWriteExecute;
 
+  bool success;
+  if (FLAG_wasm_memory_protection_keys) {
+    TRACE_HEAP(
+        "Setting rwx permissions and memory protection key %d for 0x%" PRIxPTR
+        ":0x%" PRIxPTR "\n",
+        memory_protection_key_, region.begin(), region.end());
+    success = SetPermissionsAndMemoryProtectionKey(
+        GetPlatformPageAllocator(), region, permission, memory_protection_key_);
+  } else {
     TRACE_HEAP("Setting rwx permissions for 0x%" PRIxPTR ":0x%" PRIxPTR "\n",
                region.begin(), region.end());
+    success = SetPermissions(GetPlatformPageAllocator(), region.begin(),
+                             region.size(), permission);
+  }
 
-  if (!SetPermissions(GetPlatformPageAllocator(), region.begin(), region.size(),
-                      permission)) {
-    // Highly unlikely.
+  if (V8_UNLIKELY(!success)) {
     V8::FatalProcessOutOfMemory(
         nullptr,
         "WasmCodeManager::Commit: Cannot make pre-reserved region writable");
@@ -1770,8 +1800,13 @@ void WasmCodeManager::Decommit(base::AddressRegion region) {
   USE(old_committed);
   TRACE_HEAP("Discarding system pages 0x%" PRIxPTR ":0x%" PRIxPTR "\n",
              region.begin(), region.end());
-  CHECK(allocator->SetPermissions(reinterpret_cast<void*>(region.begin()),
-                                  region.size(), PageAllocator::kNoAccess));
+  if (FLAG_wasm_memory_protection_keys) {
+    CHECK(SetPermissionsAndMemoryProtectionKey(
+        allocator, region, PageAllocator::kNoAccess, kNoMemoryProtectionKey));
+  } else {
+    CHECK(SetPermissions(allocator, region.begin(), region.size(),
+                         PageAllocator::kNoAccess));
+  }
 }
 
 void WasmCodeManager::AssignRange(base::AddressRegion region,
@@ -2235,14 +2270,29 @@ WasmCode* WasmCodeManager::LookupCode(Address pc) const {
 NativeModuleModificationScope::NativeModuleModificationScope(
     NativeModule* native_module)
     : native_module_(native_module) {
-  if (FLAG_wasm_write_protect_code_memory && native_module_) {
+  DCHECK_NOT_NULL(native_module_);
+  if (FLAG_wasm_memory_protection_keys) {
+    bool success = native_module_->SetThreadWritable(true);
+    if (!success && FLAG_wasm_write_protect_code_memory) {
+      // Fallback to mprotect-based write protection (much slower).
+      success = native_module_->SetWritable(true);
+      CHECK(success);
+    }
+  } else if (FLAG_wasm_write_protect_code_memory) {
     bool success = native_module_->SetWritable(true);
     CHECK(success);
   }
 }
 
 NativeModuleModificationScope::~NativeModuleModificationScope() {
-  if (FLAG_wasm_write_protect_code_memory && native_module_) {
+  if (FLAG_wasm_memory_protection_keys) {
+    bool success = native_module_->SetThreadWritable(false);
+    if (!success && FLAG_wasm_write_protect_code_memory) {
+      // Fallback to mprotect-based write protection (much slower).
+      success = native_module_->SetWritable(false);
+      CHECK(success);
+    }
+  } else if (FLAG_wasm_write_protect_code_memory) {
     bool success = native_module_->SetWritable(false);
     CHECK(success);
   }

src/wasm/wasm-code-manager.h

@@ -426,6 +426,12 @@ class WasmCodeAllocator {
   // Hold the {NativeModule}'s {allocation_mutex_} when calling this method.
   V8_EXPORT_PRIVATE bool SetWritable(bool writable);
 
+  // Set this thread's permission of all owned code space to read-write or
+  // read-only (if {writable} is false). Uses memory protection keys.
+  // Returns true on success. Since the permission is thread-local, there is no
+  // requirement to hold any lock when calling this method.
+  bool SetThreadWritable(bool writable);
+
   // Free memory pages of all given code objects. Used for wasm code GC.
   // Hold the {NativeModule}'s {allocation_mutex_} when calling this method.
   void FreeCode(Vector<WasmCode* const>);
@@ -577,6 +583,10 @@ class V8_EXPORT_PRIVATE NativeModule final {
     return code_allocator_.SetWritable(writable);
   }
 
+  bool SetThreadWritable(bool writable) {
+    return code_allocator_.SetThreadWritable(writable);
+  }
+
   // For cctests, where we build both WasmModule and the runtime objects
   // on the fly, and bypass the instance builder pipeline.
   void ReserveCodeTableForTesting(uint32_t max_functions);