#11116: roll back on error during add so mailbox isn't left corrupted.

2 files changed, 49 insertions, 7 deletions

diff --git a/Lib/mailbox.py b/Lib/mailbox.py
index 53f4159..928e795 100644
--- a/Lib/mailbox.py
+++ b/Lib/mailbox.py
@@ -277,8 +277,11 @@ class Maildir(Mailbox):
         tmp_file = self._create_tmp()
         try:
             self._dump_message(message, tmp_file)
-        finally:
-            _sync_close(tmp_file)
+        except BaseException:
+            tmp_file.close()
+            os.remove(tmp_file.name)
+            raise
+        _sync_close(tmp_file)
         if isinstance(message, MaildirMessage):
             subdir = message.get_subdir()
             suffix = self.colon + message.get_info()
@@ -724,9 +727,14 @@ class _singlefileMailbox(Mailbox):
     def _append_message(self, message):
         """Append message to mailbox and return (start, stop) offsets."""
         self._file.seek(0, 2)
-        self._pre_message_hook(self._file)
-        offsets = self._install_message(message)
-        self._post_message_hook(self._file)
+        before = self._file.tell()
+        try:
+            self._pre_message_hook(self._file)
+            offsets = self._install_message(message)
+            self._post_message_hook(self._file)
+        except BaseException:
+            self._file.truncate(before)
+            raise
         self._file.flush()
         self._file_length = self._file.tell()  # Record current length of mailbox
         return offsets
@@ -906,7 +914,11 @@ class MH(Mailbox):
             if self._locked:
                 _lock_file(f)
             try:
-                self._dump_message(message, f)
+                try:
+                    self._dump_message(message, f)
+                except BaseException:
+                    os.remove(new_path)
+                    raise
                 if isinstance(message, MHMessage):
                     self._dump_sequences(message, new_key)
             finally:
diff --git a/Lib/test/test_mailbox.py b/Lib/test/test_mailbox.py
index 8e4c57a..1e4f887 100644
--- a/Lib/test/test_mailbox.py
+++ b/Lib/test/test_mailbox.py
@@ -107,9 +107,22 @@ class TestMailbox(TestBase):
             'Subject: =?unknown-8bit?b?RmFsaW5hcHThciBo4Xpob3pzeuFsbO104XNz'
             'YWwuIE3hciByZW5kZWx06Ww/?=\n\n')
 
-    def test_add_nonascii_header_raises(self):
+    def test_add_nonascii_string_header_raises(self):
         with self.assertRaisesRegex(ValueError, "ASCII-only"):
             self._box.add(self._nonascii_msg)
+        self._box.flush()
+        self.assertEqual(len(self._box), 0)
+        self.assertMailboxEmpty()
+
+    def test_add_that_raises_leaves_mailbox_empty(self):
+        # XXX This test will start failing when Message learns to handle
+        # non-ASCII string headers, and a different internal failure will
+        # need to be found or manufactured.
+        with self.assertRaises(ValueError):
+            self._box.add(email.message_from_string("From: Alphöso"))
+        self.assertEqual(len(self._box), 0)
+        self._box.close()
+        self.assertMailboxEmpty()
 
     _non_latin_bin_msg = textwrap.dedent("""\
         From: foo@bar.com
@@ -174,6 +187,9 @@ class TestMailbox(TestBase):
         with self.assertWarns(DeprecationWarning):
             with self.assertRaisesRegex(ValueError, "ASCII-only"):
                 self._box.add(io.StringIO(self._nonascii_msg))
+        self.assertEqual(len(self._box), 0)
+        self._box.close()
+        self.assertMailboxEmpty()
 
     def test_remove(self):
         # Remove messages using remove()
@@ -571,6 +587,9 @@ class TestMaildir(TestMailbox):
         if os.name in ('nt', 'os2') or sys.platform == 'cygwin':
             self._box.colon = '!'
 
+    def assertMailboxEmpty(self):
+        self.assertEqual(os.listdir(os.path.join(self._path, 'tmp')), [])
+
     def test_add_MM(self):
         # Add a MaildirMessage instance
         msg = mailbox.MaildirMessage(self._template % 0)
@@ -890,6 +909,10 @@ class _TestMboxMMDF(TestMailbox):
         for lock_remnant in glob.glob(self._path + '.*'):
             support.unlink(lock_remnant)
 
+    def assertMailboxEmpty(self):
+        with open(self._path) as f:
+            self.assertEqual(f.readlines(), [])
+
     def test_add_from_string(self):
         # Add a string starting with 'From ' to the mailbox
         key = self._box.add('From foo@bar blah\nFrom: foo\n\n0')
@@ -1012,6 +1035,9 @@ class TestMH(TestMailbox):
 
     _factory = lambda self, path, factory=None: mailbox.MH(path, factory)
 
+    def assertMailboxEmpty(self):
+        self.assertEqual(os.listdir(self._path), ['.mh_sequences'])
+
     def test_list_folders(self):
         # List folders
         self._box.add_folder('one')
@@ -1144,6 +1170,10 @@ class TestBabyl(TestMailbox):
 
     _factory = lambda self, path, factory=None: mailbox.Babyl(path, factory)
 
+    def assertMailboxEmpty(self):
+        with open(self._path) as f:
+            self.assertEqual(f.readlines(), [])
+
     def tearDown(self):
         super().tearDown()
         self._box.close()