[Core] Write session files atomically so a partial write cannot destroy them

src/azure-cli-core/azure/cli/core/_session.py

@@ -3,9 +3,12 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
+import errno
 import json
 import logging
 import os
+import stat
+import tempfile
 import time
 from collections.abc import MutableMapping
 
@@ -50,9 +53,51 @@ def load(self, filename, max_age=0):
             self.save()
 
     def save(self):
-        if self.filename:
+        if not self.filename:
+            return
+
+        # Write to a temporary file in the same directory and rename it over the target, so the
+        # file on disk is never observably incomplete. Opening the target with 'w' truncates it
+        # in place, which leaves a window where a concurrent process reads zero or half a
+        # document. That reader then fails to parse it and load() overwrites it with defaults,
+        # so a process that only meant to read destroys the data.
+        target = os.path.realpath(self.filename)
+        directory = os.path.dirname(target) or os.curdir
+
+        try:
+            fd, temp_name = tempfile.mkstemp(dir=directory, prefix=os.path.basename(target) + '.',
+                                             suffix='.tmp')
+        except OSError as ex:
+            if ex.errno not in (errno.EACCES, errno.EPERM):
+                # Anything else, a full disk or a read only file system, would break the in place
+                # write too, so there is nothing to fall back to and it should surface.
+                raise
+            # The directory denies writes but the file itself is writable, which happens in locked
+            # down containers and CI images. Writing in place still works there, so keep the old
+            # behaviour rather than failing a save that used to succeed. Serialize first so that
+            # unserializable data raises before the file is opened and truncated.
+            content = json.dumps(self.data)
             with open(self.filename, 'w', encoding=self._encoding) as f:
+                f.write(content)
+            return
+
+        try:
+            with os.fdopen(fd, 'w', encoding=self._encoding) as f:
                 json.dump(self.data, f)
+            # mkstemp creates the file 0o600. Carry over the permissions of the file being
+            # replaced so a save does not silently tighten them.
+            try:
+                os.chmod(temp_name, stat.S_IMODE(os.stat(target).st_mode))
+            except OSError:
+                pass
+            os.replace(temp_name, target)
+        except BaseException:
+            # BaseException so that an interrupt does not leave the temporary file behind.
+            try:
+                os.remove(temp_name)
+            except OSError:
+                pass
+            raise
 
     def save_with_retry(self, retries=5):
         for _ in range(retries - 1):

src/azure-cli-core/azure/cli/core/tests/test_session.py

@@ -0,0 +1,162 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import errno
+import json
+import os
+import shutil
+import stat
+import tempfile
+import unittest
+from unittest import mock
+
+from azure.cli.core._session import Session
+
+
+class TestSession(unittest.TestCase):
+
+    def setUp(self):
+        self.dir = tempfile.mkdtemp()
+        self.filename = os.path.join(self.dir, 'test.json')
+
+    def tearDown(self):
+        shutil.rmtree(self.dir, ignore_errors=True)
+
+    def _session(self, data):
+        session = Session()
+        session.filename = self.filename
+        session.data = data
+        return session
+
+    def _read(self):
+        with open(self.filename, encoding='utf-8-sig') as f:
+            return json.load(f)
+
+    def _temp_files(self):
+        return [name for name in os.listdir(self.dir) if name.endswith('.tmp')]
+
+    def test_save_writes_the_data(self):
+        self._session({'a': 1}).save()
+        self.assertEqual(self._read(), {'a': 1})
+
+    def test_save_replaces_the_file_instead_of_truncating_it(self):
+        # A truncating write leaves the file empty or half written while it runs, and a concurrent
+        # reader that fails to parse it has load() overwrite it with defaults. Replacing the file
+        # means a reader always sees either the old contents or the new ones.
+        with open(self.filename, 'w', encoding='utf-8-sig') as f:
+            json.dump({'a': 1}, f)
+        inode = os.stat(self.filename).st_ino
+
+        self._session({'a': 2}).save()
+
+        self.assertNotEqual(os.stat(self.filename).st_ino, inode)
+        self.assertEqual(self._read(), {'a': 2})
+
+    def test_save_leaves_the_file_intact_when_the_data_cannot_be_serialized(self):
+        with open(self.filename, 'w', encoding='utf-8-sig') as f:
+            json.dump({'kept': True}, f)
+
+        with self.assertRaises(TypeError):
+            self._session({'a': 'fine', 'b': {1, 2}}).save()
+
+        self.assertEqual(self._read(), {'kept': True})
+        self.assertFalse(self._temp_files())
+
+    @unittest.skipIf(os.name == 'nt', 'POSIX file permissions not applicable on Windows')
+    def test_save_preserves_the_permissions_of_an_existing_file(self):
+        with open(self.filename, 'w', encoding='utf-8-sig') as f:
+            json.dump({}, f)
+        os.chmod(self.filename, 0o644)
+
+        self._session({'a': 1}).save()
+
+        self.assertEqual(stat.S_IMODE(os.stat(self.filename).st_mode), 0o644)
+
+    @unittest.skipIf(os.name == 'nt', 'Symlink test not applicable on Windows')
+    def test_save_follows_a_symlink_rather_than_replacing_it(self):
+        target = os.path.join(self.dir, 'target.json')
+        link = os.path.join(self.dir, 'link.json')
+        with open(target, 'w', encoding='utf-8-sig') as f:
+            json.dump({'a': 1}, f)
+        os.symlink(target, link)
+
+        session = Session()
+        session.filename = link
+        session.data = {'a': 2}
+        session.save()
+
+        self.assertTrue(os.path.islink(link))
+        with open(target, encoding='utf-8-sig') as f:
+            self.assertEqual(json.load(f), {'a': 2})
+
+    @unittest.skipIf(os.name == 'nt', 'directory permissions are not enforced the same way on Windows')
+    @unittest.skipIf(hasattr(os, 'geteuid') and os.geteuid() == 0, 'root ignores the directory mode')
+    def test_save_still_writes_when_the_directory_is_not_writable(self):
+        # Locked down containers and build images mount the config directory read only while
+        # leaving the file itself writable. A save that used to succeed there must keep working.
+        with open(self.filename, 'w', encoding='utf-8-sig') as f:
+            json.dump({}, f)
+        os.chmod(self.dir, 0o500)
+        try:
+            self._session({'a': 1}).save()
+            self.assertEqual(self._read(), {'a': 1})
+        finally:
+            os.chmod(self.dir, 0o700)
+
+    def test_save_does_not_fall_back_when_the_temporary_file_fails_for_another_reason(self):
+        # The in place write would also fail on a full disk, and it would lose the file doing it,
+        # so only a permission problem should reach the fallback.
+        with open(self.filename, 'w', encoding='utf-8-sig') as f:
+            json.dump({'kept': True}, f)
+
+        with mock.patch('tempfile.mkstemp', side_effect=OSError(errno.ENOSPC, 'No space left')):
+            with self.assertRaises(OSError):
+                self._session({'a': 1}).save()
+
+        self.assertEqual(self._read(), {'kept': True})
+
+    def test_save_falls_back_without_losing_the_file_when_the_data_is_bad(self):
+        # The fallback must keep the same guarantee as the atomic path: data that cannot be
+        # serialized has to fail before the target is opened.
+        with open(self.filename, 'w', encoding='utf-8-sig') as f:
+            json.dump({'kept': True}, f)
+
+        with mock.patch('tempfile.mkstemp', side_effect=OSError(errno.EACCES, 'Permission denied')):
+            with self.assertRaises(TypeError):
+                self._session({'a': 'fine', 'b': {1, 2}}).save()
+
+        self.assertEqual(self._read(), {'kept': True})
+
+    def test_save_uses_the_fallback_when_the_directory_denies_permission(self):
+        # Runs everywhere, including as root where the directory mode would not stop mkstemp.
+        with open(self.filename, 'w', encoding='utf-8-sig') as f:
+            json.dump({}, f)
+
+        with mock.patch('tempfile.mkstemp', side_effect=OSError(errno.EACCES, 'Permission denied')):
+            self._session({'a': 1}).save()
+
+        self.assertEqual(self._read(), {'a': 1})
+
+    def test_load_reads_back_what_save_wrote(self):
+        self._session({'a': 1}).save()
+
+        session = Session()
+        session.load(self.filename)
+
+        self.assertEqual(session.data, {'a': 1})
+
+    def test_load_overrides_a_file_that_cannot_be_parsed(self):
+        with open(self.filename, 'w', encoding='utf-8-sig') as f:
+            f.write('{"truncated"')
+
+        session = Session()
+        session.load(self.filename)
+
+        self.assertEqual(session.data, {})
+        self.assertEqual(self._read(), {})
+
+
+if __name__ == '__main__':
+    unittest.main()