2 files changed, 56 insertions, 1 deletions

diff --git a/compiler-rt/lib/safestack/safestack.cc b/compiler-rt/lib/safestack/safestack.cc
index 920b89b5e7c..673f5fd30c5 100644
--- a/compiler-rt/lib/safestack/safestack.cc
+++ b/compiler-rt/lib/safestack/safestack.cc
@@ -152,6 +152,7 @@ static void *thread_start(void *arg) {
 struct thread_stack_ll {
   struct thread_stack_ll *next;
   void *stack_base;
+  size_t size;
   pid_t pid;
   tid_t tid;
 };
@@ -183,7 +184,7 @@ static void thread_cleanup_handler(void *_iter) {
   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
@@ -193,6 +194,7 @@ static void thread_cleanup_handler(void *_iter) {
   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;
 
diff --git a/compiler-rt/test/safestack/pthread-stack-size.c b/compiler-rt/test/safestack/pthread-stack-size.c
new file mode 100644
index 00000000000..2fe7c680ef6
--- /dev/null
+++ b/compiler-rt/test/safestack/pthread-stack-size.c
@@ -0,0 +1,53 @@
+// 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;
+}