feat(host-agent): add McpTokenStore for MCP bearer token
This commit is contained in:
@@ -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
|
||||||
Reference in New Issue
Block a user