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

This commit is contained in:
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