[SafeStack] Allocate unsafe sigaltstack (#206463)

PR https://github.com/llvm/llvm-project/pull/196969 was approved and
merged but with `spr` the wrong base branch was set.

This merges the approved changes into `main`

---------

Co-authored-by: Paul Walker <paul.walker@arm.com>
Co-authored-by: Arseniy Obolenskiy <arseniy.obolenskiy@amd.com>

compiler-rt/include/sanitizer/safestack_interface.h

@@ -28,6 +28,21 @@ const void *SANITIZER_CDECL __safestack_get_unsafe_stack_bottom(void);
 /// Returns a pointer to the top of the unsafe stack of the current thread.
 const void *SANITIZER_CDECL __safestack_get_unsafe_stack_top(void);
 
+/// Returns a pointer to the top of the unsafe sigalt stack of the current
+/// thread.
+const void *SANITIZER_CDECL __safestack_get_unsafe_sigalt_stack_ptr(void);
+
+/// Returns a pointer to the bottom of the unsafe sigalt stack of the current
+/// thread.
+const void *SANITIZER_CDECL __safestack_get_unsafe_sigalt_stack_bottom(void);
+
+/// Returns a pointer to the top of the unsafe sigalt stack of the current
+/// thread.
+const void *SANITIZER_CDECL __safestack_get_unsafe_sigalt_stack_top(void);
+
+/// Set a new unsafe signal stack context to be used if SA_ONSTACK is set.
+int SANITIZER_CDECL __safestack_unsafe_sigaltstack(size_t ss_size);
+
 #ifdef __cplusplus
 } // extern "C"
 #endif

compiler-rt/lib/safestack/safestack.cpp

@@ -113,6 +113,14 @@ __thread void *unsafe_stack_start = nullptr;
 __thread size_t unsafe_stack_size = 0;
 __thread size_t unsafe_stack_guard = 0;
 
+// Per-thread unsafe stack information used for the unsafe stack during signal
+// handling if sigaltstack is used. Without, only the safe stack is switched by
+// the operating system. When the program indicates to use a separate stack for
+// signal handling, this should also include the unsafe stack component.
+__thread void* unsafe_sigalt_stack_ptr = nullptr;
+__thread void* unsafe_sigalt_stack_start = nullptr;
+__thread size_t unsafe_sigalt_stack_size = 0;
+
 inline void *unsafe_stack_alloc(size_t size, size_t guard) {
   SFS_CHECK(size + guard >= size);
   void *addr = Mmap(nullptr, size + guard, PROT_READ | PROT_WRITE,
@@ -134,6 +142,16 @@ inline void unsafe_stack_setup(void *start, size_t size, size_t guard) {
   unsafe_stack_guard = guard;
 }
 
+inline void unsafe_sigalt_stack_setup(void* start, size_t size) {
+  SFS_CHECK((uintptr_t)start + size >= (uintptr_t)start);
+  void* stack_ptr = (void*)((uintptr_t)start + size);
+  SFS_CHECK((((uintptr_t)stack_ptr) & (kStackAlign - 1)) == 0);
+
+  unsafe_sigalt_stack_ptr = stack_ptr;
+  unsafe_sigalt_stack_start = start;
+  unsafe_sigalt_stack_size = size;
+}
+
 /// Thread data for the cleanup handler
 pthread_key_t thread_cleanup_key;
 
@@ -225,6 +243,14 @@ void thread_cleanup_handler(void *_iter) {
   pthread_mutex_unlock(&thread_stacks_mutex);
 
   unsafe_stack_start = nullptr;
+
+  // In case the sigalt stack was allocated, we need to unmap the used memory.
+  if (unsafe_sigalt_stack_start) {
+    unsafe_sigalt_stack_ptr = nullptr;
+    Munmap(unsafe_sigalt_stack_start, unsafe_sigalt_stack_size);
+    unsafe_sigalt_stack_start = nullptr;
+    unsafe_sigalt_stack_size = 0;
+  }
 }
 
 void EnsureInterceptorsInitialized();
@@ -287,6 +313,30 @@ INTERCEPTOR(int, sigaction, int sig, const struct sigaction* act,
   return REAL(sigaction)(sig, act, oldact);
 }
 
+// Since sigaltstack is required to be async-signal-safe, we cannot simply
+// intercept it to allocate the the unsafe stack. Instead if the users wishes to
+// setup an unsafe sigalt stack can call unsafe_sigaltstack(ss_size size)
+// explicitliy.
+int setup_unsafe_sigaltstack(size_t ss_size) {
+  EnsureInterceptorsInitialized();
+
+  SFS_CHECK(ss_size);
+  ss_size = RoundUpTo(ss_size, kStackAlign);
+
+  // For now always map a new unsafe sigaltstack when setting a new
+  // sigaltstack. Potentially if the size is identical, this step can be
+  // skipped.
+  void* prev_sigalt_stack_start = unsafe_sigalt_stack_start;
+  size_t prev_sigalt_stack_size = unsafe_sigalt_stack_size;
+  void* sigalt_addr = unsafe_stack_alloc(ss_size, 0);
+  unsafe_sigalt_stack_setup(sigalt_addr, ss_size);
+  if (prev_sigalt_stack_start != nullptr) {
+    Munmap(prev_sigalt_stack_start, prev_sigalt_stack_size);
+  }
+
+  return 0;
+}
+
 pthread_mutex_t interceptor_init_mutex = PTHREAD_MUTEX_INITIALIZER;
 bool interceptors_inited = false;
 
@@ -354,6 +404,26 @@ __safestack_get_unsafe_stack_ptr() {
   return __safestack_unsafe_stack_ptr;
 }
 
+extern "C" SANITIZER_INTERFACE_ATTRIBUTE const void*
+__safestack_get_unsafe_sigalt_stack_bottom() {
+  return unsafe_sigalt_stack_start;
+}
+
+extern "C" SANITIZER_INTERFACE_ATTRIBUTE const void*
+__safestack_get_unsafe_sigalt_stack_top() {
+  return (char*)unsafe_sigalt_stack_start + unsafe_sigalt_stack_size;
+}
+
+extern "C" SANITIZER_INTERFACE_ATTRIBUTE const void*
+__safestack_get_unsafe_sigalt_stack_ptr() {
+  return unsafe_sigalt_stack_ptr;
+}
+
+extern "C" SANITIZER_INTERFACE_ATTRIBUTE int __safestack_unsafe_sigaltstack(
+    size_t ss_size) {
+  return setup_unsafe_sigaltstack(ss_size);
+}
+
 // Compatibility aliases
 extern "C" SANITIZER_INTERFACE_ATTRIBUTE void* __get_unsafe_stack_bottom() {
   return const_cast<void*>(__safestack_get_unsafe_stack_bottom());

compiler-rt/test/safestack/sigaltstack.c

@@ -49,6 +49,7 @@ void *t1_start(void *ptr) {
   sigstk.ss_size = ss_size;
   sigstk.ss_sp = ss_sp;
 
+  __safestack_unsafe_sigaltstack(sigstk.ss_size);
   sigaltstack(&sigstk, NULL);
 
   // Test that after sigaltstack is set, it signal handling still works.
@@ -64,6 +65,9 @@ int main() {
   char c[] = "hello world";
   puts(c);
 
+  // Make sure no sigaltstack is allocated by default.
+  assert(!__safestack_get_unsafe_sigalt_stack_ptr());
+
   stack_t sigstk = {};
   size_t ss_size = 4096 * 4;
   void *ss_sp = mmap(NULL, sigstk.ss_size, PROT_READ | PROT_WRITE,
@@ -71,8 +75,17 @@ int main() {
   sigstk.ss_size = ss_size;
   sigstk.ss_sp = ss_sp;
 
+  __safestack_unsafe_sigaltstack(sigstk.ss_size);
   sigaltstack(&sigstk, NULL);
 
+  // Make sure __safestack_unsafe_sigaltstack allocated the unsafe sigaltstack with the
+  // correct size.
+  assert(__safestack_get_unsafe_sigalt_stack_ptr());
+  assert(__safestack_get_unsafe_sigalt_stack_ptr() ==
+         __safestack_get_unsafe_sigalt_stack_top());
+  assert((__safestack_get_unsafe_sigalt_stack_top() -
+          __safestack_get_unsafe_sigalt_stack_bottom()) == sigstk.ss_size);
+
   // Make sure retrieving the sigaltstack works without problems.
   sigaltstack(NULL, &sigstk);
 
@@ -82,9 +95,16 @@ int main() {
   new_sigstk.ss_sp = mmap(NULL, new_sigstk.ss_size, PROT_READ | PROT_WRITE,
                           MAP_PRIVATE | MAP_ANONYMOUS | MAP_STACK, -1, 0);
 
+  __safestack_unsafe_sigaltstack(new_sigstk.ss_size);
   sigaltstack(&new_sigstk, NULL);
   munmap(ss_sp, ss_size);
 
+  // Make sure updating the size of the unsafe sigaltstack also updates when
+  // setting a new sigaltstack.
+  assert(__safestack_get_unsafe_sigalt_stack_ptr());
+  assert((__safestack_get_unsafe_sigalt_stack_top() -
+          __safestack_get_unsafe_sigalt_stack_bottom()) == new_sigstk.ss_size);
+
   struct sigaction sa;
   sa.sa_handler = signal_handler;
   sigemptyset(&sa.sa_mask);