4 files changed, 50 insertions, 4 deletions

diff --git a/extmod/moddeflate.c b/extmod/moddeflate.c
index 4afad0161..920b898b2 100644
--- a/extmod/moddeflate.c
+++ b/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;
diff --git a/tests/extmod/deflate_compress_memory_error.py b/tests/extmod/deflate_compress_memory_error.py
new file mode 100644
index 000000000..56ce06030
--- /dev/null
+++ b/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())
diff --git a/tests/extmod/deflate_compress_memory_error.py.exp b/tests/extmod/deflate_compress_memory_error.py.exp
new file mode 100644
index 000000000..606315c14
--- /dev/null
+++ b/tests/extmod/deflate_compress_memory_error.py.exp
@@ -0,0 +1,2 @@
+MemoryError
+b''
diff --git a/tests/run-tests.py b/tests/run-tests.py
index 50adb6b4d..9e7cab468 100755
--- a/tests/run-tests.py
+++ b/tests/run-tests.py
@@ -163,6 +163,7 @@ platform_tests_to_skip = {
         "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",