gh-139653: Add PyUnstable_ThreadState_SetStackProtection() (GH-139668)

Add PyUnstable_ThreadState_SetStackProtection() and
PyUnstable_ThreadState_ResetStackProtection() functions
to set the stack base address and stack size of a Python
thread state.

Co-authored-by: Petr Viktorin <encukou@gmail.com>

Author: encukou

Doc/c-api/exceptions.rst

@@ -984,6 +984,9 @@ because the :ref:`call protocol <call>` takes care of recursion handling.
    be concatenated to the :exc:`RecursionError` message caused by the recursion
    depth limit.
 
+   .. seealso::
+      The :c:func:`PyUnstable_ThreadState_SetStackProtection` function.
+
    .. versionchanged:: 3.9
       This function is now also available in the :ref:`limited API <limited-c-api>`.
 

Doc/c-api/init.rst

@@ -1531,6 +1531,43 @@ All of the following functions must be called after :c:func:`Py_Initialize`.
    .. versionadded:: 3.11
 
 
+.. c:function:: int PyUnstable_ThreadState_SetStackProtection(PyThreadState *tstate, void *stack_start_addr, size_t stack_size)
+
+   Set the stack protection start address and stack protection size
+   of a Python thread state.
+
+   On success, return ``0``.
+   On failure, set an exception and return ``-1``.
+
+   CPython implements :ref:`recursion control <recursion>` for C code by raising
+   :py:exc:`RecursionError` when it notices that the machine execution stack is close
+   to overflow. See for example the :c:func:`Py_EnterRecursiveCall` function.
+   For this, it needs to know the location of the current thread's stack, which it
+   normally gets from the operating system.
+   When the stack is changed, for example using context switching techniques like the
+   Boost library's ``boost::context``, you must call
+   :c:func:`~PyUnstable_ThreadState_SetStackProtection` to inform CPython of the change.
+
+   Call :c:func:`~PyUnstable_ThreadState_SetStackProtection` either before
+   or after changing the stack.
+   Do not call any other Python C API between the call and the stack
+   change.
+
+   See :c:func:`PyUnstable_ThreadState_ResetStackProtection` for undoing this operation.
+
+   .. versionadded:: next
+
+
+.. c:function:: void PyUnstable_ThreadState_ResetStackProtection(PyThreadState *tstate)
+
+   Reset the stack protection start address and stack protection size
+   of a Python thread state to the operating system defaults.
+
+   See :c:func:`PyUnstable_ThreadState_SetStackProtection` for an explanation.
+
+   .. versionadded:: next
+
+
 .. c:function:: PyInterpreterState* PyInterpreterState_Get(void)
 
    Get the current interpreter.

Include/cpython/pystate.h

@@ -243,6 +243,18 @@ PyAPI_FUNC(int) PyGILState_Check(void);
 */
 PyAPI_FUNC(PyObject*) _PyThread_CurrentFrames(void);
 
+// Set the stack protection start address and stack protection size
+// of a Python thread state
+PyAPI_FUNC(int) PyUnstable_ThreadState_SetStackProtection(
+    PyThreadState *tstate,
+    void *stack_start_addr,  // Stack start address
+    size_t stack_size);      // Stack size (in bytes)
+
+// Reset the stack protection start address and stack protection size
+// of a Python thread state
+PyAPI_FUNC(void) PyUnstable_ThreadState_ResetStackProtection(
+    PyThreadState *tstate);
+
 /* Routines for advanced debuggers, requested by David Beazley.
    Don't use unless you know what you are doing! */
 PyAPI_FUNC(PyInterpreterState *) PyInterpreterState_Main(void);

Include/internal/pycore_pythonrun.h

@@ -54,6 +54,12 @@ extern const char* _Py_SourceAsString(
 #  define _PyOS_STACK_MARGIN_SHIFT (_PyOS_LOG2_STACK_MARGIN + 2)
 #endif
 
+#ifdef _Py_THREAD_SANITIZER
+#  define _PyOS_MIN_STACK_SIZE (_PyOS_STACK_MARGIN_BYTES * 6)
+#else
+#  define _PyOS_MIN_STACK_SIZE (_PyOS_STACK_MARGIN_BYTES * 3)
+#endif
+
 
 #ifdef __cplusplus
 }

Include/internal/pycore_tstate.h

@@ -37,6 +37,10 @@ typedef struct _PyThreadStateImpl {
     uintptr_t c_stack_soft_limit;
     uintptr_t c_stack_hard_limit;
 
+    // PyUnstable_ThreadState_ResetStackProtection() values
+    uintptr_t c_stack_init_base;
+    uintptr_t c_stack_init_top;
+
     PyObject *asyncio_running_loop; // Strong reference
     PyObject *asyncio_running_task; // Strong reference
 

Misc/NEWS.d/next/C_API/2025-10-06-22-17-47.gh-issue-139653.6-1MOd.rst

@@ -0,0 +1,4 @@
+Add :c:func:`PyUnstable_ThreadState_SetStackProtection` and
+:c:func:`PyUnstable_ThreadState_ResetStackProtection` functions to set the
+stack protection base address and stack protection size of a Python thread
+state. Patch by Victor Stinner.

Modules/_testinternalcapi.c

@@ -2408,6 +2408,58 @@ set_vectorcall_nop(PyObject *self, PyObject *func)
     Py_RETURN_NONE;
 }
 
+
+static void
+check_threadstate_set_stack_protection(PyThreadState *tstate,
+                                       void *start, size_t size)
+{
+    assert(PyUnstable_ThreadState_SetStackProtection(tstate, start, size) == 0);
+    assert(!PyErr_Occurred());
+
+    _PyThreadStateImpl *ts = (_PyThreadStateImpl *)tstate;
+    assert(ts->c_stack_top == (uintptr_t)start + size);
+    assert(ts->c_stack_hard_limit <= ts->c_stack_soft_limit);
+    assert(ts->c_stack_soft_limit < ts->c_stack_top);
+}
+
+
+static PyObject *
+test_threadstate_set_stack_protection(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    PyThreadState *tstate = PyThreadState_GET();
+    _PyThreadStateImpl *ts = (_PyThreadStateImpl *)tstate;
+    assert(!PyErr_Occurred());
+
+    uintptr_t init_base = ts->c_stack_init_base;
+    size_t init_top = ts->c_stack_init_top;
+
+    // Test the minimum stack size
+    size_t size = _PyOS_MIN_STACK_SIZE;
+    void *start = (void*)(_Py_get_machine_stack_pointer() - size);
+    check_threadstate_set_stack_protection(tstate, start, size);
+
+    // Test a larger size
+    size = 7654321;
+    assert(size > _PyOS_MIN_STACK_SIZE);
+    start = (void*)(_Py_get_machine_stack_pointer() - size);
+    check_threadstate_set_stack_protection(tstate, start, size);
+
+    // Test invalid size (too small)
+    size = 5;
+    start = (void*)(_Py_get_machine_stack_pointer() - size);
+    assert(PyUnstable_ThreadState_SetStackProtection(tstate, start, size) == -1);
+    assert(PyErr_ExceptionMatches(PyExc_ValueError));
+    PyErr_Clear();
+
+    // Test PyUnstable_ThreadState_ResetStackProtection()
+    PyUnstable_ThreadState_ResetStackProtection(tstate);
+    assert(ts->c_stack_init_base == init_base);
+    assert(ts->c_stack_init_top == init_top);
+
+    Py_RETURN_NONE;
+}
+
+
 static PyMethodDef module_functions[] = {
     {"get_configs", get_configs, METH_NOARGS},
     {"get_recursion_depth", get_recursion_depth, METH_NOARGS},
@@ -2516,6 +2568,8 @@ static PyMethodDef module_functions[] = {
     {"emscripten_set_up_async_input_device", emscripten_set_up_async_input_device, METH_NOARGS},
 #endif
     {"set_vectorcall_nop", set_vectorcall_nop, METH_O},
+    {"test_threadstate_set_stack_protection",
+     test_threadstate_set_stack_protection, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };
 

Python/ceval.c

@@ -425,7 +425,7 @@ int pthread_attr_destroy(pthread_attr_t *a)
 #endif
 
 static void
-hardware_stack_limits(uintptr_t *top, uintptr_t *base)
+hardware_stack_limits(uintptr_t *base, uintptr_t *top)
 {
 #ifdef WIN32
     ULONG_PTR low, high;
@@ -468,23 +468,86 @@ hardware_stack_limits(uintptr_t *top, uintptr_t *base)
 #endif
 }
 
-void
-_Py_InitializeRecursionLimits(PyThreadState *tstate)
+static void
+tstate_set_stack(PyThreadState *tstate,
+                 uintptr_t base, uintptr_t top)
 {
-    uintptr_t top;
-    uintptr_t base;
-    hardware_stack_limits(&top, &base);
+    assert(base < top);
+    assert((top - base) >= _PyOS_MIN_STACK_SIZE);
+
 #ifdef _Py_THREAD_SANITIZER
     // Thread sanitizer crashes if we use more than half the stack.
     uintptr_t stacksize = top - base;
-    base += stacksize/2;
+    base += stacksize / 2;
 #endif
     _PyThreadStateImpl *_tstate = (_PyThreadStateImpl *)tstate;
     _tstate->c_stack_top = top;
     _tstate->c_stack_hard_limit = base + _PyOS_STACK_MARGIN_BYTES;
     _tstate->c_stack_soft_limit = base + _PyOS_STACK_MARGIN_BYTES * 2;
+
+#ifndef NDEBUG
+    // Sanity checks
+    _PyThreadStateImpl *ts = (_PyThreadStateImpl *)tstate;
+    assert(ts->c_stack_hard_limit <= ts->c_stack_soft_limit);
+    assert(ts->c_stack_soft_limit < ts->c_stack_top);
+#endif
+}
+
+
+void
+_Py_InitializeRecursionLimits(PyThreadState *tstate)
+{
+    uintptr_t base, top;
+    hardware_stack_limits(&base, &top);
+    assert(top != 0);
+
+    tstate_set_stack(tstate, base, top);
+    _PyThreadStateImpl *ts = (_PyThreadStateImpl *)tstate;
+    ts->c_stack_init_base = base;
+    ts->c_stack_init_top = top;
+
+    // Test the stack pointer
+#if !defined(NDEBUG) && !defined(__wasi__)
+    uintptr_t here_addr = _Py_get_machine_stack_pointer();
+    assert(ts->c_stack_soft_limit < here_addr);
+    assert(here_addr < ts->c_stack_top);
+#endif
+}
+
+
+int
+PyUnstable_ThreadState_SetStackProtection(PyThreadState *tstate,
+                                void *stack_start_addr, size_t stack_size)
+{
+    if (stack_size < _PyOS_MIN_STACK_SIZE) {
+        PyErr_Format(PyExc_ValueError,
+                     "stack_size must be at least %zu bytes",
+                     _PyOS_MIN_STACK_SIZE);
+        return -1;
+    }
+
+    uintptr_t base = (uintptr_t)stack_start_addr;
+    uintptr_t top = base + stack_size;
+    tstate_set_stack(tstate, base, top);
+    return 0;
 }
 
+
+void
+PyUnstable_ThreadState_ResetStackProtection(PyThreadState *tstate)
+{
+    _PyThreadStateImpl *ts = (_PyThreadStateImpl *)tstate;
+    if (ts->c_stack_init_top != 0) {
+        tstate_set_stack(tstate,
+                         ts->c_stack_init_base,
+                         ts->c_stack_init_top);
+        return;
+    }
+
+    _Py_InitializeRecursionLimits(tstate);
+}
+
+
 /* The function _Py_EnterRecursiveCallTstate() only calls _Py_CheckRecursiveCall()
    if the recursion_depth reaches recursion_limit. */
 int

Python/pystate.c

@@ -1583,6 +1583,9 @@ init_threadstate(_PyThreadStateImpl *_tstate,
     _tstate->c_stack_top = 0;
     _tstate->c_stack_hard_limit = 0;
 
+    _tstate->c_stack_init_base = 0;
+    _tstate->c_stack_init_top = 0;
+
     _tstate->asyncio_running_loop = NULL;
     _tstate->asyncio_running_task = NULL;