memory_maps: make impl lock-free using atomics

This reformulates how threading works for memory maps.
This gets rid of all of the locks and just uses an atomic counter
in exchange for guaranteeing that each thread registers only once at the start
(with `ia2_thread_metadata_new_for_current_thread`).

`ia2_thread_metadata_new_for_current_thread`
is called for the main thread in `ia2_start`
and in new threads in `ia2_thread_begin`.

Making this lock-free avoids any potential deadlocks and reentrancy issues,
meaning the memory map functions can be called in any context,
including in a debugger stopped at any point.

Author: kkysen

runtime/libia2/ia2.c

@@ -381,6 +381,7 @@ int protect_tls_pages(struct dl_phdr_info *info, size_t size, void *data) {
           exit(-1);
         }
 #if IA2_DEBUG_MEMORY
+        // Atomic write.
         thread_metadata->tls_addr_compartment1_first = (uintptr_t)start_round_down;
 #endif
       }
@@ -396,6 +397,7 @@ int protect_tls_pages(struct dl_phdr_info *info, size_t size, void *data) {
           exit(-1);
         }
 #if IA2_DEBUG_MEMORY
+        // Atomic write.
         thread_metadata->tls_addr_compartment1_second = (uintptr_t)after_untrusted_region_start;
 #endif
       }
@@ -409,6 +411,7 @@ int protect_tls_pages(struct dl_phdr_info *info, size_t size, void *data) {
         exit(-1);
       }
 #if IA2_DEBUG_MEMORY
+      // Atomic write.
       thread_metadata->tls_addrs[pkey] = (uintptr_t)start_round_down;
 #endif
     }

runtime/libia2/include/ia2_internal.h

@@ -49,6 +49,8 @@ struct dl_phdr_info;
 
 /// The data here is shared, so it should not be trusted for use as a pointer,
 /// but it can be used best effort for non-trusted purposes.
+///
+/// All fields should be used atomically.
 struct ia2_thread_metadata {
   pid_t tid;
   pthread_t thread;
@@ -81,9 +83,16 @@ struct ia2_thread_metadata {
 #define IA2_MAX_THREADS 512
 
 struct ia2_all_threads_metadata {
-  pthread_mutex_t lock;
-  size_t num_threads;
+  /// This is the number of threads registered,
+  /// and it is monotonically increasing by 1.
+  ///
+  /// It may be transiently higher than `IA2_MAX_THREADS`,
+  /// but will `abort` if that happens (other threads may observe a higher value).
+  _Atomic size_t num_threads;
+
   pid_t tids[IA2_MAX_THREADS];
+
+  /// Should be initialized to 0.
   struct ia2_thread_metadata thread_metadata[IA2_MAX_THREADS];
 };
 
@@ -493,7 +502,6 @@ __attribute__((__noreturn__)) void ia2_reinit_stack_err(int i);
   /* All zeroed, so this should go in `.bss` */                                                       \
   /* and only have pages lazily allocated. */                                                         \
   struct ia2_all_threads_metadata ia2_threads_metadata IA2_SHARED_DATA = {                            \
-      .lock = PTHREAD_MUTEX_INITIALIZER,                                                              \
       .num_threads = 0,                                                                               \
       .thread_metadata = {0},                                                                         \
   };

runtime/libia2/init.c

@@ -84,6 +84,7 @@ char *allocate_stack(int i) {
   ia2_log("allocating stack for compartment %d on thread %ld: %p..%p\n", i, (long)gettid(), stack, stack + STACK_SIZE);
 #if IA2_DEBUG_MEMORY
   struct ia2_thread_metadata *const thread_metadata = ia2_thread_metadata_get_for_current_thread();
+  // Atomic write.
   thread_metadata->stack_addrs[i] = (uintptr_t)stack;
 #endif
   assert(stacks[i] == NULL); // We should only be setting this once per thread compartment right after thread creation.
@@ -291,6 +292,9 @@ void ia2_start(void) {
   /* Set up global resources. */
   ia2_set_up_tags();
   create_thread_keys();
+#if IA2_DEBUG_MEMORY
+  ia2_thread_metadata_new_for_current_thread();
+#endif
   verify_tls_padding();
   /* allocate an unprotected stack for the untrusted compartment */
   allocate_stack_0();

runtime/libia2/memory_maps.c

@@ -2,6 +2,8 @@
 #include "ia2.h"
 #include "thread_name.h"
 
+#include <stdatomic.h>
+
 // Only enable this code that stores these addresses when debug logging is enabled.
 // This reduces the trusted codebase and avoids runtime overhead.
 #if IA2_DEBUG_MEMORY
@@ -10,97 +12,97 @@
 // so that it can be used in `ia2_internal.h` within `_IA2_INIT_RUNTIME`
 // to only initialize the `ia2_threads_metadata` global once.
 
-#define array_len(a) (sizeof(a) / sizeof(*(a)))
+#define min(a, b) ((a) < (b) ? (a) : (b))
 
-struct ia2_thread_metadata *ia2_all_threads_metadata_get_for_current_thread(struct ia2_all_threads_metadata *const this) {
-  const pid_t tid = gettid();
-
-  if (pthread_mutex_lock(&this->lock) != 0) {
-    perror("pthread_mutex_lock in ia2_all_threads_data_lookup failed");
+struct ia2_thread_metadata *ia2_all_threads_metadata_new_for_current_thread(struct ia2_all_threads_metadata *const this) {
+  const size_t thread = atomic_fetch_add(&this->num_threads, 1);
+  if (thread >= IA2_MAX_THREADS) {
+    fprintf(stderr, "created %zu threads, but can't store them all (max is IA2_MAX_THREADS: %zu)\n",
+            thread + 1, (size_t)IA2_MAX_THREADS);
     abort();
   }
 
-  if (this->num_threads >= array_len(this->thread_metadata)) {
-    fprintf(stderr, "created %zu threads, but can't store them all (max is IA2_MAX_THREADS)\n", this->num_threads);
-    abort();
-  }
-
-  struct ia2_thread_metadata *metadata = NULL;
-  for (size_t i = 0; i < this->num_threads; i++) {
-    if (this->tids[i] == tid) {
-      metadata = &this->thread_metadata[i];
-      goto unlock;
-    }
-  }
-
-  metadata = &this->thread_metadata[this->num_threads];
-  this->tids[this->num_threads] = tid;
-  this->num_threads++;
+  const pid_t tid = gettid();
+  this->tids[thread] = tid;
 
+  struct ia2_thread_metadata *metadata = &this->thread_metadata[thread];
   metadata->tid = tid;
   metadata->thread = pthread_self();
+  return metadata;
+}
 
-unlock:
-  if (pthread_mutex_unlock(&this->lock) != 0) {
-    perror("pthread_mutex_unlock in ia2_all_threads_data_lookup failed");
-    abort();
+struct ia2_thread_metadata *ia2_all_threads_metadata_get_for_current_thread(struct ia2_all_threads_metadata *const this) {
+  const pid_t tid = gettid();
+
+  // We won't see threads created/registered after this,
+  // but `ia2_all_threads_metadata_new_for_current_thread`
+  // was supposed to be called first for this function to find it.
+  const size_t num_threads = min(IA2_MAX_THREADS, atomic_load(&this->num_threads));
+
+  for (size_t thread = 0; thread < num_threads; thread++) {
+    if (this->tids[thread] == tid) {
+      return &this->thread_metadata[thread];
+    }
   }
 
-  return metadata;
+  fprintf(stderr,
+          "ia2_thread_metadata not found for thread %ld\n"
+          "ia2_thread_metadata_new_for_current_thread must not have been previously called on this thread\n",
+          (long)tid);
+  abort();
 }
 
 struct ia2_addr_location ia2_all_threads_metadata_find_addr(struct ia2_all_threads_metadata *const this, const uintptr_t addr) {
-  struct ia2_addr_location location = {
-      .name = NULL,
-      .thread_metadata = NULL,
-      .compartment = -1,
-  };
-  if (pthread_mutex_lock(&this->lock) != 0) {
-    perror("pthread_mutex_lock in ia2_all_threads_data_find_addr failed");
-    goto ret;
-  }
+  // We won't see threads created/registered after this,
+  // but this is supposed to be best effort, so that's okay.
+  const size_t num_threads = min(IA2_MAX_THREADS, atomic_load(&this->num_threads));
+
   for (size_t thread = 0; thread < this->num_threads; thread++) {
     const pid_t tid = this->tids[thread];
     const struct ia2_thread_metadata *const thread_metadata = &this->thread_metadata[thread];
 
     if (addr == thread_metadata->tls_addr_compartment1_first || addr == thread_metadata->tls_addr_compartment1_second) {
-      location.name = "tls";
-      location.thread_metadata = thread_metadata;
-      location.compartment = 1;
-      goto unlock;
+      return (struct ia2_addr_location){
+          .name = "tls",
+          .thread_metadata = thread_metadata,
+          .compartment = 1,
+      };
     }
 
     for (int compartment = 0; compartment < IA2_MAX_COMPARTMENTS; compartment++) {
       if (addr == thread_metadata->stack_addrs[compartment]) {
-        location.name = "stack";
-        location.thread_metadata = thread_metadata;
-        location.compartment = compartment;
-        goto unlock;
+        return (struct ia2_addr_location){
+            .name = "stack",
+            .thread_metadata = thread_metadata,
+            .compartment = compartment,
+        };
       }
       if (addr == thread_metadata->tls_addrs[compartment]) {
-        location.name = "tls";
-        location.thread_metadata = thread_metadata;
-        location.compartment = compartment;
-        goto unlock;
+        return (struct ia2_addr_location){
+            .name = "tls",
+            .thread_metadata = thread_metadata,
+            .compartment = compartment,
+        };
       }
     }
   }
 
-  goto unlock;
-
-unlock:
-  if (pthread_mutex_unlock(&this->lock) != 0) {
-    perror("pthread_mutex_unlock in ia2_all_threads_data_find_addr failed");
-  }
-ret:
-  return location;
+  return (struct ia2_addr_location){
+      .name = NULL,
+      .thread_metadata = NULL,
+      .compartment = -1,
+  };
 }
 
 // Moved `ia2_threads_metadata` from here to `ia2_internal.h`
 // so that it can be used in `_IA2_INIT_RUNTIME`
 // to only initialize the `ia2_threads_metadata` global once.
 extern struct ia2_all_threads_metadata ia2_threads_metadata;
 
+struct ia2_thread_metadata *ia2_thread_metadata_new_for_current_thread(void) {
+  return ia2_all_threads_metadata_new_for_current_thread(&ia2_threads_metadata);
+}
+
 struct ia2_thread_metadata *ia2_thread_metadata_get_for_current_thread(void) {
   return ia2_all_threads_metadata_get_for_current_thread(&ia2_threads_metadata);
 }

runtime/libia2/memory_maps.h

@@ -11,14 +11,26 @@
 // so that it can be used in `ia2_internal.h` within `_IA2_INIT_RUNTIME`
 // to only initialize the `ia2_threads_metadata` global once.
 
-/// Find the `struct ia2_thread_metadata*` for the current thread,
-/// adding (but not allocating) one if there isn't one yet.
-/// If there is no memory for more or some unexpected error,
-/// it `abort`s.  `NULL` is never returned.
-/// This is a purely lookup and/or additive operation,
-/// so the lifetime of the returned `struct ia2_thread_metadata*` is infinite,
-/// and since it's thread-specific,
-/// it is thread-safe to read and write.
+/// Allocate and initialize a new `ia2_thread_metadata` for the current thread.
+/// Importantly, this may only be called once per thread.
+///
+/// The returned pointer is stable, non-`NULL`,
+/// and will not be moved or deallocated/uninitialized.
+/// Operations on the `ia2_thread_metadata` must be atomic.
+///
+/// If too many threads are created, this will `abort`
+/// and `IA2_MAX_THREADS` can be increased.
+struct ia2_thread_metadata *ia2_thread_metadata_new_for_current_thread(void);
+
+/// Find the `ia2_thread_metadata` for the current thread.
+///
+/// `ia2_thread_metadata_new_for_current_thread`
+/// must have been previously called for this thread,
+/// or else the `ia2_thread_metadata` will not be found and this will `abort`.
+///
+/// The returned pointer is stable, non-`NULL`,
+/// and will not be moved or deallocated/uninitialized.
+/// Operations on the `ia2_thread_metadata` must be atomic.
 struct ia2_thread_metadata *ia2_thread_metadata_get_for_current_thread(void);
 
 struct ia2_addr_location {
@@ -29,6 +41,8 @@ struct ia2_addr_location {
   const char *name;
 
   /// The metadata of the thread this address belongs to.
+  ///
+  /// Fields must be read atomically.
   const struct ia2_thread_metadata *thread_metadata;
 
   /// The compartment this address is in.

runtime/libia2/threads.c

@@ -31,7 +31,8 @@ void *ia2_thread_begin(void *arg) {
   munmap(arg, sizeof(struct ia2_thread_thunk));
 
 #if IA2_DEBUG_MEMORY
-  struct ia2_thread_metadata *const thread_metadata = ia2_thread_metadata_get_for_current_thread();
+  struct ia2_thread_metadata *const thread_metadata = ia2_thread_metadata_new_for_current_thread();
+  // Atomic write.
   thread_metadata->start_fn = fn;
 #endif