mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-06-23 21:25:33 -04:00
* 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:
+17
-2
@@ -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"""
|
||||
|
||||
@@ -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"}
|
||||
Reference in New Issue
Block a user