memory_maps: make impl lock-free using atomics

runtime/libia2/ia2.c

@@ -323,7 +323,7 @@ int protect_tls_pages(struct dl_phdr_info *info, size_t size, void *data) {
   const int pkey = search_args->pkey;
 
 #if IA2_DEBUG_MEMORY
-  struct ia2_thread_metadata *const thread_metadata = ia2_thread_metadata_get_current_thread();
+  struct ia2_thread_metadata *const thread_metadata = ia2_thread_metadata_get_for_current_thread();
 #endif
 
   // Protect TLS segment.
@@ -381,9 +381,8 @@ int protect_tls_pages(struct dl_phdr_info *info, size_t size, void *data) {
           exit(-1);
         }
 #if IA2_DEBUG_MEMORY
-        if (thread_metadata) {
-          thread_metadata->tls_addr_compartment1_first = (uintptr_t)start_round_down;
-        }
+        // Atomic write.
+        thread_metadata->tls_addr_compartment1_first = (uintptr_t)start_round_down;
 #endif
       }
       uint64_t after_untrusted_region_start = untrusted_stackptr_addr + 0x1000;
@@ -398,9 +397,8 @@ int protect_tls_pages(struct dl_phdr_info *info, size_t size, void *data) {
           exit(-1);
         }
 #if IA2_DEBUG_MEMORY
-        if (thread_metadata) {
-          thread_metadata->tls_addr_compartment1_second = (uintptr_t)after_untrusted_region_start;
-        }
+        // Atomic write.
+        thread_metadata->tls_addr_compartment1_second = (uintptr_t)after_untrusted_region_start;
 #endif
       }
     } else {
@@ -413,9 +411,8 @@ int protect_tls_pages(struct dl_phdr_info *info, size_t size, void *data) {
         exit(-1);
       }
 #if IA2_DEBUG_MEMORY
-      if (thread_metadata) {
-        thread_metadata->tls_addrs[pkey] = (uintptr_t)start_round_down;
-      }
+      // 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

@@ -83,10 +83,9 @@ 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_current_thread();
-  if (thread_metadata) {
-    thread_metadata->stack_addrs[i] = (uintptr_t)stack;
-  }
+  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.
   stacks[i] = stack;
@@ -293,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,96 +12,99 @@
 // 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)))
-
-struct ia2_thread_metadata *ia2_all_threads_metadata_lookup(struct ia2_all_threads_metadata *const this) {
-  const pid_t tid = gettid();
+#define min(a, b) ((a) < (b) ? (a) : (b))
 
-  struct ia2_thread_metadata *metadata = NULL;
-  if (pthread_mutex_lock(&this->lock) != 0) {
-    perror("pthread_mutex_lock in ia2_all_threads_data_lookup failed");
-    goto ret;
-  }
-  for (size_t i = 0; i < this->num_threads; i++) {
-    if (this->tids[i] == tid) {
-      metadata = &this->thread_metadata[i];
-      goto unlock;
-    }
-  }
-  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);
-    goto unlock;
+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();
   }
 
-  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;
+}
+
+struct ia2_thread_metadata *ia2_all_threads_metadata_get_for_current_thread(struct ia2_all_threads_metadata *const this) {
+  const pid_t tid = gettid();
 
-  goto unlock;
+  // 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));
 
-unlock:
-  if (pthread_mutex_unlock(&this->lock) != 0) {
-    perror("pthread_mutex_unlock in ia2_all_threads_data_lookup failed");
+  for (size_t thread = 0; thread < num_threads; thread++) {
+    if (this->tids[thread] == tid) {
+      return &this->thread_metadata[thread];
+    }
   }
-ret:
-  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) {
+      return (struct ia2_addr_location){
+          .name = "tls",
+          .thread_metadata = thread_metadata,
+          .compartment = 1,
+      };
+    }
+
     for (int compartment = 0; compartment < IA2_MAX_COMPARTMENTS; compartment++) {
-      const struct ia2_thread_metadata *const thread_metadata = &this->thread_metadata[thread];
       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;
-      }
-      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 = 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_get_current_thread(void) {
-  return ia2_all_threads_metadata_lookup(&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);
 }
 
 struct ia2_addr_location ia2_addr_location_find(const uintptr_t addr) {
@@ -112,6 +117,7 @@ static void label_memory_map(FILE *log, uintptr_t start_addr) {
   const struct ia2_addr_location location = ia2_addr_location_find(start_addr);
   const struct ia2_thread_metadata *metadata = location.thread_metadata;
 
+  // If `location.name` is non-`NULL`, then `location` was found.
   if (location.name) {
     Dl_info dl_info = {0};
     const bool has_dl_info = dladdr((void *)metadata->start_fn, &dl_info);

runtime/libia2/memory_maps.h

@@ -11,14 +11,27 @@
 // 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 an error, `NULL` is 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.
-struct ia2_thread_metadata *ia2_thread_metadata_get_current_thread(void);
+/// 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 {
   /// A descriptive name of what this address points to.
@@ -28,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,10 +31,9 @@ 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_current_thread();
-  if (thread_metadata) {
-    thread_metadata->start_fn = fn;
-  }
+  struct ia2_thread_metadata *const thread_metadata = ia2_thread_metadata_new_for_current_thread();
+  // Atomic write.
+  thread_metadata->start_fn = fn;
 #endif
 
   init_stacks_and_setup_tls();