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

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

@@ -6,6 +6,8 @@
 import json
 import logging
 import os
+import stat
+import tempfile
 import time
 from collections.abc import MutableMapping
 
@@ -50,9 +52,45 @@ 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:
+            # The directory is not 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.
             with open(self.filename, 'w', encoding=self._encoding) as f:
                 json.dump(self.data, f)
+            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,120 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import json
+import os
+import stat
+import tempfile
+import unittest
+
+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 _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())
+
+    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.skipUnless(hasattr(os, 'symlink'), 'requires symlink support')
+    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')
+    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_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()