From 5099ae4fe2ec3774c6d7b8dc0e65af23f18b4325 Mon Sep 17 00:00:00 2001 From: benya Date: Sun, 15 Feb 2026 01:12:41 +0300 Subject: [PATCH] Fix critical race conditions and unsafe disk report command --- lock_utils.py | 10 ++++++-- services/disk_report.py | 48 ++++++++++++++++++++++++++++++++++---- services/runtime_state.py | 49 ++++++++++++++++++++++++++++----------- 3 files changed, 87 insertions(+), 20 deletions(-) diff --git a/lock_utils.py b/lock_utils.py index 0770639..eba52c3 100644 --- a/lock_utils.py +++ b/lock_utils.py @@ -1,4 +1,5 @@ from pathlib import Path +import os import time LOCK_DIR = Path("/var/run/tg-bot") @@ -11,9 +12,14 @@ def lock_path(name: str) -> Path: def acquire_lock(name: str) -> bool: p = lock_path(name) - if p.exists(): + try: + fd = os.open(str(p), os.O_CREAT | os.O_EXCL | os.O_WRONLY) + except FileExistsError: return False - p.write_text(str(time.time())) + try: + os.write(fd, str(time.time()).encode("ascii", errors="ignore")) + finally: + os.close(fd) return True diff --git a/services/disk_report.py b/services/disk_report.py index 2734603..b94609f 100644 --- a/services/disk_report.py +++ b/services/disk_report.py @@ -1,11 +1,48 @@ import os +import re from typing import Any from services.runner import run_cmd def _top_dirs_cmd(path: str, limit: int) -> list[str]: - return ["bash", "-lc", f"du -xhd1 {path} 2>/dev/null | sort -h | tail -n {limit}"] + _ = limit + return ["du", "-x", "-h", "-d", "1", path] + + +_SIZE_RE = re.compile(r"^\s*([0-9]+(?:\.[0-9]+)?)([KMGTP]?)(i?B?)?$", re.IGNORECASE) + + +def _size_to_bytes(value: str) -> float: + m = _SIZE_RE.match(value.strip()) + if not m: + return -1.0 + num = float(m.group(1)) + unit = (m.group(2) or "").upper() + mul = { + "": 1, + "K": 1024, + "M": 1024**2, + "G": 1024**3, + "T": 1024**4, + "P": 1024**5, + }.get(unit, 1) + return num * mul + + +def _format_top_dirs(raw: str, limit: int) -> str: + rows: list[tuple[float, str]] = [] + for line in raw.splitlines(): + line = line.strip() + if not line: + continue + parts = line.split(maxsplit=1) + if len(parts) != 2: + continue + size, name = parts + rows.append((_size_to_bytes(size), f"{size}\t{name}")) + rows.sort(key=lambda x: x[0]) + return "\n".join(line for _sz, line in rows[-max(1, limit):]) async def build_disk_report(cfg: dict[str, Any], mount: str, usage: int) -> str: @@ -15,24 +52,27 @@ async def build_disk_report(cfg: dict[str, Any], mount: str, usage: int) -> str: rc, out = await run_cmd(_top_dirs_cmd(mount, limit), timeout=30) if rc == 0 and out.strip(): + top_out = _format_top_dirs(out, limit) lines.append("") lines.append("Top directories:") - lines.append(out.strip()) + lines.append(top_out) docker_dir = cfg.get("disk_report", {}).get("docker_dir", "/var/lib/docker") if docker_dir and os.path.exists(docker_dir): rc2, out2 = await run_cmd(_top_dirs_cmd(docker_dir, limit), timeout=30) if rc2 == 0 and out2.strip(): + top_out2 = _format_top_dirs(out2, limit) lines.append("") lines.append(f"Docker dir: {docker_dir}") - lines.append(out2.strip()) + lines.append(top_out2) logs_dir = cfg.get("disk_report", {}).get("logs_dir", "/var/log") if logs_dir and os.path.exists(logs_dir): rc3, out3 = await run_cmd(_top_dirs_cmd(logs_dir, limit), timeout=30) if rc3 == 0 and out3.strip(): + top_out3 = _format_top_dirs(out3, limit) lines.append("") lines.append(f"Logs dir: {logs_dir}") - lines.append(out3.strip()) + lines.append(top_out3) return "\n".join(lines) diff --git a/services/runtime_state.py b/services/runtime_state.py index e4a409f..92a16a7 100644 --- a/services/runtime_state.py +++ b/services/runtime_state.py @@ -1,9 +1,13 @@ import json import os +import threading +import tempfile from typing import Any, Dict _PATH = "/var/server-bot/runtime.json" _STATE: Dict[str, Any] = {} +_LOCK = threading.RLock() +_LOADED = False def configure(path: str | None): @@ -13,40 +17,57 @@ def configure(path: str | None): def _load_from_disk(): - global _STATE + global _STATE, _LOADED if not os.path.exists(_PATH): _STATE = {} + _LOADED = True return try: with open(_PATH, "r", encoding="utf-8") as f: _STATE = json.load(f) except Exception: _STATE = {} + _LOADED = True def _save(): - os.makedirs(os.path.dirname(_PATH), exist_ok=True) + directory = os.path.dirname(_PATH) or "." + os.makedirs(directory, exist_ok=True) try: - with open(_PATH, "w", encoding="utf-8") as f: - json.dump(_STATE, f) + fd, tmp_path = tempfile.mkstemp(prefix=".runtime.", suffix=".json", dir=directory) + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + json.dump(_STATE, f, ensure_ascii=False) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, _PATH) + finally: + if os.path.exists(tmp_path): + try: + os.unlink(tmp_path) + except Exception: + pass except Exception: pass def get_state() -> Dict[str, Any]: - if not _STATE: - _load_from_disk() - return _STATE + with _LOCK: + if not _LOADED: + _load_from_disk() + return _STATE def set_state(key: str, value: Any): - if not _STATE: - _load_from_disk() - _STATE[key] = value - _save() + with _LOCK: + if not _LOADED: + _load_from_disk() + _STATE[key] = value + _save() def get(key: str, default: Any = None) -> Any: - if not _STATE: - _load_from_disk() - return _STATE.get(key, default) + with _LOCK: + if not _LOADED: + _load_from_disk() + return _STATE.get(key, default)