bpo-44698: Fix undefined behaviour in complex exponentiation. (GH-27278)

3 files changed, 37 insertions, 8 deletions

diff --git a/Lib/test/test_complex.py b/Lib/test/test_complex.py
index af39ee8..badb234 100644
--- a/Lib/test/test_complex.py
+++ b/Lib/test/test_complex.py
@@ -1,4 +1,5 @@
 import unittest
+import sys
 from test import support
 from test.test_grammar import (VALID_UNDERSCORE_LITERALS,
                                INVALID_UNDERSCORE_LITERALS)
@@ -248,6 +249,26 @@ class ComplexTest(unittest.TestCase):
         b = 5.1+2.3j
         self.assertRaises(ValueError, pow, a, b, 0)
 
+        # Check some boundary conditions; some of these used to invoke
+        # undefined behaviour (https://bugs.python.org/issue44698). We're
+        # not actually checking the results of these operations, just making
+        # sure they don't crash (for example when using clang's
+        # UndefinedBehaviourSanitizer).
+        values = (sys.maxsize, sys.maxsize+1, sys.maxsize-1,
+                  -sys.maxsize, -sys.maxsize+1, -sys.maxsize+1)
+        for real in values:
+            for imag in values:
+                with self.subTest(real=real, imag=imag):
+                    c = complex(real, imag)
+                    try:
+                        c ** real
+                    except OverflowError:
+                        pass
+                    try:
+                        c ** c
+                    except OverflowError:
+                        pass
+
     def test_boolcontext(self):
         for i in range(100):
             self.assertTrue(complex(random() + 1e-6, random() + 1e-6))
diff --git a/Misc/NEWS.d/next/Core and Builtins/2021-07-21-15-26-56.bpo-44698.DA4_0o.rst b/Misc/NEWS.d/next/Core and Builtins/2021-07-21-15-26-56.bpo-44698.DA4_0o.rst
new file mode 100644
index 0000000..ed38963
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2021-07-21-15-26-56.bpo-44698.DA4_0o.rst
@@ -0,0 +1 @@
+Fix undefined behaviour in complex object exponentiation.
\ No newline at end of file
diff --git a/Objects/complexobject.c b/Objects/complexobject.c
index 91e06a8..05cae32 100644
--- a/Objects/complexobject.c
+++ b/Objects/complexobject.c
@@ -514,8 +514,6 @@ static PyObject *
 complex_pow(PyObject *v, PyObject *w, PyObject *z)
 {
     Py_complex p;
-    Py_complex exponent;
-    long int_exponent;
     Py_complex a, b;
     TO_COMPLEX(v, a);
     TO_COMPLEX(w, b);
@@ -525,12 +523,21 @@ complex_pow(PyObject *v, PyObject *w, PyObject *z)
         return NULL;
     }
     errno = 0;
-    exponent = b;
-    int_exponent = (long)exponent.real;
-    if (exponent.imag == 0. && exponent.real == int_exponent)
-        p = c_powi(a, int_exponent);
-    else
-        p = _Py_c_pow(a, exponent);
+    // Check if w is an integer value that fits inside a C long, so we can
+    // use a faster algorithm. TO_COMPLEX(w, b), above, already handled the
+    // conversion from larger longs, as well as other types.
+    if (PyLong_Check(w)) {
+        int overflow = 0;
+        long int_exponent = PyLong_AsLongAndOverflow(w, &overflow);
+        if (int_exponent == -1 && PyErr_Occurred())
+            return NULL;
+        if (overflow == 0)
+            p = c_powi(a, int_exponent);
+        else
+            p = _Py_c_pow(a, b);
+    } else {
+        p = _Py_c_pow(a, b);
+    }
 
     Py_ADJUST_ERANGE2(p.real, p.imag);
     if (errno == EDOM) {