Skip to content

Commit df0ee5a

Browse files
[3.14] gh-160324: Keep io.StringIO state after MemoryError (GH-160325) (#161042)
* gh-160324: Keep `io.StringIO` state after `MemoryError` (GH-160325) Allocate before consuming the accumulated writes in realize() and make_intermediate(), and reset the size when the contents are lost, so that later calls do not index outside the internal buffer. Shrink the buffer of the accumulating writer before finishing it, so that getvalue(), read() and the other methods leave the contents unchanged when they fail with MemoryError instead of dropping them. ANT-2026-MPHGX2FL (cherry picked from commit 83b047e) Co-authored-by: Gregory P. Smith <68491+gpshead@users.noreply.github.com> * Avoid _PyUnicodeWriter_SetBuffer() and _PyUnicode_ResizeCompact() They do not exist in 3.14. Use PyUnicode_Resize(), which leaves the buffer unchanged on failure, and update the writer fields directly. * Do not copy the intermediate string in make_intermediate() In 3.14, writing a string to a new PyUnicodeWriter copies it, so a MemoryError there lost the contents of the StringIO. Disable overallocation for this write so the writer shares the string (copy-on-write) instead, as it does in main. --------- Co-authored-by: Gregory P. Smith <68491+gpshead@users.noreply.github.com> Co-authored-by: Gregory P. Smith <greg@krypto.org>
1 parent 576f556 commit df0ee5a

3 files changed

Lines changed: 98 additions & 18 deletions

File tree

‎Lib/test/test_memoryio.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1067,6 +1067,39 @@ def __str__(self):
10671067
memio2.write(MyStr("world"))
10681068
self.assertEqual(memio2.getvalue(), "hello world")
10691069

1070+
@support.nomemtest
1071+
def test_memory_error_consistent_state(self):
1072+
# A MemoryError while copying the accumulated writes into the
1073+
# internal buffer must leave the size of the stream consistent with
1074+
# the buffer, so that using the object again does not read or write
1075+
# outside it.
1076+
piece = "0123456789" * 10
1077+
text = piece * 50
1078+
calls = [("readline",), ("read", 5), ("read",), ("getvalue",),
1079+
("truncate", 3), ("write", "ab"), ("__next__",)]
1080+
for name, *args in calls:
1081+
for n in range(3):
1082+
with self.subTest(method=name, n=n):
1083+
memio = self.ioclass()
1084+
for _ in range(50):
1085+
memio.write(piece)
1086+
memio.seek(0)
1087+
method = getattr(memio, name)
1088+
try:
1089+
with support.inject_memory_error_cm(n, n + 1):
1090+
method(*args)
1091+
except MemoryError:
1092+
pass
1093+
else:
1094+
continue
1095+
# The contents are unchanged.
1096+
self.assertEqual(memio.getvalue(), text)
1097+
self.assertEqual(memio.seek(0, 2), len(text))
1098+
memio.seek(0)
1099+
self.assertEqual(memio.readline(), text)
1100+
memio.write("!")
1101+
self.assertEqual(memio.getvalue(), text + "!")
1102+
10701103

10711104
class CStringIOPickleTest(PyStringIOPickleTest):
10721105
UnsupportedOperation = io.UnsupportedOperation
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Fix a crash or garbage data when an :class:`io.StringIO` object is used again
2+
after one of its methods failed with :exc:`MemoryError`. The object now keeps
3+
its contents.

‎Modules/_io/stringio.c‎

Lines changed: 62 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -129,25 +129,60 @@ resize_buffer(stringio *self, size_t size)
129129
return -1;
130130
}
131131

132+
/* Shrink the buffer of the writer to its contents, so that the following
133+
PyUnicodeWriter_Finish() has nothing to allocate and cannot fail.
134+
If this fails, the writer is unchanged. */
135+
static int
136+
shrink_writer(stringio *self)
137+
{
138+
_PyUnicodeWriter *writer = (_PyUnicodeWriter *)self->writer;
139+
if (writer->readonly || writer->buffer == NULL
140+
|| writer->pos == 0 || writer->pos == writer->size)
141+
{
142+
return 0;
143+
}
144+
/* On failure, writer->buffer is left unchanged. */
145+
if (PyUnicode_Resize(&writer->buffer, writer->pos) < 0) {
146+
return -1;
147+
}
148+
PyObject *buffer = writer->buffer;
149+
writer->data = PyUnicode_DATA(buffer);
150+
writer->kind = PyUnicode_KIND(buffer);
151+
writer->maxchar = PyUnicode_MAX_CHAR_VALUE(buffer);
152+
writer->size = PyUnicode_GET_LENGTH(buffer);
153+
return 0;
154+
}
155+
132156
static PyObject *
133157
make_intermediate(stringio *self)
134158
{
135-
PyObject *intermediate = PyUnicodeWriter_Finish(self->writer);
136-
self->writer = NULL;
137-
self->state = STATE_REALIZED;
138-
if (intermediate == NULL)
159+
/* Create the new writer and shrink the old one first: if this fails,
160+
the object keeps its contents. */
161+
PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);
162+
if (writer == NULL)
139163
return NULL;
164+
if (shrink_writer(self) < 0) {
165+
PyUnicodeWriter_Discard(writer);
166+
return NULL;
167+
}
140168

141-
self->writer = PyUnicodeWriter_Create(0);
142-
if (self->writer == NULL) {
143-
Py_DECREF(intermediate);
169+
PyObject *intermediate = PyUnicodeWriter_Finish(self->writer);
170+
self->writer = writer;
171+
if (intermediate == NULL) {
172+
/* The contents are lost, the new writer is empty. */
173+
self->string_size = 0;
144174
return NULL;
145175
}
146-
if (PyUnicodeWriter_WriteStr(self->writer, intermediate)) {
176+
/* Without overallocation, writing to an empty writer shares the string
177+
(copy-on-write) instead of copying it, so it does not allocate. */
178+
((_PyUnicodeWriter *)writer)->overallocate = 0;
179+
int res = PyUnicodeWriter_WriteStr(writer, intermediate);
180+
((_PyUnicodeWriter *)writer)->overallocate = 1;
181+
if (res < 0) {
147182
Py_DECREF(intermediate);
183+
self->string_size = 0;
148184
return NULL;
149185
}
150-
self->state = STATE_ACCUMULATING;
151186
return intermediate;
152187
}
153188

@@ -160,23 +195,30 @@ realize(stringio *self)
160195
if (self->state == STATE_REALIZED)
161196
return 0;
162197
assert(self->state == STATE_ACCUMULATING);
198+
199+
/* Allocate the internal buffer and shrink the writer first: if this
200+
fails, the object stays in the accumulating state and keeps its
201+
contents. */
202+
if (resize_buffer(self, self->string_size) < 0)
203+
return -1;
204+
if (shrink_writer(self) < 0)
205+
return -1;
163206
self->state = STATE_REALIZED;
164207

165208
intermediate = PyUnicodeWriter_Finish(self->writer);
166209
self->writer = NULL;
167-
if (intermediate == NULL)
210+
if (intermediate == NULL) {
211+
/* The contents are lost, the internal buffer contains nothing. */
212+
self->string_size = 0;
168213
return -1;
214+
}
169215

170-
/* Append the intermediate string to the internal buffer.
171-
The length should be equal to the current cursor position.
172-
*/
216+
/* Copy the intermediate string to the internal buffer. */
173217
len = PyUnicode_GET_LENGTH(intermediate);
174-
if (resize_buffer(self, len) < 0) {
175-
Py_DECREF(intermediate);
176-
return -1;
177-
}
218+
assert(len == self->string_size);
178219
if (!PyUnicode_AsUCS4(intermediate, self->buf, len, 0)) {
179220
Py_DECREF(intermediate);
221+
self->string_size = 0;
180222
return -1;
181223
}
182224

@@ -344,7 +386,9 @@ _io_StringIO_read_impl(stringio *self, Py_ssize_t size)
344386
/* Optimization for seek(0); read() */
345387
if (self->state == STATE_ACCUMULATING && self->pos == 0 && size == n) {
346388
PyObject *result = make_intermediate(self);
347-
self->pos = self->string_size;
389+
if (result != NULL) {
390+
self->pos = self->string_size;
391+
}
348392
return result;
349393
}
350394

0 commit comments

Comments
 (0)