extmod/moddeflate: Keep DeflateIO state consistent on window alloc fail.

Allocation of a large compression window may fail, and in that case keep
the `DeflateIO` state consistent so its other methods (such as `close()`)
still work.  Consistency is kept by only updating the `self->write` member
if the window allocation succeeds.

Thanks to @jimmo for finding the bug.

Signed-off-by: Damien George <damien@micropython.org>

Author: dpgeorge

extmod/moddeflate.c

@@ -168,16 +168,20 @@ static bool deflateio_init_write(mp_obj_deflateio_t *self) {
 
     const mp_stream_p_t *stream = mp_get_stream_raise(self->stream, MP_STREAM_OP_WRITE);
 
-    self->write = m_new_obj(mp_obj_deflateio_write_t);
-    self->write->input_len = 0;
-
     int wbits = self->window_bits;
     if (wbits == 0) {
         // Same default wbits for all formats.
         wbits = DEFLATEIO_DEFAULT_WBITS;
     }
+
+    // Allocate the large window before allocating the mp_obj_deflateio_write_t, in case the
+    // window allocation fails the mp_obj_deflateio_t object will remain in a consistent state.
     size_t window_len = 1 << wbits;
-    self->write->window = m_new(uint8_t, window_len);
+    uint8_t *window = m_new(uint8_t, window_len);
+
+    self->write = m_new_obj(mp_obj_deflateio_write_t);
+    self->write->window = window;
+    self->write->input_len = 0;
 
     uzlib_lz77_init(&self->write->lz77, self->write->window, window_len);
     self->write->lz77.dest_write_data = self;

tests/extmod/deflate_compress_memory_error.py

@@ -0,0 +1,39 @@
+# Test deflate.DeflateIO compression, with out-of-memory errors.
+
+try:
+    # Check if deflate is available.
+    import deflate
+    import io
+except ImportError:
+    print("SKIP")
+    raise SystemExit
+
+# Check if compression is enabled.
+if not hasattr(deflate.DeflateIO, "write"):
+    print("SKIP")
+    raise SystemExit
+
+# Create a compressor object.
+b = io.BytesIO()
+g = deflate.DeflateIO(b, deflate.RAW, 15)
+
+# Then, use up most of the heap.
+l = []
+while True:
+    try:
+        l.append(bytearray(1000))
+    except:
+        break
+l.pop()
+
+# Try to compress.  This will try to allocate a large window and fail.
+try:
+    g.write('test')
+except MemoryError:
+    print("MemoryError")
+
+# Should still be able to close the stream.
+g.close()
+
+# The underlying output stream should be unchanged.
+print(b.getvalue())

tests/extmod/deflate_compress_memory_error.py.exp

@@ -0,0 +1,2 @@
+MemoryError
+b''

tests/run-tests.py

@@ -163,6 +163,7 @@ def open(self, path, mode):
         "extmod/asyncio_threadsafeflag.py",
         "extmod/asyncio_wait_for_fwd.py",
         "extmod/binascii_a2b_base64.py",
+        "extmod/deflate_compress_memory_error.py",  # tries to allocate unlimited memory
         "extmod/re_stack_overflow.py",
         "extmod/time_res.py",
         "extmod/vfs_posix.py",