Prevent JSON file corruption with atomic writes - updated

Author: ubaskota

botocore/utils.py

@@ -3583,17 +3583,15 @@ def __setitem__(self, cache_key, value):
             temp_fd, temp_path = tempfile.mkstemp(
                 dir=self._working_dir, suffix='.tmp'
             )
-            if hasattr(os, 'fchmod'):
-                os.fchmod(temp_fd, 0o600)
             with os.fdopen(temp_fd, 'w') as f:
                 temp_fd = None
                 f.write(file_content)
                 f.flush()
                 os.fsync(f.fileno())
-            
+
             os.replace(temp_path, full_key)
             temp_path = None
-            
+
         except Exception:
             if temp_fd is not None:
                 try:

tests/unit/test_utils.py

@@ -17,6 +17,7 @@
 import os
 import shutil
 import tempfile
+import threading
 from contextlib import contextmanager
 from sys import getrefcount
 
@@ -56,7 +57,7 @@
 from botocore.utils import (
     ArgumentGenerator,
     ArnParser,
-    CachedProperty, 
+    CachedProperty,
     ContainerMetadataFetcher,
     IMDSRegionProvider,
     InstanceMetadataFetcher,
@@ -3684,6 +3685,7 @@ def test_get_token_from_environment_returns_none(
     monkeypatch.delenv(env_var, raising=False)
     assert get_token_from_environment(signing_name) is None
 
+
 class TestJSONFileCacheAtomicWrites(unittest.TestCase):
     """Test atomic write operations in JSONFileCache."""
 
@@ -3696,94 +3698,68 @@ def tearDown(self):
 
     @mock.patch('os.replace')
     def test_uses_tempfile_and_replace_for_atomic_write(self, mock_replace):
-
         self.cache['test_key'] = {'data': 'test_value'}
         mock_replace.assert_called_once()
 
-        temp_path, final_path = mock_replace.call_args[0]
-    
-        self.assertIn('.tmp', temp_path)
-        self.assertTrue(final_path.endswith('test_key.json'))
+        call_args = mock_replace.call_args[0]
+        temp_path = call_args[0]
 
-    def test_concurrent_writes_same_key(self):
+        assert '.tmp' in temp_path
+
+    def test_concurrent_writes_to_multiple_temp_files(self):
         """Test concurrent writes to same key don't cause corruption."""
-        import threading
-        
-        key = 'concurrent_test'
         errors = []
-        temp_files_used = []
-        original_mkstemp = tempfile.mkstemp
-
-        def track_temp_files(*args, **kwargs):
-            fd, path = original_mkstemp(*args, **kwargs)
-            temp_files_used.append(path)
-            return fd, path
 
         def write_worker(thread_id):
             try:
+                key = f'concurrent_test_{thread_id}'
                 for i in range(3):
                     self.cache[key] = {'thread': thread_id, 'iteration': i}
-                    if os.name == 'nt':
-                        time.sleep(0.01)
             except Exception as e:
                 errors.append(f'Thread {thread_id}: {e}')
 
-        with mock.patch('tempfile.mkstemp', side_effect=track_temp_files):
-            threads = [
-                threading.Thread(target=write_worker, args=(i,))
-                for i in range(3)
-            ]
-        
-            for thread in threads:
-                thread.start()
-            for thread in threads:
-                thread.join()
-
-        # On Windows, file locking can cause expected write errors
-        # so we allow errors but ensure the key exists in cache.
-        if errors and os.name == 'nt':
-            print(f"Windows file locking warnings: {errors}")
-            self.assertIn(key, self.cache)
-        else:
-            self.assertEqual(len(errors), 0, f'Concurrent write errors: {errors}')
-        
-        # Verify each write used a separate temporary file
-        self.assertEqual(len(temp_files_used), 9)
-        self.assertEqual(
-            len(set(temp_files_used)),
-            9,
-            'Concurrent writes should use separate temp files',
-        )
-
-        # Verify final data is valid
-        final_data = self.cache[key]
-        self.assertIsInstance(final_data, dict)
-        self.assertIn('thread', final_data)
-        self.assertIn('iteration', final_data)
+        threads = [
+            threading.Thread(target=write_worker, args=(i,)) for i in range(3)
+        ]
+
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+
+        self.assertEqual(len(errors), 0, f'Concurrent write errors: {errors}')
+
+        for thread_id in range(3):
+            key = f'concurrent_test_{thread_id}'
+            final_data = self.cache[key]
+            self.assertIsInstance(final_data, dict)
+            self.assertEqual(final_data['thread'], thread_id)
+            self.assertIn('thread', final_data)
+            self.assertIn('iteration', final_data)
 
     def test_atomic_write_preserves_data_on_failure(self):
         """Test write failures don't corrupt existing data."""
         key = 'atomic_test'
         original_data = {'status': 'original'}
-        
+
         self.cache[key] = original_data
-        
+
         # Mock write failure
         original_dumps = self.cache._dumps
         self.cache._dumps = mock.Mock(side_effect=ValueError('Write failed'))
-        
+
         with self.assertRaises(ValueError):
             self.cache[key] = {'status': 'should_fail'}
-        
+
         self.cache._dumps = original_dumps
-        
+
         # Verify original data intact
         self.assertEqual(self.cache[key], original_data)
 
     def test_no_temp_files_after_write(self):
         """Test temporary files cleaned up after writes."""
         self.cache['test'] = {'data': 'value'}
-        
+
         temp_files = [
             f for f in os.listdir(self.temp_dir) if f.endswith('.tmp')
         ]