Skip to content

Commit 01806d5

Browse files
authored
bpo-43260: io: Prevent large data remains in textio buffer. (GH-24592)
When very large data remains in TextIOWrapper, flush() may fail forever. So prevent that data larger than chunk_size is remained in TextIOWrapper internal buffer. Co-Authored-By: Eryk Sun
1 parent 84f7afe commit 01806d5

File tree

3 files changed

+46
-3
lines changed

3 files changed

+46
-3
lines changed

Lib/test/test_io.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3767,6 +3767,33 @@ def test_del__CHUNK_SIZE_SystemError(self):
37673767
with self.assertRaises(AttributeError):
37683768
del t._CHUNK_SIZE
37693769

3770+
def test_internal_buffer_size(self):
3771+
# bpo-43260: TextIOWrapper's internal buffer should not store
3772+
# data larger than chunk size.
3773+
chunk_size = 8192 # default chunk size, updated later
3774+
3775+
class MockIO(self.MockRawIO):
3776+
def write(self, data):
3777+
if len(data) > chunk_size:
3778+
raise RuntimeError
3779+
return super().write(data)
3780+
3781+
buf = MockIO()
3782+
t = self.TextIOWrapper(buf, encoding="ascii")
3783+
chunk_size = t._CHUNK_SIZE
3784+
t.write("abc")
3785+
t.write("def")
3786+
# default chunk size is 8192 bytes so t don't write data to buf.
3787+
self.assertEqual([], buf._write_stack)
3788+
3789+
with self.assertRaises(RuntimeError):
3790+
t.write("x"*(chunk_size+1))
3791+
3792+
self.assertEqual([b"abcdef"], buf._write_stack)
3793+
t.write("ghi")
3794+
t.write("x"*chunk_size)
3795+
self.assertEqual([b"abcdef", b"ghi", b"x"*chunk_size], buf._write_stack)
3796+
37703797

37713798
class PyTextIOWrapperTest(TextIOWrapperTest):
37723799
io = pyio
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix TextIOWrapper can not flush internal buffer forever after very large
2+
text is written.

Modules/_io/textio.c

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1585,6 +1585,8 @@ _textiowrapper_writeflush(textio *self)
15851585
ret = PyObject_CallMethodOneArg(self->buffer, _PyIO_str_write, b);
15861586
} while (ret == NULL && _PyIO_trap_eintr());
15871587
Py_DECREF(b);
1588+
// NOTE: We cleared buffer but we don't know how many bytes are actually written
1589+
// when an error occurred.
15881590
if (ret == NULL)
15891591
return -1;
15901592
Py_DECREF(ret);
@@ -1642,7 +1644,10 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text)
16421644

16431645
/* XXX What if we were just reading? */
16441646
if (self->encodefunc != NULL) {
1645-
if (PyUnicode_IS_ASCII(text) && is_asciicompat_encoding(self->encodefunc)) {
1647+
if (PyUnicode_IS_ASCII(text) &&
1648+
// See bpo-43260
1649+
PyUnicode_GET_LENGTH(text) <= self->chunk_size &&
1650+
is_asciicompat_encoding(self->encodefunc)) {
16461651
b = text;
16471652
Py_INCREF(b);
16481653
}
@@ -1651,8 +1656,9 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text)
16511656
}
16521657
self->encoding_start_of_stream = 0;
16531658
}
1654-
else
1659+
else {
16551660
b = PyObject_CallMethodOneArg(self->encoder, _PyIO_str_encode, text);
1661+
}
16561662

16571663
Py_DECREF(text);
16581664
if (b == NULL)
@@ -1677,6 +1683,14 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text)
16771683
self->pending_bytes_count = 0;
16781684
self->pending_bytes = b;
16791685
}
1686+
else if (self->pending_bytes_count + bytes_len > self->chunk_size) {
1687+
// Prevent to concatenate more than chunk_size data.
1688+
if (_textiowrapper_writeflush(self) < 0) {
1689+
Py_DECREF(b);
1690+
return NULL;
1691+
}
1692+
self->pending_bytes = b;
1693+
}
16801694
else if (!PyList_CheckExact(self->pending_bytes)) {
16811695
PyObject *list = PyList_New(2);
16821696
if (list == NULL) {
@@ -1696,7 +1710,7 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text)
16961710
}
16971711

16981712
self->pending_bytes_count += bytes_len;
1699-
if (self->pending_bytes_count > self->chunk_size || needflush ||
1713+
if (self->pending_bytes_count >= self->chunk_size || needflush ||
17001714
text_needflush) {
17011715
if (_textiowrapper_writeflush(self) < 0)
17021716
return NULL;

0 commit comments

Comments
 (0)