feat: detect stack overflows on all platforms and threads

src/runtime/stack_overflow.cpp

@@ -4,38 +4,66 @@ Released under Apache 2.0 license as described in the file LICENSE.
 
 Author: Sebastian Ullrich
 
-Install a segfault signal handler and abort with custom message if address is within stack guard.
-Inspired by https://github.com/rust-lang/rust/blob/master/src/libstd/sys/unix/stack_overflow.rs
+Print a nicer error message on stack overflow.
+Port of the corresponding Rust code (see links below).
 */
-#if !defined(LEAN_WINDOWS) && !defined(__APPLE__)
+#ifdef LEAN_WINDOWS
+#include <windows.h>
+#else
 #include <csignal>
+#include <pthread.h>
+#include <unistd.h>
+#endif
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
-#include <pthread.h>
-#include <unistd.h>
+#include "runtime/stack_overflow.h"
 
 namespace lean {
-// signal stack of the main thread
-static stack_t g_signal_stack;
+// stack guard of the main thread
+static stack_guard * g_stack_guard;
+
+#ifdef LEAN_WINDOWS
+// https://github.com/rust-lang/rust/blob/master/src/libstd/sys/windows/stack_overflow.rs
+
+LONG WINAPI stack_overflow_handler(struct _EXCEPTION_POINTERS * info) {
+    if (info->ExceptionRecord->ExceptionCode == EXCEPTION_STACK_OVERFLOW) {
+        fprintf(stderr, "\nStack overflow detected. Aborting.\n");
+        abort();
+    } else {
+        return EXCEPTION_CONTINUE_SEARCH;
+    }
+}
+
+stack_guard::stack_guard() {
+    // reserve some stack space for the handler
+    ULONG sz = 0x5000;
+    SetThreadStackGuarantee(&sz);
+}
+
+stack_guard::~stack_guard() {}
+#else
+// Install a segfault signal handler and abort with custom message if address is within stack guard.
+// https://github.com/rust-lang/rust/blob/master/src/libstd/sys/unix/stack_overflow.rs
+
 
 // https://github.com/rust-lang/rust/blob/7c8dbd969dd0ef2af6d8bab9e03ba7ce6cac41a2/src/libstd/sys/unix/thread.rs#L293
 bool is_within_stack_guard(void * addr) {
+    char * stackaddr;
+#ifdef __APPLE__
+    stackaddr = static_cast<char *>(pthread_get_stackaddr_np(pthread_self())) - pthread_get_stacksize_np(pthread_self());
+#else
     pthread_attr_t attr;
     if (pthread_attr_init(&attr) != 0) return false;
     pthread_getattr_np(pthread_self(), &attr);
-    void * stackaddr;
-    size_t stacksize, guardsize;
-    pthread_attr_getstack(&attr, &stackaddr, &stacksize);
-    pthread_attr_getguardsize(&attr, &guardsize);
+    size_t stacksize;
+    pthread_attr_getstack(&attr, reinterpret_cast<void **>(&stackaddr), &stacksize);
     pthread_attr_destroy(&attr);
-    if (guardsize == 0) {
-        // probably the main thread, make an educated guess
-        // https://github.com/rust-lang/rust/blob/7c8dbd969dd0ef2af6d8bab9e03ba7ce6cac41a2/src/libstd/sys/unix/thread.rs#L343-L347
-        guardsize = static_cast<size_t>(sysconf(_SC_PAGESIZE));
-    }
+#endif
+    // close enough; the actual guard might be bigger, but it's unlikely a Lean function will have stack frames that big
+    size_t guardsize = static_cast<size_t>(sysconf(_SC_PAGESIZE));
     // the stack guard is *below* the stack
-    return stackaddr > addr && addr >= static_cast<char *>(stackaddr) - guardsize;
+    return stackaddr - guardsize <= addr && addr < stackaddr;
 }
 
 extern "C" void segv_handler(int signum, siginfo_t * info, void *) {
@@ -51,40 +79,36 @@ extern "C" void segv_handler(int signum, siginfo_t * info, void *) {
     }
 }
 
-// We need a separate signal stack since we can't use the overflowed stack
-void setup_signal_stack(stack_t & stack) {
-    stack.ss_sp = malloc(SIGSTKSZ);
-    if (stack.ss_sp == nullptr) return;
-    stack.ss_size = SIGSTKSZ;
-    stack.ss_flags = 0;
-    sigaltstack(&stack, nullptr);
+stack_guard::stack_guard() {
+    m_signal_stack.ss_sp = malloc(SIGSTKSZ);
+    if (m_signal_stack.ss_sp == nullptr) return;
+    m_signal_stack.ss_size = SIGSTKSZ;
+    m_signal_stack.ss_flags = 0;
+    sigaltstack(&m_signal_stack, nullptr);
 }
 
-void free_signal_stack(stack_t & stack) {
-    if (!stack.ss_sp) {
-        return;
-    }
-    stack.ss_flags = SS_DISABLE;
-    sigaltstack(&stack, nullptr);
-    free(stack.ss_sp);
+stack_guard::~stack_guard() {
+    if (!m_signal_stack.ss_sp) return;
+    m_signal_stack.ss_flags = SS_DISABLE;
+    sigaltstack(&m_signal_stack, nullptr);
+    free(m_signal_stack.ss_sp);
 }
+#endif
 
 void initialize_stack_overflow() {
-    setup_signal_stack(g_signal_stack);
+    g_stack_guard = new stack_guard();
+#ifdef LEAN_WINDOWS
+    AddVectoredExceptionHandler(0, stack_overflow_handler);
+#else
     struct sigaction action;
     memset(&action, 0, sizeof(struct sigaction));
     action.sa_flags = SA_SIGINFO | SA_ONSTACK;
     action.sa_sigaction = segv_handler;
     sigaction(SIGSEGV, &action, nullptr);
+#endif
 }
 
 void finalize_stack_overflow() {
-    free_signal_stack(g_signal_stack);
+    delete g_stack_guard;
 }
 }
-#else
-namespace lean {
-void initialize_stack_overflow() {}
-void finalize_stack_overflow() {}
-}
-#endif

src/runtime/stack_overflow.h

@@ -5,8 +5,21 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Author: Sebastian Ullrich
 */
 #pragma once
+#ifndef LEAN_WINDOWS
+#include <csignal>
+#endif
 
 namespace lean {
+class stack_guard {
+#ifndef LEAN_WINDOWS
+    // We need a separate signal stack since we can't use the overflowed stack
+    stack_t m_signal_stack;
+#endif
+public:
+    stack_guard();
+    ~stack_guard();
+};
+
 void initialize_stack_overflow();
 void finalize_stack_overflow();
 }

src/runtime/thread.cpp

@@ -17,6 +17,7 @@ Author: Leonardo de Moura
 #include "runtime/interrupt.h"
 #include "runtime/exception.h"
 #include "runtime/alloc.h"
+#include "runtime/stack_overflow.h"
 
 #ifndef LEAN_DEFAULT_THREAD_STACK_SIZE
 #define LEAN_DEFAULT_THREAD_STACK_SIZE 8*1024*1024 // 8Mb
@@ -108,6 +109,7 @@ struct lthread::imp {
     bool                      m_joined = false;
 
     static void * _main(void * p) {
+        stack_guard guard;
         thread_main(p);
         return nullptr;
     }

tests/compiler/StackOverflow.lean

@@ -0,0 +1,3 @@
+partial def foo : Nat → Nat | n => foo n + 1
+@[neverExtract]
+def main : IO Unit := IO.println $ foo 0

tests/compiler/StackOverflow.lean.expected.out

@@ -0,0 +1,2 @@
+
+Stack overflow detected. Aborting.

tests/compiler/StackOverflowTask.lean

@@ -0,0 +1,3 @@
+partial def foo : Nat → Nat | n => foo n + 1
+@[neverExtract]
+def main : IO Unit := IO.println $ Task.get $ Task.mk $ fun _ => foo 0

tests/compiler/StackOverflowTask.lean.expected.out

@@ -0,0 +1,2 @@
+
+Stack overflow detected. Aborting.