feat(host-agent): add McpTokenStore for MCP bearer token

2026-07-21 13:55:24 +08:00
parent 29b9a8c39a
commit c7faee8da3
2 changed files with 225 additions and 0 deletions
@@ -0,0 +1,118 @@
+"""Bearer-token persistence for the host-agent MCP server.
+
+The token is generated on first start and persisted to a JSON file with
+0o600 permissions (POSIX) alongside the host identity. Rotation = delete
+the file and restart host-agent.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import secrets
+import tempfile
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+_TOKEN_BYTES = 32
+
+
+class McpTokenStoreError(RuntimeError):
+    """Raised when the MCP token file cannot be read or written."""
+
+
+@dataclass(frozen=True)
+class McpToken:
+    version: int
+    token: str
+    created_at: datetime
+
+
+class McpTokenStore:
+    def __init__(
+        self,
+        path: Path,
+        *,
+        now: Callable[[], datetime] | None = None,
+    ) -> None:
+        self._path = Path(path)
+        self._now = now or (lambda: datetime.now(UTC))
+
+    def load_or_create(self) -> McpToken:
+        if self._path.exists():
+            return self._read_existing()
+        return self._generate_and_write()
+
+    def verify(self, presented: str) -> bool:
+        try:
+            token = self.load_or_create()
+        except McpTokenStoreError:
+            return False
+        import hmac
+
+        return hmac.compare_digest(token.token, presented)
+
+    def _read_existing(self) -> McpToken:
+        try:
+            data = json.loads(self._path.read_text())
+        except (OSError, json.JSONDecodeError) as exc:
+            raise McpTokenStoreError(
+                f"cannot read MCP token file {self._path}: {exc}"
+            ) from exc
+        if not isinstance(data, dict):
+            raise McpTokenStoreError("MCP token file is not a JSON object")
+        try:
+            return McpToken(
+                version=int(data["version"]),
+                token=str(data["token"]),
+                created_at=datetime.fromisoformat(str(data["created_at"])),
+            )
+        except (KeyError, TypeError, ValueError) as exc:
+            raise McpTokenStoreError(
+                f"MCP token file schema invalid: {exc}"
+            ) from exc
+
+    def _generate_and_write(self) -> McpToken:
+        token = McpToken(
+            version=1,
+            token=secrets.token_urlsafe(_TOKEN_BYTES),
+            created_at=self._now(),
+        )
+        payload = {
+            "version": token.version,
+            "token": token.token,
+            "created_at": token.created_at.isoformat(),
+        }
+        try:
+            self._atomic_write(json.dumps(payload, indent=2))
+        except OSError as exc:
+            raise McpTokenStoreError(
+                f"cannot write MCP token file {self._path}: {exc}"
+            ) from exc
+        return token
+
+    def _atomic_write(self, content: str) -> None:
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        # Atomic on POSIX; on Windows os.replace is also atomic per docs.
+        fd, tmp_name = tempfile.mkstemp(
+            prefix=".host_mcp_token.",
+            suffix=".tmp",
+            dir=str(self._path.parent),
+        )
+        try:
+            with os.fdopen(fd, "w", encoding="utf-8") as fh:
+                fh.write(content)
+            os.chmod(tmp_name, 0o600)
+            os.replace(tmp_name, self._path)
+        except BaseException:
+            try:
+                os.unlink(tmp_name)
+            except OSError:
+                pass
+            raise
@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+import json
+import os
+import stat
+import sys
+from datetime import datetime
+from pathlib import Path
+
+import pytest
+
+from host_agent.mcp_token import McpToken, McpTokenStore, McpTokenStoreError
+
+
+def test_load_or_create_generates_when_missing(tmp_path: Path) -> None:
+    store = McpTokenStore(tmp_path / "host_mcp_token.json")
+    token = store.load_or_create()
+    assert token.version == 1
+    assert len(token.token) >= 40  # secrets.token_urlsafe(32) -> ~43 chars
+    assert isinstance(token.created_at, datetime)
+    # File now exists.
+    assert (tmp_path / "host_mcp_token.json").exists()
+
+
+def test_load_or_create_is_idempotent(tmp_path: Path) -> None:
+    store = McpTokenStore(tmp_path / "host_mcp_token.json")
+    first = store.load_or_create()
+    second = McpTokenStore(tmp_path / "host_mcp_token.json").load_or_create()
+    assert first.token == second.token
+
+
+def test_load_or_create_writes_json_schema(tmp_path: Path) -> None:
+    path = tmp_path / "host_mcp_token.json"
+    McpTokenStore(path).load_or_create()
+    data = json.loads(path.read_text())
+    assert set(data) == {"version", "token", "created_at"}
+    assert data["version"] == 1
+    assert isinstance(data["token"], str)
+    # created_at is ISO 8601.
+    datetime.fromisoformat(data["created_at"])
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="POSIX perms only")
+def test_load_or_create_sets_posix_permissions(tmp_path: Path) -> None:
+    path = tmp_path / "host_mcp_token.json"
+    McpTokenStore(path).load_or_create()
+    mode = stat.S_IMODE(os.fstat(os.open(path, os.O_RDONLY)).st_mode)
+    assert mode == 0o600
+
+
+def test_verify_accepts_correct_token(tmp_path: Path) -> None:
+    store = McpTokenStore(tmp_path / "host_mcp_token.json")
+    token = store.load_or_create()
+    assert store.verify(token.token) is True
+
+
+def test_verify_rejects_wrong_token(tmp_path: Path) -> None:
+    store = McpTokenStore(tmp_path / "host_mcp_token.json")
+    store.load_or_create()
+    assert store.verify("wrong") is False
+
+
+def test_load_or_create_raises_on_corrupt_json(tmp_path: Path) -> None:
+    path = tmp_path / "host_mcp_token.json"
+    path.write_text("{not valid json")
+    with pytest.raises(McpTokenStoreError):
+        McpTokenStore(path).load_or_create()
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="POSIX chmod enforcement only")
+def test_load_or_create_raises_on_unwritable_dir(tmp_path: Path) -> None:
+    unwritable = tmp_path / "ro"
+    unwritable.mkdir()
+    os.chmod(unwritable, 0o500)  # r-x for owner
+    try:
+        with pytest.raises(McpTokenStoreError):
+            McpTokenStore(unwritable / "host_mcp_token.json").load_or_create()
+    finally:
+        os.chmod(unwritable, 0o700)  # restore so cleanup works
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32",
+    reason="POSIX atomic-rename semantics only",
+)
+def test_load_or_create_concurrent_calls_do_not_corrupt(
+    tmp_path: Path,
+) -> None:
+    """Two store instances racing to create: both end up reading the same token."""
+    import threading
+
+    path = tmp_path / "host_mcp_token.json"
+    results: list[McpToken] = []
+    barrier = threading.Barrier(2)
+
+    def worker() -> None:
+        barrier.wait()
+        store = McpTokenStore(path)
+        results.append(store.load_or_create())
+
+    threads = [threading.Thread(target=worker) for _ in range(2)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+    assert len(results) == 2
+    assert results[0].token == results[1].token