[compiler-rt] [safestack] store unsafe stack pointer in TCB (PR #85878)

[via llvm-commits]

https://github.com/0x00004000 created https://github.com/llvm/llvm-project/pull/85878

I modified the safestack runtime library to store the unsafe stack pointer directly in the TCB. By storing the pointer in the TCB, instead of a thread-local variable, access to the unsafe stack pointer can be faster by eliminating an additional level of indirection.

>From f0b592058c5b9e5ddd7b4a2ded679a03442fdfe2 Mon Sep 17 00:00:00 2001
From: 0x00004000 <jam at cock.pm>
Date: Tue, 19 Mar 2024 23:25:04 +0000
Subject: [PATCH] [safestack] store unsafe stack pointer in TCB

I modified the safestack runtime library to store the unsafe stack pointer directly in the TCB. By storing the pointer in the TCB, instead of a thread-local variable, access to the unsafe stack pointer can be faster by eliminating an additional level of indirection.
---
 compiler-rt/lib/safestack/safestack.cpp | 31 +++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/compiler-rt/lib/safestack/safestack.cpp b/compiler-rt/lib/safestack/safestack.cpp
index 0751f3988b9c1a..5b1f7d6bf0ecbd 100644
--- a/compiler-rt/lib/safestack/safestack.cpp
+++ b/compiler-rt/lib/safestack/safestack.cpp
@@ -37,6 +37,24 @@ __attribute__((visibility(
     "default"))) __thread void *__safestack_unsafe_stack_ptr = nullptr;
 }
 
+#if defined(__linux__) && defined(__x86_64__)
+#include <stdint.h>
+#include <sys/types.h>
+
+typedef struct {
+  void *tcb;        /* Pointer to the TCB.  Not necessarily the
+                       thread descriptor used by libpthread.  */
+  uint64_t *stack_guard;  /* Stack guard value for the thread.  */
+  void *unsafe_stack_ptr; /* Unsafe stack pointer for the thread. */
+} tcbhead_t;
+
+extern "C" {
+static __thread tcbhead_t tcb __attribute__((tls_model("initial-exec")));
+}
+#elif defined(__FreeBSD__)
+// TODO: Add TCB structure modification for FreeBSD
+#endif
+
 namespace {
 
 // TODO: The runtime library does not currently protect the safe stack beyond
@@ -103,7 +121,14 @@ inline void unsafe_stack_setup(void *start, size_t size, size_t guard) {
   void *stack_ptr = (char *)start + size;
   SFS_CHECK((((size_t)stack_ptr) & (kStackAlign - 1)) == 0);
 
+#if defined(__linux__) && defined(__x86_64__)
+  tcb.unsafe_stack_ptr = stack_ptr;
+#elif defined(__FreeBSD__)
+  // TODO: Store unsafe_stack_ptr in TCB for FreeBSD
+#else
   __safestack_unsafe_stack_ptr = stack_ptr;
+#endif
+
   unsafe_stack_start = start;
   unsafe_stack_size = size;
   unsafe_stack_guard = guard;
@@ -306,5 +331,11 @@ extern "C"
 
 extern "C"
     __attribute__((visibility("default"))) void *__get_unsafe_stack_ptr() {
+#if defined(__linux__) && defined(__x86_64__)
+  return tcb.unsafe_stack_ptr;
+#elif defined(__FreeBSD__)
+  // TODO: Retrieve unsafe_stack_ptr from TCB for FreeBSD
+#else
   return __safestack_unsafe_stack_ptr;
+#endif
 }