mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-07 01:48:53 +00:00
fix(config): harden repository writes
This commit is contained in:
parent
f5371a6c5a
commit
27f7549c84
@ -7,6 +7,7 @@ import json
|
|||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
import stat
|
import stat
|
||||||
|
import threading
|
||||||
import uuid
|
import uuid
|
||||||
from collections.abc import Callable
|
from collections.abc import Callable
|
||||||
from contextlib import suppress
|
from contextlib import suppress
|
||||||
@ -23,6 +24,11 @@ ConfigMutator = Callable[[Config], None]
|
|||||||
|
|
||||||
_ENV_REF_PATTERN = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}")
|
_ENV_REF_PATTERN = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}")
|
||||||
|
|
||||||
|
# Several WebUI config actions run in worker threads via asyncio.to_thread().
|
||||||
|
# Repository instances for the same path therefore share a process-local lock.
|
||||||
|
_path_locks_guard = threading.Lock()
|
||||||
|
_path_locks: dict[Path, threading.RLock] = {}
|
||||||
|
|
||||||
|
|
||||||
class ConfigConflictError(RuntimeError):
|
class ConfigConflictError(RuntimeError):
|
||||||
"""Raised when a caller tries to update a stale configuration snapshot."""
|
"""Raised when a caller tries to update a stale configuration snapshot."""
|
||||||
@ -55,6 +61,12 @@ class ConfigCommit:
|
|||||||
changed_paths: frozenset[str]
|
changed_paths: frozenset[str]
|
||||||
|
|
||||||
|
|
||||||
|
def _lock_for_path(path: Path) -> threading.RLock:
|
||||||
|
key = path.expanduser().resolve(strict=False)
|
||||||
|
with _path_locks_guard:
|
||||||
|
return _path_locks.setdefault(key, threading.RLock())
|
||||||
|
|
||||||
|
|
||||||
def _revision_for_bytes(raw: bytes | None) -> str:
|
def _revision_for_bytes(raw: bytes | None) -> str:
|
||||||
if raw is None:
|
if raw is None:
|
||||||
return "missing"
|
return "missing"
|
||||||
@ -100,10 +112,13 @@ def _write_config_atomic(path: Path, data: dict[str, Any]) -> None:
|
|||||||
existing_mode = stat.S_IMODE(path.stat().st_mode) if path.exists() else None
|
existing_mode = stat.S_IMODE(path.stat().st_mode) if path.exists() else None
|
||||||
try:
|
try:
|
||||||
with open(tmp, "w", encoding="utf-8") as handle:
|
with open(tmp, "w", encoding="utf-8") as handle:
|
||||||
handle.write(content)
|
|
||||||
if existing_mode is not None:
|
if existing_mode is not None:
|
||||||
with suppress(OSError):
|
try:
|
||||||
os.chmod(tmp, existing_mode)
|
os.chmod(tmp, existing_mode)
|
||||||
|
except OSError:
|
||||||
|
if os.name != "nt":
|
||||||
|
raise
|
||||||
|
handle.write(content)
|
||||||
handle.flush()
|
handle.flush()
|
||||||
os.fsync(handle.fileno())
|
os.fsync(handle.fileno())
|
||||||
os.replace(tmp, path)
|
os.replace(tmp, path)
|
||||||
@ -140,10 +155,12 @@ class FileConfigRepository:
|
|||||||
|
|
||||||
def __init__(self, path: str | Path):
|
def __init__(self, path: str | Path):
|
||||||
self.path = Path(path).expanduser().resolve(strict=False)
|
self.path = Path(path).expanduser().resolve(strict=False)
|
||||||
|
self._lock = _lock_for_path(self.path)
|
||||||
|
|
||||||
def load_raw(self) -> PersistedConfigSnapshot:
|
def load_raw(self) -> PersistedConfigSnapshot:
|
||||||
"""Load the persisted form without resolving secret references."""
|
"""Load the persisted form without resolving secret references."""
|
||||||
return _read_snapshot(self.path)
|
with self._lock:
|
||||||
|
return _read_snapshot(self.path)
|
||||||
|
|
||||||
def load_effective(self) -> EffectiveConfigSnapshot:
|
def load_effective(self) -> EffectiveConfigSnapshot:
|
||||||
"""Load an isolated runtime snapshot with ``${VAR}`` references resolved."""
|
"""Load an isolated runtime snapshot with ``${VAR}`` references resolved."""
|
||||||
@ -158,16 +175,17 @@ class FileConfigRepository:
|
|||||||
expected_revision: str | None = None,
|
expected_revision: str | None = None,
|
||||||
) -> PersistedConfigSnapshot:
|
) -> PersistedConfigSnapshot:
|
||||||
"""Atomically save a complete config, optionally rejecting stale writes."""
|
"""Atomically save a complete config, optionally rejecting stale writes."""
|
||||||
current = _read_snapshot(self.path)
|
with self._lock:
|
||||||
if expected_revision is not None and current.revision != expected_revision:
|
current = _read_snapshot(self.path)
|
||||||
raise ConfigConflictError(
|
if expected_revision is not None and current.revision != expected_revision:
|
||||||
f"Config changed since revision {expected_revision}; "
|
raise ConfigConflictError(
|
||||||
f"current revision is {current.revision}"
|
f"Config changed since revision {expected_revision}; "
|
||||||
)
|
f"current revision is {current.revision}"
|
||||||
data = _config_data(config)
|
)
|
||||||
_validate_config_data(data, self.path)
|
data = _config_data(config)
|
||||||
_write_config_atomic(self.path, data)
|
_validate_config_data(data, self.path)
|
||||||
return _read_snapshot(self.path)
|
_write_config_atomic(self.path, data)
|
||||||
|
return _read_snapshot(self.path)
|
||||||
|
|
||||||
def update(
|
def update(
|
||||||
self,
|
self,
|
||||||
@ -176,25 +194,26 @@ class FileConfigRepository:
|
|||||||
expected_revision: str | None = None,
|
expected_revision: str | None = None,
|
||||||
) -> ConfigCommit:
|
) -> ConfigCommit:
|
||||||
"""Atomically apply a mutation to the latest persisted config."""
|
"""Atomically apply a mutation to the latest persisted config."""
|
||||||
before = _read_snapshot(self.path)
|
with self._lock:
|
||||||
if expected_revision is not None and before.revision != expected_revision:
|
before = _read_snapshot(self.path)
|
||||||
raise ConfigConflictError(
|
if expected_revision is not None and before.revision != expected_revision:
|
||||||
f"Config changed since revision {expected_revision}; "
|
raise ConfigConflictError(
|
||||||
f"current revision is {before.revision}"
|
f"Config changed since revision {expected_revision}; "
|
||||||
)
|
f"current revision is {before.revision}"
|
||||||
|
)
|
||||||
|
|
||||||
before_data = _config_data(before.config)
|
before_data = _config_data(before.config)
|
||||||
draft = before.config.model_copy(deep=True)
|
draft = before.config.model_copy(deep=True)
|
||||||
mutator(draft)
|
mutator(draft)
|
||||||
after_data = _config_data(draft)
|
after_data = _config_data(draft)
|
||||||
changed = frozenset(_changed_paths(before_data, after_data))
|
changed = frozenset(_changed_paths(before_data, after_data))
|
||||||
if not changed:
|
if not changed:
|
||||||
return ConfigCommit(before, before, changed)
|
return ConfigCommit(before, before, changed)
|
||||||
|
|
||||||
_validate_config_data(after_data, self.path)
|
_validate_config_data(after_data, self.path)
|
||||||
_write_config_atomic(self.path, after_data)
|
_write_config_atomic(self.path, after_data)
|
||||||
after = _read_snapshot(self.path)
|
after = _read_snapshot(self.path)
|
||||||
return ConfigCommit(before, after, changed)
|
return ConfigCommit(before, after, changed)
|
||||||
|
|
||||||
|
|
||||||
def resolve_config_env_vars(config: Config) -> Config:
|
def resolve_config_env_vars(config: Config) -> Config:
|
||||||
|
|||||||
@ -2,6 +2,8 @@ import json
|
|||||||
import os
|
import os
|
||||||
import socket
|
import socket
|
||||||
import stat
|
import stat
|
||||||
|
import threading
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from unittest.mock import patch
|
from unittest.mock import patch
|
||||||
|
|
||||||
@ -83,15 +85,37 @@ def test_update_rejects_stale_expected_revision(tmp_path: Path) -> None:
|
|||||||
assert repository.load_raw().config.api.port == 9001
|
assert repository.load_raw().config.api.port == 9001
|
||||||
|
|
||||||
|
|
||||||
def test_repositories_for_same_path_read_latest_config(tmp_path: Path) -> None:
|
def test_repositories_for_same_path_serialize_threaded_updates(tmp_path: Path) -> None:
|
||||||
path = tmp_path / "config.json"
|
path = tmp_path / "config.json"
|
||||||
first = FileConfigRepository(path)
|
first = FileConfigRepository(path)
|
||||||
second = FileConfigRepository(path)
|
second = FileConfigRepository(path)
|
||||||
|
first_mutator_entered = threading.Event()
|
||||||
|
release_first = threading.Event()
|
||||||
|
second_mutator_entered = threading.Event()
|
||||||
|
|
||||||
first.update(lambda config: setattr(config.api, "port", 9001))
|
def update_first() -> None:
|
||||||
second.update(
|
def mutate(config):
|
||||||
lambda config: setattr(config.agents.defaults, "timezone", "Asia/Shanghai")
|
first_mutator_entered.set()
|
||||||
)
|
assert release_first.wait(timeout=5)
|
||||||
|
config.api.port = 9001
|
||||||
|
|
||||||
|
first.update(mutate)
|
||||||
|
|
||||||
|
def update_second() -> None:
|
||||||
|
def mutate(config):
|
||||||
|
second_mutator_entered.set()
|
||||||
|
config.agents.defaults.timezone = "Asia/Shanghai"
|
||||||
|
|
||||||
|
second.update(mutate)
|
||||||
|
|
||||||
|
with ThreadPoolExecutor(max_workers=2) as pool:
|
||||||
|
first_future = pool.submit(update_first)
|
||||||
|
assert first_mutator_entered.wait(timeout=5)
|
||||||
|
second_future = pool.submit(update_second)
|
||||||
|
assert not second_mutator_entered.wait(timeout=0.1)
|
||||||
|
release_first.set()
|
||||||
|
first_future.result(timeout=5)
|
||||||
|
second_future.result(timeout=5)
|
||||||
|
|
||||||
config = first.load_raw().config
|
config = first.load_raw().config
|
||||||
assert config.api.port == 9001
|
assert config.api.port == 9001
|
||||||
@ -150,6 +174,40 @@ def test_atomic_save_preserves_existing_file_mode(tmp_path: Path) -> None:
|
|||||||
assert stat.S_IMODE(path.stat().st_mode) == 0o600
|
assert stat.S_IMODE(path.stat().st_mode) == 0o600
|
||||||
|
|
||||||
|
|
||||||
|
def test_atomic_save_restores_mode_before_writing(tmp_path: Path, monkeypatch) -> None:
|
||||||
|
path = tmp_path / "config.json"
|
||||||
|
path.write_text('{"api": {"port": 9000}}', encoding="utf-8")
|
||||||
|
path.chmod(0o600)
|
||||||
|
observed_sizes: list[int] = []
|
||||||
|
real_chmod = os.chmod
|
||||||
|
|
||||||
|
def capture_size_before_chmod(target: str | bytes | Path, mode: int) -> None:
|
||||||
|
observed_sizes.append(Path(target).stat().st_size)
|
||||||
|
real_chmod(target, mode)
|
||||||
|
|
||||||
|
monkeypatch.setattr("nanobot.config.repository.os.chmod", capture_size_before_chmod)
|
||||||
|
|
||||||
|
FileConfigRepository(path).update(lambda config: setattr(config.api, "port", 9001))
|
||||||
|
|
||||||
|
assert observed_sizes == [0]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(os.name == "nt", reason="Windows handles chmod differently")
|
||||||
|
def test_atomic_save_keeps_previous_file_when_mode_restore_fails(tmp_path: Path) -> None:
|
||||||
|
path = tmp_path / "config.json"
|
||||||
|
path.write_text('{"api": {"port": 9000}}', encoding="utf-8")
|
||||||
|
path.chmod(0o600)
|
||||||
|
|
||||||
|
with patch("nanobot.config.repository.os.chmod", side_effect=OSError("chmod failed")):
|
||||||
|
with pytest.raises(OSError, match="chmod failed"):
|
||||||
|
FileConfigRepository(path).update(
|
||||||
|
lambda config: setattr(config.api, "port", 9001)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert json.loads(path.read_text(encoding="utf-8"))["api"]["port"] == 9000
|
||||||
|
assert list(tmp_path.glob(".config.json.*.tmp")) == []
|
||||||
|
|
||||||
|
|
||||||
def test_loading_config_does_not_change_process_network_policy(tmp_path: Path) -> None:
|
def test_loading_config_does_not_change_process_network_policy(tmp_path: Path) -> None:
|
||||||
path = tmp_path / "config.json"
|
path = tmp_path / "config.json"
|
||||||
path.write_text(json.dumps({"tools": {"ssrfWhitelist": []}}), encoding="utf-8")
|
path.write_text(json.dumps({"tools": {"ssrfWhitelist": []}}), encoding="utf-8")
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user