[3.11] GH-96569: Avoid undefined behavior (#96616)

Co-authored-by: Michael Droettboom <mdboom@gmail.com>

3 files changed, 25 insertions, 16 deletions

diff --git a/Include/internal/pycore_frame.h b/Include/internal/pycore_frame.h
index efdcdbc..ecc8346 100644
--- a/Include/internal/pycore_frame.h
+++ b/Include/internal/pycore_frame.h
@@ -192,17 +192,25 @@ _PyFrame_LocalsToFast(_PyInterpreterFrame *frame, int clear);
 extern _PyInterpreterFrame *
 _PyThreadState_BumpFramePointerSlow(PyThreadState *tstate, size_t size);
 
+static inline bool
+_PyThreadState_HasStackSpace(PyThreadState *tstate, size_t size)
+{
+    assert(
+        (tstate->datastack_top == NULL && tstate->datastack_limit == NULL)
+        ||
+        (tstate->datastack_top != NULL && tstate->datastack_limit != NULL)
+    );
+    return tstate->datastack_top != NULL &&
+        size < (size_t)(tstate->datastack_limit - tstate->datastack_top);
+}
+
 static inline _PyInterpreterFrame *
 _PyThreadState_BumpFramePointer(PyThreadState *tstate, size_t size)
 {
-    PyObject **base = tstate->datastack_top;
-    if (base) {
-        PyObject **top = base + size;
-        assert(tstate->datastack_limit);
-        if (top < tstate->datastack_limit) {
-            tstate->datastack_top = top;
-            return (_PyInterpreterFrame *)base;
-        }
+    if (_PyThreadState_HasStackSpace(tstate, size)) {
+        _PyInterpreterFrame *res = (_PyInterpreterFrame *)tstate->datastack_top;
+        tstate->datastack_top += size;
+        return res;
     }
     return _PyThreadState_BumpFramePointerSlow(tstate, size);
 }
diff --git a/Misc/NEWS.d/next/Core and Builtins/2022-09-05-16-43-44.gh-issue-96569.9lmTCC.rst b/Misc/NEWS.d/next/Core and Builtins/2022-09-05-16-43-44.gh-issue-96569.9lmTCC.rst
new file mode 100644
index 0000000..1dc4c08
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2022-09-05-16-43-44.gh-issue-96569.9lmTCC.rst
@@ -0,0 +1 @@
+Remove two cases of undefined behavior, by adding NULL checks.
diff --git a/Python/pystate.c b/Python/pystate.c
index a0c12ba..3a8140b 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -2178,16 +2178,16 @@ push_chunk(PyThreadState *tstate, int size)
 _PyInterpreterFrame *
 _PyThreadState_BumpFramePointerSlow(PyThreadState *tstate, size_t size)
 {
-    assert(size < INT_MAX/sizeof(PyObject *));
-    PyObject **base = tstate->datastack_top;
-    PyObject **top = base + size;
-    if (top >= tstate->datastack_limit) {
-        base = push_chunk(tstate, (int)size);
+    if (_PyThreadState_HasStackSpace(tstate, size)) {
+        _PyInterpreterFrame *res = (_PyInterpreterFrame *)tstate->datastack_top;
+        tstate->datastack_top += size;
+        return res;
     }
-    else {
-        tstate->datastack_top = top;
+    if (size > INT_MAX/2) {
+        PyErr_NoMemory();
+        return NULL;
     }
-    return (_PyInterpreterFrame *)base;
+    return (_PyInterpreterFrame *)push_chunk(tstate, (int)size);
 }
 
 void