update

Serialize the replacement to a temporary file before overwriting the existing MH message, so invalid input or input read errors leave the original message unchanged.

Co-Authored-By: lipengyu <lipengyu@kylinos.cn>

Author: lpyu001

Lib/mailbox.py

@@ -18,6 +18,8 @@
 import email.generator
 import io
 import contextlib
+import shutil
+import tempfile
 from types import GenericAlias
 try:
     import fcntl
@@ -1069,44 +1071,22 @@ def __setitem__(self, key, message):
                 raise KeyError('No message with key: %s' % key)
             else:
                 raise
-        file_closed = False
         try:
             if self._locked:
                 _lock_file(f)
             try:
-                new_file = _create_temporary(path)
-                try:
+                with tempfile.TemporaryFile(mode='w+b') as new_file:
                     self._dump_message(message, new_file)
-                    _sync_close(new_file)
-                    info = os.fstat(f.fileno())
-                    try:
-                        os.chown(new_file.name, info.st_uid, info.st_gid)
-                    except (AttributeError, OSError):
-                        pass
-                    os.chmod(new_file.name, info.st_mode)
-                    if os.name == 'nt':
-                        # Windows cannot replace an open file.
-                        f.close()
-                        file_closed = True
-                    os.replace(new_file.name, path)
-                except BaseException:
-                    try:
-                        new_file.close()
-                    except OSError:
-                        pass
-                    try:
-                        os.remove(new_file.name)
-                    except OSError:
-                        pass
-                    raise
+                    new_file.seek(0)
+                    os.close(os.open(path, os.O_WRONLY | os.O_TRUNC))
+                    shutil.copyfileobj(new_file, f)
                 if isinstance(message, MHMessage):
                     self._dump_sequences(message, key)
             finally:
                 if self._locked:
                     _unlock_file(f)
         finally:
-            if not file_closed:
-                f.close()
+            _sync_close(f)
 
     def get_message(self, key):
         """Return a Message representation or raise a KeyError."""

Lib/test/test_mailbox.py

@@ -1377,18 +1377,29 @@ def readline(self):
         self._box = self._factory(self._path)
         self.assertEqual(self._box.get_bytes(key), original)
 
-    @unittest.skipUnless(hasattr(os, 'chown'), 'requires os.chown')
+    @unittest.skipUnless(os.name == 'posix', 'requires POSIX permissions')
     def test_set_item_preserves_mode(self):
         key = self._box.add(self._template % 'original')
         path = os.path.join(self._path, str(key))
-        mode = os.stat(path).st_mode | stat.S_ISUID
+        mode = 0o640
         os.chmod(path, mode)
-        if os.stat(path).st_mode != mode:
-            self.skipTest('filesystem does not support set-user-ID mode')
+        if stat.S_IMODE(os.stat(path).st_mode) != mode:
+            self.skipTest('filesystem does not support POSIX permissions')
 
         self._box[key] = self._template % 'replacement'
 
-        self.assertEqual(os.stat(path).st_mode, mode)
+        self.assertEqual(stat.S_IMODE(os.stat(path).st_mode), mode)
+
+    def test_set_item_with_open_file(self):
+        key = self._box.add(self._template % 'original')
+        replacement = self._template % 'replacement'
+        self._box.lock()
+        try:
+            with self._box.get_file(key):
+                self._box[key] = replacement
+            self.assertEqual(self._box.get_bytes(key), replacement.encode('ascii'))
+        finally:
+            self._box.unlock()
 
     def test_list_folders(self):
         # List folders