Skip to content

Commit 83b047e

Browse files
authored
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
1 parent e80962a commit 83b047e

3 files changed

Lines changed: 90 additions & 17 deletions

File tree

‎Lib/test/test_io/test_memoryio.py‎

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

1240+
@support.nomemtest
1241+
def test_memory_error_consistent_state(self):
1242+
# A MemoryError while copying the accumulated writes into the
1243+
# internal buffer must leave the size of the stream consistent with
1244+
# the buffer, so that using the object again does not read or write
1245+
# outside it.
1246+
piece = "0123456789" * 10
1247+
text = piece * 50
1248+
calls = [("readline",), ("read", 5), ("read",), ("getvalue",),
1249+
("truncate", 3), ("write", "ab"), ("__next__",)]
1250+
for name, *args in calls:
1251+
for n in range(3):
1252+
with self.subTest(method=name, n=n):
1253+
memio = self.ioclass()
1254+
for _ in range(50):
1255+
memio.write(piece)
1256+
memio.seek(0)
1257+
method = getattr(memio, name)
1258+
try:
1259+
with support.inject_memory_error_cm(n, n + 1):
1260+
method(*args)
1261+
except MemoryError:
1262+
pass
1263+
else:
1264+
continue
1265+
# The contents are unchanged.
1266+
self.assertEqual(memio.getvalue(), text)
1267+
self.assertEqual(memio.seek(0, 2), len(text))
1268+
memio.seek(0)
1269+
self.assertEqual(memio.readline(), text)
1270+
memio.write("!")
1271+
self.assertEqual(memio.getvalue(), text + "!")
1272+
12401273

12411274
class CStringIOPickleTest(PyStringIOPickleTest):
12421275
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: 54 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#include "Python.h"
22
#include <stddef.h> // offsetof()
33
#include "pycore_object.h"
4+
#include "pycore_unicodeobject.h" // _PyUnicode_ResizeCompact()
45
#include "pycore_weakref.h" // FT_CLEAR_WEAKREFS()
56
#include "_iomodule.h"
67

@@ -143,25 +144,52 @@ resize_buffer(stringio *self, size_t size)
143144
return -1;
144145
}
145146

147+
/* Shrink the buffer of the writer to its contents, so that the following
148+
PyUnicodeWriter_Finish() has nothing to allocate and cannot fail.
149+
If this fails, the writer is unchanged. */
150+
static int
151+
shrink_writer(stringio *self)
152+
{
153+
_PyUnicodeWriter *writer = (_PyUnicodeWriter *)self->writer;
154+
if (writer->readonly || writer->buffer == NULL
155+
|| writer->pos == 0 || writer->pos == writer->size)
156+
{
157+
return 0;
158+
}
159+
PyObject *buffer = _PyUnicode_ResizeCompact(writer->buffer, writer->pos);
160+
if (buffer == NULL) {
161+
return -1;
162+
}
163+
writer->buffer = NULL;
164+
_PyUnicodeWriter_SetBuffer(writer, buffer);
165+
return 0;
166+
}
167+
146168
static PyObject *
147169
make_intermediate(stringio *self)
148170
{
149-
PyObject *intermediate = PyUnicodeWriter_Finish(self->writer);
150-
self->writer = NULL;
151-
self->state = STATE_REALIZED;
152-
if (intermediate == NULL)
171+
/* Create the new writer and shrink the old one first: if this fails,
172+
the object keeps its contents. */
173+
PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);
174+
if (writer == NULL)
153175
return NULL;
176+
if (shrink_writer(self) < 0) {
177+
PyUnicodeWriter_Discard(writer);
178+
return NULL;
179+
}
154180

155-
self->writer = PyUnicodeWriter_Create(0);
156-
if (self->writer == NULL) {
157-
Py_DECREF(intermediate);
181+
PyObject *intermediate = PyUnicodeWriter_Finish(self->writer);
182+
self->writer = writer;
183+
if (intermediate == NULL) {
184+
/* The contents are lost, the new writer is empty. */
185+
self->string_size = 0;
158186
return NULL;
159187
}
160188
if (PyUnicodeWriter_WriteStr(self->writer, intermediate)) {
161189
Py_DECREF(intermediate);
190+
self->string_size = 0;
162191
return NULL;
163192
}
164-
self->state = STATE_ACCUMULATING;
165193
return intermediate;
166194
}
167195

@@ -174,23 +202,30 @@ realize(stringio *self)
174202
if (self->state == STATE_REALIZED)
175203
return 0;
176204
assert(self->state == STATE_ACCUMULATING);
205+
206+
/* Allocate the internal buffer and shrink the writer first: if this
207+
fails, the object stays in the accumulating state and keeps its
208+
contents. */
209+
if (resize_buffer(self, self->string_size) < 0)
210+
return -1;
211+
if (shrink_writer(self) < 0)
212+
return -1;
177213
self->state = STATE_REALIZED;
178214

179215
intermediate = PyUnicodeWriter_Finish(self->writer);
180216
self->writer = NULL;
181-
if (intermediate == NULL)
217+
if (intermediate == NULL) {
218+
/* The contents are lost, the internal buffer contains nothing. */
219+
self->string_size = 0;
182220
return -1;
221+
}
183222

184-
/* Append the intermediate string to the internal buffer.
185-
The length should be equal to the current cursor position.
186-
*/
223+
/* Copy the intermediate string to the internal buffer. */
187224
len = PyUnicode_GET_LENGTH(intermediate);
188-
if (resize_buffer(self, len) < 0) {
189-
Py_DECREF(intermediate);
190-
return -1;
191-
}
225+
assert(len == self->string_size);
192226
if (!PyUnicode_AsUCS4(intermediate, self->buf, len, 0)) {
193227
Py_DECREF(intermediate);
228+
self->string_size = 0;
194229
return -1;
195230
}
196231

@@ -358,7 +393,9 @@ _io_StringIO_read_impl(stringio *self, Py_ssize_t size)
358393
/* Optimization for seek(0); read() */
359394
if (self->state == STATE_ACCUMULATING && self->pos == 0 && size == n) {
360395
PyObject *result = make_intermediate(self);
361-
self->pos = self->string_size;
396+
if (result != NULL) {
397+
self->pos = self->string_size;
398+
}
362399
return result;
363400
}
364401

0 commit comments

Comments
 (0)