mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-06-28 23:52:09 -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
|
keys stay encrypted. Loading via load() first would decrypt them and
|
||||||
write them back as plaintext, which then fails to decrypt on the next
|
write them back as plaintext, which then fails to decrypt on the next
|
||||||
load() and silently drops those providers.
|
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 = self._load_raw()
|
||||||
keys[provider] = self.encrypt_api_key(api_key)
|
keys[provider] = self.encrypt_api_key(api_key)
|
||||||
with open(self.api_keys_file, 'w', encoding="utf-8") as f:
|
tmp_file = self.api_keys_file + ".tmp"
|
||||||
json.dump(keys, f)
|
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]:
|
def load(self) -> Dict[str, str]:
|
||||||
"""Load and decrypt API keys"""
|
"""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