[3.15] gh-158585: Fix _PyBytes_Resize() if hash value is already computed (#159555)

* gh-158585: Fix _PyBytes_Resize() if hash value is already computed (#158589)

_PyBytes_Resize() now creates a new bytes object if the hash value
was already computed.

bytes_resize_inplace() no longer sets the hash value to -1, since
this function must not be called if the hash value was already
computed.

Add tests on PyUnicode_Resize().

(cherry picked from commit 6e7738b)

* gh-158585: Fix refleak in _testcapi.bytes_resize() (#158806)

(cherry picked from commit 587b7a5)

Author: vstinner

Lib/test/test_capi/test_bytes.py

@@ -237,39 +237,60 @@ def test_resize(self):
         """Test _PyBytes_Resize()"""
         _resize = _testcapi.bytes_resize
 
-        def resize(obj, size, new):
-            result = _resize(obj, size, new)
-            if 1 <= len(result):
-                if new or size != len(obj):
-                    # gh-156995: Make sure that the result is a fresh object.
-                    # Previously, _PyBytes_Resize(&obj, 1) returned a singleton
-                    # if _PyObject_IsUniquelyReferenced() is false.
-                    self.assertEqual(sys.getrefcount(result), 1)
-                    self.assertFalse(sys._is_immortal(result))
-            else:
+        def assert_is_fresh_copy(result, refcnt, is_new_obj):
+            self.assertEqual(refcnt, 1)
+            self.assertTrue(is_new_obj)
+            self.assertFalse(sys._is_immortal(result))
+
+        def resize(obj, size, new, compute_hash=False):
+            old_size = len(obj)
+            result, refcnt, is_new_obj = _resize(obj, size, new,
+                                                 compute_hash=compute_hash)
+
+            if size == old_size:
+                # Return the same object unchanged
+                self.assertFalse(is_new_obj)
+            elif old_size == 0:
+                assert_is_fresh_copy(result, refcnt, is_new_obj)
+            elif size == 0:
                 # check that the result is the empty bytes string singleton
+                self.assertEqual(result, b'')
                 self.assertTrue(sys._is_immortal(result))
+                self.assertTrue(is_new_obj)
+            elif (not new) or compute_hash:
+                # gh-156995: Make sure that the result is a fresh object.
+                # Previously, _PyBytes_Resize(&obj, 1) returned a singleton
+                # if _PyObject_IsUniquelyReferenced() is false.
+                assert_is_fresh_copy(result, refcnt, is_new_obj)
+            else:
+                # An in-place resize can return the same memory address, or
+                # not. 'is_new_obj' cannot be tested.
+                self.assertEqual(refcnt, 1)
+                self.assertFalse(sys._is_immortal(result))
+
             return result
 
         for new in True, False:
-            with self.subTest(new=new):
-                self.assertEqual(resize(b'abc', 0, new), b'')
-                self.assertEqual(resize(b'abc', 1, new), b'a')
-                self.assertEqual(resize(b'abc', 2, new), b'ab')
-                self.assertEqual(resize(b'abc', 3, new), b'abc')
-                b = resize(b'abc', 4, new)
-                self.assertEqual(len(b), 4)
-                self.assertEqual(b[:3], b'abc')
-
-                self.assertEqual(resize(b'a', 0, new), b'')
-                self.assertEqual(resize(b'a', 1, new), b'a')
-                b = resize(b'a', 2, new)
-                self.assertEqual(len(b), 2)
-                self.assertEqual(b[:1], b'a')
-
-                self.assertEqual(resize(b'', 0, new), b'')
-                self.assertEqual(len(resize(b'', 1, new)), 1)
-                self.assertEqual(len(resize(b'', 2, new)), 2)
+            for compute_hash in True, False:
+                with self.subTest(new=new, compute_hash=compute_hash):
+                    self.assertEqual(resize(b'abc', 0, new, compute_hash), b'')
+                    self.assertEqual(resize(b'abc', 1, new, compute_hash), b'a')
+                    self.assertEqual(resize(b'abc', 2, new, compute_hash), b'ab')
+                    self.assertEqual(resize(b'abc', 3, new, compute_hash), b'abc')
+
+                    b = resize(b'abc', 4, new, compute_hash)
+                    self.assertEqual(len(b), 4)
+                    self.assertEqual(b[:3], b'abc')
+
+                    self.assertEqual(resize(b'a', 0, new, compute_hash), b'')
+                    self.assertEqual(resize(b'a', 1, new, compute_hash), b'a')
+                    b = resize(b'a', 2, new, compute_hash)
+                    self.assertEqual(len(b), 2)
+                    self.assertEqual(b[:1], b'a')
+
+                    self.assertEqual(resize(b'', 0, new, compute_hash), b'')
+                    self.assertEqual(len(resize(b'', 1, new, compute_hash)), 1)
+                    self.assertEqual(len(resize(b'', 2, new, compute_hash)), 2)
 
         self.assertRaises(SystemError, resize, b'abc', -1, False)
         self.assertRaises(SystemError, resize, bytearray(b'abc'), 3, False)

Lib/test/test_capi/test_unicode.py

@@ -168,28 +168,75 @@ def test_write_macro(self):
 
     def test_resize(self):
         """Test PyUnicode_Resize()"""
-        resize = _testlimitedcapi.unicode_resize
+        _resize = _testlimitedcapi.unicode_resize
+        resize_null = _testlimitedcapi.unicode_resize_null
+
+        def resize(s, length, new=True, compute_hash=False):
+            if s is not NULL and isinstance(s, str):
+                old_length = len(s)
+            else:
+                old_length = 0
+            result, int_result, refcnt, is_new_obj = _resize(s, length,
+                                                             new, compute_hash)
+            self.assertEqual(int_result, 0)
+
+            if length == old_length:
+                # Return the same object unchanged
+                self.assertFalse(is_new_obj)
+            elif length == 0:
+                # Get the empty Unicode string
+                self.assertEqual(result, '')
+                self.assertTrue(sys._is_immortal(result))
+                self.assertTrue(is_new_obj)
+            elif (not new) or compute_hash:
+                # Get a fresh copy
+                self.assertEqual(refcnt, 1)
+                self.assertTrue(is_new_obj)
+                self.assertFalse(sys._is_immortal(result))
+            else:
+                # In-size replace can return the same address, or not.
+                # So 'is_new_obj' cannot be tested.
+                self.assertFalse(sys._is_immortal(result))
+
+            return result
 
         strings = [
             # all strings have exactly 3 characters
             'abc', '\xa1\xa2\xa3', '\u4f60\u597d\u4e16',
             '\U0001f600\U0001f601\U0001f602'
         ]
-        for s in strings:
-            self.assertEqual(resize(s, 3), (s, 0))
-            self.assertEqual(resize(s, 2), (s[:2], 0))
-            self.assertEqual(resize(s, 4), (s + '\0', 0))
-            self.assertEqual(resize(s, 10), (s + '\0'*7, 0))
-            self.assertEqual(resize(s, 0), ('', 0))
-            self.assertRaises(MemoryError, resize, s, PY_SSIZE_T_MAX)
-            self.assertRaises(SystemError, resize, s, -1)
-            self.assertRaises(SystemError, resize, s, PY_SSIZE_T_MIN)
+        for new in (True, False):
+            for compute_hash in (True, False):
+                for s in strings:
+                    with self.subTest(new=new, compute_hash=compute_hash, s=s):
+                        self.assertEqual(resize(s, 3, new, compute_hash),
+                                         s)
+                        self.assertEqual(resize(s, 2, new, compute_hash),
+                                         s[:2])
+                        self.assertEqual(resize(s, 4, new, compute_hash),
+                                         s + '\0')
+                        self.assertEqual(resize(s, 10, new, compute_hash),
+                                         s + '\0'*7)
+                        self.assertEqual(resize(s, 0, new, compute_hash),
+                                         '')
+
+                        with self.assertRaises(MemoryError):
+                            resize(s, PY_SSIZE_T_MAX, new, compute_hash)
+                        with self.assertRaises(SystemError):
+                            resize(s, -1, new, compute_hash)
+                        with self.assertRaises(SystemError):
+                            resize(s, PY_SSIZE_T_MIN, new, compute_hash)
+
         self.assertRaises(SystemError, resize, b'abc', 0)
         self.assertRaises(SystemError, resize, [], 0)
         self.assertRaises(SystemError, resize, NULL, 0)
         # TODO: Test PyUnicode_Resize() with non-modifiable and legacy unicode
         # and with NULL as the address.
 
+        # Test PyUnicode_Resize(NULL, length)
+        self.assertRaises(SystemError, resize_null, 0)
+        self.assertRaises(SystemError, resize_null, 123)
+
     def test_append(self):
         """Test PyUnicode_Append()"""
         append = _testlimitedcapi.unicode_append

Modules/_testcapi/bytes.c

@@ -11,14 +11,19 @@
 
 /* Test _PyBytes_Resize() */
 static PyObject *
-bytes_resize(PyObject *Py_UNUSED(module), PyObject *args)
+bytes_resize(PyObject *Py_UNUSED(module), PyObject *args, PyObject *kwargs)
 {
+    static char *kwlist[] = {"obj", "newsize", "new", "compute_hash", NULL};
     PyObject *obj;
     Py_ssize_t newsize;
     int new;
+    int compute_hash = 0;
 
-    if (!PyArg_ParseTuple(args, "Onp", &obj, &newsize, &new))
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs,
+                                     "Onp|p", kwlist,
+                                     &obj, &newsize, &new, &compute_hash)) {
         return NULL;
+    }
 
     NULLABLE(obj);
     if (new) {
@@ -34,13 +39,25 @@ bytes_resize(PyObject *Py_UNUSED(module), PyObject *args)
     else {
         Py_XINCREF(obj);
     }
+
+    if (compute_hash) {
+        if (PyObject_Hash(obj) == -1) {
+            Py_DECREF(obj);
+            return NULL;
+        }
+    }
+
+    PyObject *old_obj = obj;
     if (_PyBytes_Resize(&obj, newsize) < 0) {
         assert(obj == NULL);
+        return NULL;
     }
     else {
         assert(obj != NULL);
     }
-    return obj;
+
+    Py_ssize_t refcnt = Py_REFCNT(obj);
+    return Py_BuildValue("Nnp", obj, refcnt, obj != old_obj);
 }
 
 
@@ -526,7 +543,7 @@ test_byteswriter_ptr(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
 
 
 static PyMethodDef test_methods[] = {
-    {"bytes_resize", bytes_resize, METH_VARARGS},
+    {"bytes_resize", _PyCFunction_CAST(bytes_resize), METH_VARARGS | METH_KEYWORDS},
     {"bytes_join", bytes_join, METH_VARARGS},
     {"byteswriter_abc", byteswriter_abc, METH_NOARGS},
     {"byteswriter_resize", byteswriter_resize, METH_NOARGS},

Modules/_testlimitedcapi/unicode.c

@@ -183,24 +183,59 @@ unicode_resize(PyObject *self, PyObject *args)
     PyObject *obj, *copy;
     Py_ssize_t length;
     int result;
+    int new = 1;
+    int compute_hash = 0;
 
-    if (!PyArg_ParseTuple(args, "On", &obj, &length)) {
+    if (!PyArg_ParseTuple(args, "On|pp", &obj, &length, &new, &compute_hash)) {
         return NULL;
     }
 
     NULLABLE(obj);
-    if (!(copy = unicode_copy(obj)) && obj) {
-        return NULL;
+    Py_ssize_t old_len = obj ? PyUnicode_GetLength(obj) : 0;
+    if (obj != NULL && new) {
+        copy = unicode_copy(obj);
+        if (!copy) {
+            return NULL;
+        }
+        obj = copy;
     }
-    result = PyUnicode_Resize(&copy, length);
+    else {
+        Py_XINCREF(obj);
+    }
+
+    if (compute_hash && obj != NULL) {
+        if (PyObject_Hash(obj) == -1) {
+            Py_DECREF(obj);
+            return NULL;
+        }
+    }
+
+    PyObject *old_obj = obj;
+    result = PyUnicode_Resize(&obj, length);
     if (result == -1 && PyErr_Occurred()) {
-        Py_XDECREF(copy);
+        Py_XDECREF(obj);
         return NULL;
     }
-    if (obj && PyUnicode_Check(obj) && length > PyUnicode_GetLength(obj)) {
-        unicode_fill(copy, PyUnicode_GetLength(obj), length, 0U);
+    if (obj && PyUnicode_Check(obj) && length > old_len) {
+        unicode_fill(obj, old_len, length, 0U);
     }
-    return Py_BuildValue("(Ni)", copy, result);
+    Py_ssize_t refcnt = Py_REFCNT(obj);
+    return Py_BuildValue("(Ninp)", obj, result, refcnt, obj != old_obj);
+}
+
+// Test PyUnicode_Resize(NULL, length)
+static PyObject *
+unicode_resize_null(PyObject *self, PyObject *args)
+{
+    Py_ssize_t length;
+    if (!PyArg_ParseTuple(args, "n", &length)) {
+        return NULL;
+    }
+
+    if (PyUnicode_Resize(NULL, length)) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
 }
 
 /* Test PyUnicode_Append() */
@@ -1882,6 +1917,7 @@ static PyMethodDef TestMethods[] = {
     {"test_widechar",            test_widechar,                  METH_NOARGS},
     {"unicode_writechar",        unicode_writechar,              METH_VARARGS},
     {"unicode_resize",           unicode_resize,                 METH_VARARGS},
+    {"unicode_resize_null",      unicode_resize_null,            METH_VARARGS},
     {"unicode_append",           unicode_append,                 METH_VARARGS},
     {"unicode_appendanddel",     unicode_appendanddel,           METH_VARARGS},
     {"unicode_fromstringandsize",unicode_fromstringandsize,      METH_VARARGS},