From 72024950cf766c7c1a80d37d1b01bfb602ca12ab Mon Sep 17 00:00:00 2001 From: deacon-botdoctor <291411030+deacon-botdoctor@users.noreply.github.com> Date: Tue, 28 Jul 2026 08:19:24 -0500 Subject: [PATCH] fix(gateway): harden shutdown message flush --- gateway/shutdown_flush.py | 71 ++++++++++++++++++++++----- tests/gateway/test_shutdown_flush.py | 72 +++++++++++++++++++++++++++- 2 files changed, 130 insertions(+), 13 deletions(-) diff --git a/gateway/shutdown_flush.py b/gateway/shutdown_flush.py index 8f0bb478be36..3ffa74ea2564 100644 --- a/gateway/shutdown_flush.py +++ b/gateway/shutdown_flush.py @@ -22,7 +22,10 @@ from __future__ import annotations import json import logging +import os import time +import uuid +from pathlib import Path from typing import Any, Dict, Optional logger = logging.getLogger(__name__) @@ -31,13 +34,59 @@ logger = logging.getLogger(__name__) def _get_flush_dir(): """Return the pending-messages flush directory under the active HERMES_HOME.""" from hermes_constants import get_hermes_home - from pathlib import Path flush_dir = get_hermes_home() / "pending_messages" - flush_dir.mkdir(parents=True, exist_ok=True) + flush_dir.mkdir(parents=True, exist_ok=True, mode=0o700) + if os.name == "posix": + os.chmod(flush_dir, 0o700) return flush_dir +def _fsync_directory(path: Path) -> None: + """Persist a directory entry on platforms that support directory fsync.""" + if os.name != "posix": + return + directory_fd = os.open(path, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + + +def _write_payload(flush_dir: Path, payload: Dict[str, Any]) -> None: + """Atomically write one private, uniquely named recovery payload.""" + file_id = uuid.uuid4().hex + final_path = flush_dir / f"pending-{file_id}.json" + temp_path = flush_dir / f".pending-{file_id}.tmp" + file_descriptor = -1 + + try: + file_descriptor = os.open( + temp_path, + os.O_WRONLY | os.O_CREAT | os.O_EXCL, + 0o600, + ) + with os.fdopen(file_descriptor, "w", encoding="utf-8") as handle: + file_descriptor = -1 # The file object now owns the descriptor. + json.dump(payload, handle, ensure_ascii=False, default=str) + handle.flush() + os.fsync(handle.fileno()) + + os.replace(temp_path, final_path) + except Exception: + if file_descriptor >= 0: + os.close(file_descriptor) + temp_path.unlink(missing_ok=True) + raise + + try: + _fsync_directory(flush_dir) + except OSError as exc: + # The atomically published file is still the only recovery copy. + # Keep it even if this filesystem cannot persist directory entries. + logger.debug("Failed to fsync pending-message directory: %s", exc) + + def flush_pending_to_file( pending: Dict[str, Any], *, @@ -72,16 +121,14 @@ def flush_pending_to_file( serialised = _serialise_value(value) if serialised is None: continue - safe_key = session_key.replace("/", "_").replace("\\", "_") - path = flush_dir / f"{safe_key}_{ts}.json" - path.write_text( - json.dumps({ + _write_payload( + flush_dir, + { "session_key": session_key, "reason": reason, "ts": ts, "data": serialised, - }, ensure_ascii=False, default=str), - encoding="utf-8", + }, ) flushed += 1 except Exception as exc: @@ -148,8 +195,6 @@ def recover_pending_to_db( int Number of messages recovered. """ - from pathlib import Path - flush_dir = _get_flush_dir() flush_files = sorted(flush_dir.glob("*.json")) if not flush_files: @@ -170,7 +215,11 @@ def recover_pending_to_db( data = payload.get("data", {}) text = data.get("text", "") if not text or not session_key: - path.unlink(missing_ok=True) + logger.warning( + "Cannot recover structurally invalid pending message from %s; " + "the flush file has been preserved", + path, + ) continue # The session_key is a gateway routing key (e.g. diff --git a/tests/gateway/test_shutdown_flush.py b/tests/gateway/test_shutdown_flush.py index efbed00f372a..92b070964142 100644 --- a/tests/gateway/test_shutdown_flush.py +++ b/tests/gateway/test_shutdown_flush.py @@ -1,6 +1,8 @@ """Tests for gateway/shutdown_flush.py — pending message durability (#72680).""" import json +import os +import stat import time from pathlib import Path from unittest.mock import MagicMock @@ -44,6 +46,56 @@ def test_flush_writes_string_pending_to_file(tmp_path, monkeypatch): assert payload["session_key"] == "agent:main:telegram:supergroup:123" assert payload["reason"] == "shutdown" assert payload["data"]["text"] == "hello world" + assert ":" not in files[0].name + assert "telegram" not in files[0].name + + +def test_flush_same_session_twice_does_not_overwrite(tmp_path, monkeypatch): + flush_dir = _make_flush_dir(tmp_path) + monkeypatch.setattr("gateway.shutdown_flush._get_flush_dir", lambda: flush_dir) + monkeypatch.setattr("gateway.shutdown_flush.time.time", lambda: 1234) + + pending = {"agent:main:telegram:supergroup:123": "hello world"} + assert flush_pending_to_file(pending, reason="shutdown") == 1 + assert flush_pending_to_file(pending, reason="shutdown") == 1 + + files = list(flush_dir.glob("*.json")) + assert len(files) == 2 + assert all( + json.loads(path.read_text(encoding="utf-8"))["data"]["text"] == "hello world" + for path in files + ) + + +def test_flush_write_failure_leaves_no_recovery_file(tmp_path, monkeypatch): + flush_dir = _make_flush_dir(tmp_path) + monkeypatch.setattr("gateway.shutdown_flush._get_flush_dir", lambda: flush_dir) + + def fail_replace(source, destination): + raise OSError("simulated replace failure") + + monkeypatch.setattr("gateway.shutdown_flush.os.replace", fail_replace) + + assert flush_pending_to_file({"session": "message"}, reason="test") == 0 + assert list(flush_dir.iterdir()) == [] + + +def test_flush_directory_fsync_failure_keeps_recovery_file(tmp_path, monkeypatch): + flush_dir = _make_flush_dir(tmp_path) + monkeypatch.setattr("gateway.shutdown_flush._get_flush_dir", lambda: flush_dir) + + def fail_directory_fsync(path): + raise OSError("simulated directory fsync failure") + + monkeypatch.setattr( + "gateway.shutdown_flush._fsync_directory", fail_directory_fsync + ) + + assert flush_pending_to_file({"session": "message"}, reason="test") == 1 + [flush_file] = list(flush_dir.glob("*.json")) + assert json.loads(flush_file.read_text(encoding="utf-8"))["data"] == { + "text": "message" + } def test_flush_writes_message_event_to_file(tmp_path, monkeypatch): @@ -135,7 +187,7 @@ def test_recover_skips_file_without_session_id(tmp_path, monkeypatch): assert flush_file.exists() -def test_recover_deletes_empty_text_file(tmp_path, monkeypatch): +def test_recover_preserves_structurally_invalid_file(tmp_path, monkeypatch): flush_dir = _make_flush_dir(tmp_path) monkeypatch.setattr( "gateway.shutdown_flush._get_flush_dir", lambda: flush_dir @@ -153,7 +205,7 @@ def test_recover_deletes_empty_text_file(tmp_path, monkeypatch): count = recover_pending_to_db(mock_db) assert count == 0 - assert not flush_file.exists() + assert flush_file.exists() def test_serialise_string(): @@ -197,3 +249,19 @@ def test_get_flush_dir_uses_get_hermes_home(tmp_path, monkeypatch): result = mod._get_flush_dir() assert captured.get("called") is True assert result == tmp_path / "pending_messages" + + +@pytest.mark.skipif( + os.name != "posix", + reason="mode assertions require POSIX permissions", +) +def test_get_flush_dir_and_files_are_private(tmp_path, monkeypatch): + import gateway.shutdown_flush as mod + + monkeypatch.setattr("hermes_constants.get_hermes_home", lambda: tmp_path) + flush_dir = mod._get_flush_dir() + assert stat.S_IMODE(flush_dir.stat().st_mode) == 0o700 + + assert flush_pending_to_file({"session": "message"}, reason="test") == 1 + [flush_file] = list(flush_dir.glob("*.json")) + assert stat.S_IMODE(flush_file.stat().st_mode) == 0o600