8000 bpo-43260: io: Prevent large data remains in textio buffer. (GH-24592) · python/cpython@d51436f · GitHub
[go: up one dir, main page]

Skip to content

Commit d51436f

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 (cherry picked from commit 01806d5)
1 parent 44fe320 commit d51436f

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
@@ -3747,6 +3747,33 @@ def test_del__CHUNK_SIZE_SystemError(self):
37473747
with self.assertRaises(AttributeError):
37483748
del t._CHUNK_SIZE
37493749

3750+
def test_internal_buffer_size(self):
3751+
# bpo-43260: TextIOWrapper's internal buffer should not store
3752+
# data larger than chunk size.
3753+
chunk_size = 8192 # default chunk size, updated later
3754+
3755+
class MockIO(self.MockRawIO):
3756+
def write(self, data):
3757+
if len(data) > chunk_size:
3758+
raise RuntimeError
3759+
return super().write(data)
3760+
3761+
buf = MockIO()
3762+
t = self.TextIOWrapper(buf, encoding="ascii")
3763+
chunk_size = t._CHUNK_SIZE
3764+
t.write("abc")
3765+
t.write("def")
3766+
# default chunk size is 8192 bytes so t don't write data to buf.
3767+
self.assertEqual([], buf._write_stack)
3768+
3769+
with self.assertRaises(RuntimeError):
3770+
t.write("x"*(chunk_size+1))
3771+
3772+
self.assertEqual([b"abcdef"], buf._write_stack)
3773+
t.write("ghi")
3774+
t.write("x"*chunk_size)
3775+
self.assertEqual([b"abcdef", b"ghi", b"x"*chunk_size], buf._write_stack)
3776+
37503777

37513778
class PyTextIOWrapperTest(TextIOWrapperTest):
37523779
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
@@ -1602,6 +1602,8 @@ _textiowrapper_writeflush(textio *self)
16021602
ret = PyObject_CallMethodOneArg(self->buffer, _PyIO_str_write, b);
16031603
} while (ret == NULL && _PyIO_trap_eintr());
16041604
Py_DECREF(b);
1605+
// NOTE: We cleared buffer but we don't know how many bytes are actually written
1606+
// when an error occurred.
16051607
if (ret == NULL)
16061608
return -1;
16071609
Py_DECREF(ret);
@@ -1659,7 +1661,10 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text)
16591661

16601662
/* XXX What if we were just reading? */
16611663
if (self->encodefunc != NULL) {
1662-
if (PyUnicode_IS_ASCII(text) && is_asciicompat_encoding(self->encodefunc)) {
1664+
if (PyUnicode_IS_ASCII(text) &&
1665+
// See bpo-43260
1666+
PyUnicode_GET_LENGTH(text) <= self->chunk_size &&
1667+
is_asciicompat_encoding(self->encodefunc)) {
16631668
b = text;
16641669
Py_INCREF(b);
16651670
}
@@ -1668,8 +1673,9 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text)
16681673
}
16691674
self->encoding_start_of_stream = 0;
16701675
}
1671-
else
1676+
else {
16721677
b = PyObject_CallMethodOneArg(self->encoder, _PyIO_str_encode, text);
1678+
}
16731679

16741680
Py_DECREF(text);
16751681
if (b == NULL)
@@ -1694,6 +1700,14 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text)
16941700
self->pending_bytes_count = 0;
16951701
self->pending_bytes = b;
16961702
}
1703+
else if (self->pending_bytes_count + bytes_len > self->chunk_size) {
1704+
// Prevent to concatenate more than chunk_size data.
1705+
if (_textiowrapper_writeflush(self) < 0) {
1706+
Py_DECREF(b);
1707+
return NULL;
1708+
}
1709+
self->pending_bytes = b;
1710+
}
16971711
else if (!PyList_CheckExact(self->pending_bytes)) {
16981712
PyObject *list = PyList_New(2);
16991713
if (list == NULL) {
@@ -1713,7 +1727,7 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text)
17131727
}
17141728

17151729
self->pending_bytes_count += bytes_len;
1716-
if (self->pending_bytes_count > self->chunk_size || needflush ||
1730+
if (self->pending_bytes_count >= self->chunk_size || needflush ||
17171731
text_needflush) {
17181732
if (_textiowrapper_writeflush(self) < 0)
17191733
return NULL;

0 commit comments

Comments
 (0)
0