gh-112625: Protect bytearray from being freed by misbehaving iterator inside bytearray.join (GH-112626)

chilaxan · web-flow · commit 0e732d0997cf · 2023-12-04T08:15:43.000Z
diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
@@ -2039,6 +2039,23 @@ def test_bytearray_extend_error(self):
         bad_iter = map(int, "X")
         self.assertRaises(ValueError, array.extend, bad_iter)
 
+    def test_bytearray_join_with_misbehaving_iterator(self):
+        # Issue #112625
+        array = bytearray(b',')
+        def iterator():
+            array.clear()
+            yield b'A'
+            yield b'B'
+        self.assertRaises(BufferError, array.join, iterator())
+
+    def test_bytearray_join_with_custom_iterator(self):
+        # Issue #112625
+        array = bytearray(b',')
+        def iterator():
+            yield b'A'
+            yield b'B'
+        self.assertEqual(bytearray(b'A,B'), array.join(iterator()))
+
     def test_construct_singletons(self):
         for const in None, Ellipsis, NotImplemented:
             tp = type(const)
diff --git a/Misc/NEWS.d/next/Core and Builtins/2023-12-03-19-34-51.gh-issue-112625.QWTlwS.rst b/Misc/NEWS.d/next/Core and Builtins/2023-12-03-19-34-51.gh-issue-112625.QWTlwS.rst
@@ -0,0 +1 @@
+Fixes a bug where a bytearray object could be cleared while iterating over an argument in the ``bytearray.join()`` method that could result in reading memory after it was freed.
diff --git a/Objects/bytearrayobject.c b/Objects/bytearrayobject.c
@@ -2007,7 +2007,10 @@ static PyObject *
 bytearray_join(PyByteArrayObject *self, PyObject *iterable_of_bytes)
 /*[clinic end generated code: output=a8516370bf68ae08 input=aba6b1f9b30fcb8e]*/
 {
-    return stringlib_bytes_join((PyObject*)self, iterable_of_bytes);
+    self->ob_exports++; // this protects `self` from being cleared/resized if `iterable_of_bytes` is a custom iterator
+    PyObject* ret = stringlib_bytes_join((PyObject*)self, iterable_of_bytes);
+    self->ob_exports--; // unexport `self`
+    return ret;
 }
 
 /*[clinic input]

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+Fixes a bug where a bytearray object could be cleared while iterating over an argument in the ``bytearray.join()`` method that could result in reading memory after it was freed.
Original file line number	Diff line number	Diff line change
`@@ -2007,7 +2007,10 @@ static PyObject *`
`2007`	`2007`	`bytearray_join(PyByteArrayObject self, PyObject iterable_of_bytes)`
`2008`	`2008`	`/[clinic end generated code: output=a8516370bf68ae08 input=aba6b1f9b30fcb8e]/`
`2009`	`2009`	`{`
`2010`		`- return stringlib_bytes_join((PyObject*)self, iterable_of_bytes);`
	`2010`	+ self->ob_exports++; // this protects `self` from being cleared/resized if `iterable_of_bytes` is a custom iterator
	`2011`	`+ PyObject* ret = stringlib_bytes_join((PyObject*)self, iterable_of_bytes);`
	`2012`	+ self->ob_exports--; // unexport `self`
	`2013`	`+ return ret;`
`2011`	`2014`	`}`
`2012`	`2015`
`2013`	`2016`	`/*[clinic input]`