fix: use atomic write in APIKeyManager.save() to prevent credential data loss (#4591) (#4597)

* fix: use atomic write in APIKeyManager.save() to prevent data loss

Opening api_keys.json with 'w' truncates the file before writing, so a
crash, disk-full, or mid-write error leaves all stored provider API keys
corrupted. Switch to atomic write (temp file + fsync + os.replace) so
the original file is always intact on any failure.

Fixes #4591

* chore: trigger CI re-run

* chore: update PR description

* chore: fix how-to-test section for description check

---------

Co-authored-by: michaelxer <michaelxer@users.noreply.github.com>
This commit is contained in:
Michael
2026-06-24 04:28:53 +07:00
committed by GitHub
parent 2e16394b41
commit 72c0bde8a9
2 changed files with 96 additions and 2 deletions
+17 -2
View File
@@ -81,11 +81,26 @@ class APIKeyManager:
keys stay encrypted. Loading via load() first would decrypt them and
write them back as plaintext, which then fails to decrypt on the next
load() and silently drops those providers.
Uses atomic write (temp file + os.replace) so a crash, disk-full, or
mid-write error never truncates the existing keys file.
"""
keys = self._load_raw()
keys[provider] = self.encrypt_api_key(api_key)
with open(self.api_keys_file, 'w', encoding="utf-8") as f:
json.dump(keys, f)
tmp_file = self.api_keys_file + ".tmp"
try:
with open(tmp_file, 'w', encoding="utf-8") as f:
json.dump(keys, f)
f.flush()
os.fsync(f.fileno())
os.replace(tmp_file, self.api_keys_file)
except OSError:
# Clean up temp file on failure; re-raise so callers see the error
try:
os.remove(tmp_file)
except OSError:
pass
raise
def load(self) -> Dict[str, str]:
"""Load and decrypt API keys"""
+79
View File
@@ -0,0 +1,79 @@
"""Test that APIKeyManager.save() uses atomic write to prevent data loss."""
import os
import json
import pytest
from unittest.mock import patch, mock_open
from src.api_key_manager import APIKeyManager
def test_save_creates_atomic_tmp_file(tmp_path):
"""Verify save() writes to a temp file and replaces atomically."""
mgr = APIKeyManager(str(tmp_path))
mgr.save("openai", "sk-test")
# The final file should exist with the correct content
assert os.path.exists(mgr.api_keys_file)
with open(mgr.api_keys_file, "r", encoding="utf-8") as f:
keys = json.load(f)
assert "openai" in keys
# The temp file should NOT remain after successful save
tmp_file = mgr.api_keys_file + ".tmp"
assert not os.path.exists(tmp_file)
def test_save_preserves_existing_keys_atomically(tmp_path):
"""Verify atomic save doesn't corrupt other providers' keys."""
mgr = APIKeyManager(str(tmp_path))
mgr.save("openai", "sk-openai")
mgr.save("anthropic", "sk-anthropic")
loaded = mgr.load()
assert loaded["openai"] == "sk-openai"
assert loaded["anthropic"] == "sk-anthropic"
def test_save_preserves_original_on_write_failure(tmp_path):
"""If the temp file write fails, the original keys file must survive intact."""
mgr = APIKeyManager(str(tmp_path))
mgr.save("openai", "sk-original")
# Now attempt a save that will fail during json.dump
with patch("builtins.open", side_effect=OSError("disk full")):
with pytest.raises(OSError, match="disk full"):
mgr.save("anthropic", "sk-new")
# Original file must still be intact with the original key
loaded = mgr.load()
assert loaded == {"openai": "sk-original"}
assert "anthropic" not in loaded
def test_save_cleans_up_tmp_on_failure(tmp_path):
"""Temp file should be removed if the write fails."""
mgr = APIKeyManager(str(tmp_path))
mgr.save("openai", "sk-original")
tmp_file = mgr.api_keys_file + ".tmp"
# Force a failure after the temp file is opened
original_open = open
def failing_open(*args, **kwargs):
f = original_open(*args, **kwargs)
if args and isinstance(args[0], str) and args[0].endswith(".tmp"):
# Close the file then raise
f.close()
raise OSError("simulated write failure")
return f
with patch("builtins.open", side_effect=failing_open):
with pytest.raises(OSError):
mgr.save("anthropic", "sk-new")
# Temp file should be cleaned up
assert not os.path.exists(tmp_file)
# Original should be intact
loaded = mgr.load()
assert loaded == {"openai": "sk-original"}