From 397697aa8368f45bd0794e22a25083e9e18e3cfe Mon Sep 17 00:00:00 2001
From: Damien George <damien@micropython.org>
Date: Tue, 5 Sep 2023 18:18:47 +1000
Subject: [PATCH] py/persistentcode: Always close reader even if an exception
 is raised.

Fixes issue #3874.

Signed-off-by: Damien George <damien@micropython.org>
---
 py/persistentcode.c            | 7 ++++++-
 tests/extmod/vfs_userfs.py     | 7 +++++++
 tests/extmod/vfs_userfs.py.exp | 6 ++++++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/py/persistentcode.c b/py/persistentcode.c
index fdc87d5cc8..e1218e0b8a 100644
--- a/py/persistentcode.c
+++ b/py/persistentcode.c
@@ -391,6 +391,10 @@ STATIC mp_raw_code_t *load_raw_code(mp_reader_t *reader, mp_module_context_t *co
 }
 
 void mp_raw_code_load(mp_reader_t *reader, mp_compiled_module_t *cm) {
+    // Set exception handler to close the reader if an exception is raised.
+    MP_DEFINE_NLR_JUMP_CALLBACK_FUNCTION_1(ctx, reader->close, reader->data);
+    nlr_push_jump_callback(&ctx.callback, mp_call_function_1_from_nlr_jump_callback);
+
     byte header[4];
     read_bytes(reader, header, sizeof(header));
     byte arch = MPY_FEATURE_DECODE_ARCH(header[2]);
@@ -435,7 +439,8 @@ void mp_raw_code_load(mp_reader_t *reader, mp_compiled_module_t *cm) {
     cm->n_obj = n_obj;
     #endif
 
-    reader->close(reader->data);
+    // Deregister exception handler and close the reader.
+    nlr_pop_jump_callback(true);
 }
 
 void mp_raw_code_load_mem(const byte *buf, size_t len, mp_compiled_module_t *context) {
diff --git a/tests/extmod/vfs_userfs.py b/tests/extmod/vfs_userfs.py
index 91c355c3df..36f1088870 100644
--- a/tests/extmod/vfs_userfs.py
+++ b/tests/extmod/vfs_userfs.py
@@ -69,6 +69,7 @@ user_files = {
     "/usermod1.py": b"print('in usermod1')\nimport usermod2",
     "/usermod2.py": b"print('in usermod2')",
     "/usermod3.py": b"syntax error",
+    "/usermod4.mpy": b"syntax error",
 }
 os.mount(UserFS(user_files), "/userfs")
 
@@ -86,6 +87,12 @@ try:
 except SyntaxError:
     print("SyntaxError in usermod3")
 
+# import a .mpy file with a syntax error (file should be closed on error)
+try:
+    import usermod4
+except ValueError:
+    print("ValueError in usermod4")
+
 # unmount and undo path addition
 os.umount("/userfs")
 sys.path.pop()
diff --git a/tests/extmod/vfs_userfs.py.exp b/tests/extmod/vfs_userfs.py.exp
index be54da47c2..05cff0452c 100644
--- a/tests/extmod/vfs_userfs.py.exp
+++ b/tests/extmod/vfs_userfs.py.exp
@@ -15,3 +15,9 @@ stat /usermod3.py
 open /usermod3.py rb
 ioctl 4 0
 SyntaxError in usermod3
+stat /usermod4
+stat /usermod4.py
+stat /usermod4.mpy
+open /usermod4.mpy rb
+ioctl 4 0
+ValueError in usermod4