Fix exit handler segfault with union PKRU for destructor wrappers

Library destructors need to call libc functions (e.g., __cxa_finalize) which
access libc's .bss section protected by pkey 1. Previously, destructor wrappers
switched exclusively to the target compartment (e.g., pkey 2), causing SEGV_PKUERR
when the destructor called back into libc.

Solution: Use union PKRU value (0xffffffc0 for pkeys 1+2) that allows simultaneous
access to both the exit compartment (libc, pkey 1) and target compartment (pkey 2)
during destructor execution. This prevents protection key violations when destructors
call libc functions via PLT without requiring wrappers for every libc call.

Changes:
- tools/rewriter/GenCallAsm.cpp: Add optional union_pkey parameter to emit_set_pkru()
- tools/rewriter/GenCallAsm.h: Add use_union_pkru parameter to emit_asm_wrapper()
- tools/rewriter/SourceRewriter.cpp: Pass use_union_pkru=true for destructor wrappers
- runtime/libia2/init.c: Re-enable ia2_setup_destructors() call
- runtime/libia2/CMakeLists.txt: Add trace_exit.c to build
- runtime/libia2/include/ia2_internal.h: Add ia2_trace_exit_record() declaration

Author: oinoom

runtime/libia2/CMakeLists.txt

@@ -1,7 +1,16 @@
 cmake_minimum_required(VERSION 4.0)
 project(libia2)
 
-add_library(libia2 ia2.c init.c threads.c main.c exit.c memory_maps.c thread_name.c)
+add_library(libia2
+    ia2.c
+    init.c
+    destructor_runtime.c
+    threads.c
+    main.c
+    exit.c
+    memory_maps.c
+    thread_name.c
+    trace_exit.c)
 target_compile_options(libia2 PRIVATE "-fPIC")
 
 if (LIBIA2_AARCH64)

runtime/libia2/destructor_runtime.c

@@ -0,0 +1,156 @@
+#ifndef _GNU_SOURCE
+#define _GNU_SOURCE
+#endif
+
+#include "ia2_destructor_runtime.h"
+
+#include "ia2.h"
+#include "ia2_internal.h"
+
+#include <dlfcn.h>
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+struct ia2_destructor_runtime_entry {
+  void (*wrapper)(void);
+  const struct ia2_destructor_metadata_record *record;
+  uint32_t pkru_value;
+};
+
+static struct ia2_destructor_runtime_entry *destructor_entries;
+static size_t destructor_entry_count;
+static bool initialized;
+
+static uint32_t compute_union_pkru(int exit_pkey, int target_pkey) {
+  uint32_t mask = ~0u;
+  mask &= ~0b11u; // Always allow shared key 0.
+  mask &= ~(0b11u << (2 * exit_pkey));
+  mask &= ~(0b11u << (2 * target_pkey));
+  return mask;
+}
+
+static void ensure_initialized(void) {
+  if (initialized) {
+    return;
+  }
+  initialized = true;
+
+  if (ia2_destructor_metadata_count == 0) {
+    return;
+  }
+
+  destructor_entries =
+      calloc(ia2_destructor_metadata_count, sizeof(struct ia2_destructor_runtime_entry));
+  if (!destructor_entries) {
+    fprintf(stderr, "ia2: failed to allocate destructor metadata cache\n");
+    return;
+  }
+
+  for (unsigned int i = 0; i < ia2_destructor_metadata_count; ++i) {
+    const struct ia2_destructor_metadata_record *record = &ia2_destructor_metadata[i];
+
+    dlerror();
+    void *symbol = dlsym(RTLD_DEFAULT, record->wrapper);
+    const char *dlerr = dlerror();
+    if (dlerr != NULL || symbol == NULL) {
+      fprintf(stderr, "ia2: failed to resolve destructor wrapper '%s': %s\n",
+              record->wrapper, dlerr ? dlerr : "symbol not found");
+      continue;
+    }
+
+    struct ia2_destructor_runtime_entry *entry = &destructor_entries[destructor_entry_count++];
+    entry->wrapper = (void (*)(void))symbol;
+    entry->record = record;
+    if (record->uses_union_pkru) {
+      entry->pkru_value = compute_union_pkru(ia2_destructor_exit_pkey, record->compartment_pkey);
+    } else {
+      entry->pkru_value = PKRU(record->compartment_pkey);
+    }
+  }
+}
+
+void ia2_destructor_runtime_init(void) {
+  ensure_initialized();
+}
+
+bool ia2_destructor_metadata_lookup(void (*wrapper)(void),
+                                    const struct ia2_destructor_metadata_record **out_record,
+                                    uint32_t *out_pkru_value) {
+  ensure_initialized();
+
+  if (!wrapper || destructor_entry_count == 0) {
+    return false;
+  }
+
+  for (size_t i = 0; i < destructor_entry_count; ++i) {
+    if (destructor_entries[i].wrapper == wrapper) {
+      if (out_record) {
+        *out_record = destructor_entries[i].record;
+      }
+      if (out_pkru_value) {
+        *out_pkru_value = destructor_entries[i].pkru_value;
+      }
+      return true;
+    }
+  }
+  return false;
+}
+
+uint32_t ia2_destructor_pkru_for(void (*wrapper)(void), int target_compartment_pkey) {
+  uint32_t pkru_value;
+  if (ia2_destructor_metadata_lookup(wrapper, NULL, &pkru_value)) {
+    return pkru_value;
+  }
+  return PKRU(target_compartment_pkey);
+}
+
+// Helper to read/write PKRU, guarded for non-x86 architectures.
+#if defined(__x86_64__)
+static inline uint32_t read_pkru(void) {
+  uint32_t pkru;
+  __asm__ volatile("xor %%ecx, %%ecx\n"
+                   "rdpkru\n"
+                   : "=a"(pkru)
+                   :
+                   : "ecx", "edx");
+  return pkru;
+}
+
+static inline void write_pkru(uint32_t pkru) {
+  __asm__ volatile("xor %%ecx, %%ecx\n"
+                   "xor %%edx, %%edx\n"
+                   "wrpkru\n"
+                   :
+                   : "a"(pkru)
+                   : "ecx", "edx");
+}
+#endif
+
+uint32_t ia2_destructor_enter(void *wrapper_addr, int target_pkey) {
+#if defined(__x86_64__)
+  uint32_t original_pkru = read_pkru();
+  uint32_t desired_pkru = ia2_destructor_pkru_for((void (*)(void))wrapper_addr, target_pkey);
+  if (desired_pkru != original_pkru) {
+    write_pkru(desired_pkru);
+#ifdef IA2_TRACE_EXIT
+    ia2_trace_exit_record((int)ia2_get_compartment(), target_pkey, desired_pkru);
+#endif
+  }
+  return original_pkru;
+#else
+  (void)wrapper_addr;
+  (void)target_pkey;
+  return 0;
+#endif
+}
+
+void ia2_destructor_leave(uint32_t original_pkru) {
+#if defined(__x86_64__)
+  write_pkru(original_pkru);
+#else
+  (void)original_pkru;
+#endif
+}

runtime/libia2/include/ia2_destructor_runtime.h

@@ -0,0 +1,53 @@
+#pragma once
+
+#include <stdbool.h>
+#include <stdint.h>
+
+struct ia2_destructor_metadata_record {
+  const char *wrapper;
+  const char *target;
+  int compartment_pkey;
+  int uses_union_pkru;
+};
+
+extern const struct ia2_destructor_metadata_record ia2_destructor_metadata[];
+extern const unsigned int ia2_destructor_metadata_count;
+extern const int ia2_destructor_exit_pkey;
+
+/// Initialize runtime bookkeeping for destructor metadata.
+///
+/// Must be called after all compartments have been set up so the rewriter
+/// generated symbols are visible via dlsym(). Safe to call multiple times; the
+/// function short-circuits if re-invoked.
+void ia2_destructor_runtime_init(void);
+
+/// Compute the PKRU mask that should be applied while running a destructor
+/// wrapper.
+///
+/// `wrapper` must be the exact function pointer registered with `__cxa_atexit`.
+/// When metadata is present the returned value reflects whether the
+/// destructor requires the union (exit compartment + target compartment)
+/// permissions; otherwise the caller gets the default single-compartment mask
+/// derived from `target_compartment`.
+uint32_t ia2_destructor_pkru_for(void (*wrapper)(void), int target_compartment_pkey);
+
+/// Lookup helper primarily intended for diagnostics and tracing.
+///
+/// Returns true when metadata was found for `wrapper` and fills the out
+/// parameters. Callers may pass NULL for any output they do not care about.
+bool ia2_destructor_metadata_lookup(void (*wrapper)(void),
+                                    const struct ia2_destructor_metadata_record **out_record,
+                                    uint32_t *out_pkru_value);
+
+/// Destructor wrapper enter hook - called before switching compartments.
+///
+/// Reads the current PKRU, computes the desired value for the wrapper based on
+/// metadata, applies it, and returns the original PKRU so the caller can
+/// restore it afterwards. `wrapper_addr` must point at the wrapper function
+/// itself; `target_pkey` is used as a fallback if metadata is unavailable.
+uint32_t ia2_destructor_enter(void *wrapper_addr, int target_pkey);
+
+/// Destructor wrapper leave hook - called before returning from wrapper.
+///
+/// Restores the PKRU value captured by ia2_destructor_enter.
+void ia2_destructor_leave(uint32_t original_pkru);

runtime/libia2/include/ia2_internal.h

@@ -18,6 +18,7 @@ struct dl_phdr_info;
 #include <link.h>
 #include <locale.h>
 #include <stdbool.h>
+#include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -448,3 +449,4 @@ __attribute__((__noreturn__)) void ia2_reinit_stack_err(int i);
 void **ia2_stackptr_for_tag(size_t tag);
 void **ia2_stackptr_for_compartment(int compartment);
 void ia2_setup_destructors(void);
+void ia2_trace_exit_record(int caller_pkey, int target_pkey, uint32_t pkru_value);

runtime/libia2/init.c

@@ -3,6 +3,7 @@
 #endif
 #include "ia2.h"
 #include "ia2_internal.h"
+#include "ia2_destructor_runtime.h"
 #include "memory_maps.h"
 #include "thread_name.h"
 #include <dlfcn.h>
@@ -287,10 +288,7 @@ void ia2_start(void) {
   ia2_log("initializing ia2 runtime\n");
   /* Get the user config before doing anything else */
   ia2_main();
-  // DISABLED FOR TESTING: Compartment destructors cause exit handler violations
-  // when libc is protected in compartment 1. The destructor wrappers switch
-  // to compartment 0, causing SEGV_PKUERR when accessing libc's .bss
-  // ia2_setup_destructors();
+  ia2_setup_destructors();
   /* Set up global resources. */
   ia2_set_up_tags();
   create_thread_keys();
@@ -311,5 +309,6 @@ void ia2_start(void) {
       exit(rc);
     }
   }
+  ia2_destructor_runtime_init();
   mark_init_finished();
 }

runtime/libia2/main.c

@@ -48,7 +48,6 @@ __asm__(
     // NOTE: Removed switch to compartment 0 to allow exit handlers to run
     // in compartment 1 (where libc lives). This prevents SEGV_PKUERR when
     // exit() tries to acquire __exit_funcs_lock in libc's .bss section.
-    // See: tests/dl_debug_test/*_ANALYSIS.md for details
     // "xor %ecx,%ecx\n"
     // "xor %edx,%edx\n"
     // "mov_pkru_eax 0\n"

runtime/libia2/trace_exit.c

@@ -0,0 +1,16 @@
+#include "ia2.h"
+
+#include <stdint.h>
+#include <stdio.h>
+
+void ia2_trace_exit_record(int caller_pkey, int target_pkey, uint32_t pkru_value) {
+#if IA2_TRACE_EXIT
+  fprintf(stderr,
+          "[TRACE_EXIT] caller_pkru=0x%08x caller_pkey=%d target_pkey=%d\n",
+          pkru_value, caller_pkey, target_pkey);
+#else
+  (void)caller_pkey;
+  (void)target_pkey;
+  (void)pkru_value;
+#endif
+}