SafeStack: Use correct unsafe stack sizes

[vlad.tsyrklevich]

Summary

When deallocating thread stacks, we use one thread's unsafe stack size
to deallocate another thread's unsafe stack; however, the two sizes may
differ. Record an unsafe stack's size in the thread stack linked list.

[eugenis]

Nice catch!

lib/safestack/safestack.cc

static void *thread_start(void *arg) {
 
   return start_routine(start_routine_arg);
 }
 
 /// Linked list used to store exiting threads stack/thread information.
 struct thread_stack_ll {
   struct thread_stack_ll *next;
   void *stack_base;
+  size_t size;
   pid_t pid;
   tid_t tid;
 };
 
 /// Linked list of unsafe stacks for threads that are exiting. We delay
 /// unmapping them until the thread exits.
 static thread_stack_ll *thread_stacks = nullptr;
 static pthread_mutex_t thread_stacks_mutex = PTHREAD_MUTEX_INITIALIZER;
static void thread_cleanup_handler(void *_iter) {
   pid_t pid = getpid();
   tid_t tid = GetTid();
 
   // Free stacks for dead threads
   thread_stack_ll **stackp = &temp_stacks;
   while (*stackp) {
     thread_stack_ll *stack = *stackp;
     if (stack->pid != pid || TgKill(stack->pid, stack->tid, 0) == -ESRCH) {
-      UnmapOrDie(stack->stack_base, unsafe_stack_size + unsafe_stack_guard);
+      UnmapOrDie(stack->stack_base, stack->size);
       *stackp = stack->next;
       free(stack);
     } else
       stackp = &stack->next;
   }
 
   thread_stack_ll *cur_stack =
       (thread_stack_ll *)malloc(sizeof(thread_stack_ll));
   cur_stack->stack_base = (char *)unsafe_stack_start - unsafe_stack_guard;
+  cur_stack->size = unsafe_stack_size + unsafe_stack_guard;
   cur_stack->pid = pid;
   cur_stack->tid = tid;
 
   pthread_mutex_lock(&thread_stacks_mutex);
   // Merge thread_stacks with the current thread's stack and any remaining
   // temp_stacks
   *stackp = thread_stacks;
   cur_stack->next = temp_stacks;

test/safestack/pthread-stack-size.c

+// RUN: %clang_safestack %s -pthread -o %t
+// RUN: %run %t
+
+// Test unsafe stack deallocation with custom stack sizes, in particular ensure
+// that we correctly deallocate small stacks and don't accidentally deallocate
+// adjacent memory.
+
+#include <pthread.h>
+#include <stdlib.h>
+#include <string.h>
+#include <unistd.h>
+
+volatile int step = 0;
+
+void *wait_until(void *ptr) {
+  while ((int)ptr != step)
+    usleep(1000);
+
+  volatile char buf[64];
+  buf[0] = 0;
+
+  return NULL;
+}
+
+int main(int argc, char **argv) {
+  pthread_t t1, t2, t3;
+
+  pthread_attr_t small_stack_attr;
+  pthread_attr_init(&small_stack_attr);
+  pthread_attr_setstacksize(&small_stack_attr, 65536);
+
+  if (pthread_create(&t3, NULL, wait_until, (void *)3))
+    abort();
+  if (pthread_create(&t1, &small_stack_attr, wait_until, (void *)1))
+    abort();
+  if (pthread_create(&t2, NULL, wait_until, (void *)2))
+    abort();
+
+  step = 1;
+  if (pthread_join(t1, NULL))
+    abort();
+
+  step = 2;
+  if (pthread_join(t2, NULL))
+    abort();
+
+  step = 3;
+  if (pthread_join(t3, NULL))
+    abort();
+
+  pthread_attr_destroy(&small_stack_attr);
+  return 0;
+}