From fa42487e45620534440256c9b29e4526f3137de9 Mon Sep 17 00:00:00 2001
From: Damien George <damien@micropython.org>
Date: Tue, 18 Mar 2025 00:34:20 +1100
Subject: [PATCH] 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>
---
 extmod/moddeflate.c                           | 12 ++++--
 tests/extmod/deflate_compress_memory_error.py | 39 +++++++++++++++++++
 .../deflate_compress_memory_error.py.exp      |  2 +
 tests/run-tests.py                            |  1 +
 4 files changed, 50 insertions(+), 4 deletions(-)
 create mode 100644 tests/extmod/deflate_compress_memory_error.py
 create mode 100644 tests/extmod/deflate_compress_memory_error.py.exp

diff --git a/extmod/moddeflate.c b/extmod/moddeflate.c
index 4afad01618..920b898b2c 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 0000000000..56ce060308
--- /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 0000000000..606315c146
--- /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 50adb6b4d3..9e7cab4689 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",