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,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