357 lines
12 KiB
Python
357 lines
12 KiB
Python
"""Tests for the OAuth identity resolver: no mix-up between identities, upstream refresh."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from mcp.server.auth.provider import AccessToken
|
|
|
|
from aegis_gitea_mcp.auth import STATIC_CLIENT_ID
|
|
from aegis_gitea_mcp.errors import ToolError
|
|
from aegis_gitea_mcp.identity import Identity
|
|
from aegis_gitea_mcp.oauth.resolver import (
|
|
OAUTH_ONLY,
|
|
SIGN_IN_AGAIN,
|
|
OAuthGiteaExecutor,
|
|
OAuthIdentityResolver,
|
|
SessionIdentity,
|
|
)
|
|
from aegis_gitea_mcp.oauth.store import OAuthStore, SessionRecord, TokenRecord
|
|
from aegis_gitea_mcp.oauth.upstream import UpstreamError, UpstreamInvalidGrant, UpstreamTokens
|
|
from tests.test_oauth_helpers import (
|
|
ENCRYPTION_KEY,
|
|
GITEA_TOKEN,
|
|
USER_REFRESH,
|
|
USER_TOKEN,
|
|
FakeClock,
|
|
oauth_settings,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _env(mock_env: None) -> None:
|
|
"""The audit logger reads the process-wide settings."""
|
|
|
|
|
|
class FakeUpstream:
|
|
"""Stands in for GiteaUpstream.refresh."""
|
|
|
|
def __init__(self, *outcomes: UpstreamTokens | Exception, delay: float = 0.0) -> None:
|
|
self.outcomes = list(outcomes)
|
|
self.calls: list[str] = []
|
|
self.delay = delay
|
|
|
|
async def refresh(self, refresh_token: str) -> UpstreamTokens:
|
|
self.calls.append(refresh_token)
|
|
if self.delay:
|
|
await asyncio.sleep(self.delay)
|
|
outcome = self.outcomes.pop(0)
|
|
if isinstance(outcome, Exception):
|
|
raise outcome
|
|
return outcome
|
|
|
|
|
|
class FakeAdmin:
|
|
def __init__(self) -> None:
|
|
self.seen: list[AccessToken | None] = []
|
|
|
|
async def resolve(self, access: AccessToken | None) -> Identity:
|
|
self.seen.append(access)
|
|
return Identity(login="root", mode="admin-token", token=GITEA_TOKEN)
|
|
|
|
|
|
def make_store(tmp_path: Path, clock: FakeClock, **session: Any) -> OAuthStore:
|
|
store = OAuthStore.load(tmp_path, ENCRYPTION_KEY, clock)
|
|
values: dict[str, Any] = {
|
|
"session_id": "s1",
|
|
"client_id": "c1",
|
|
"login": "alice",
|
|
"is_admin": False,
|
|
"gitea_access_token": USER_TOKEN,
|
|
"gitea_refresh_token": USER_REFRESH,
|
|
"gitea_expires_at": clock.now + 3600,
|
|
"scopes": ["gitea"],
|
|
"resource": "https://mcp.example.com/mcp",
|
|
"created_at": clock.now,
|
|
}
|
|
values.update(session)
|
|
store.sessions["s1"] = SessionRecord(**values)
|
|
store.tokens["h"] = TokenRecord("access", "s1", "c1", ["gitea"], clock.now + 3600)
|
|
return store
|
|
|
|
|
|
def oauth_access(session_id: Any = "s1", subject: str = "alice", **claims: Any) -> AccessToken:
|
|
return AccessToken(
|
|
token="sha256:x",
|
|
client_id="c1",
|
|
scopes=["gitea"],
|
|
subject=subject,
|
|
claims={"auth": "oauth", "session_id": session_id, **claims},
|
|
)
|
|
|
|
|
|
STATIC = AccessToken(token="static", client_id=STATIC_CLIENT_ID, scopes=["gitea"])
|
|
|
|
|
|
def resolver(
|
|
store: OAuthStore, upstream: Any = None, admin: Any = None, clock: Any = None
|
|
) -> OAuthIdentityResolver:
|
|
return OAuthIdentityResolver(
|
|
store, upstream or FakeUpstream(), admin, clock=clock or FakeClock()
|
|
)
|
|
|
|
|
|
async def test_oauth_token_resolves_to_the_session_user(tmp_path: Path) -> None:
|
|
clock = FakeClock()
|
|
identity = await resolver(make_store(tmp_path, clock), clock=clock).resolve(oauth_access())
|
|
|
|
assert isinstance(identity, SessionIdentity)
|
|
assert (identity.login, identity.mode, identity.token) == ("alice", "oauth", USER_TOKEN)
|
|
assert identity.redact == (USER_REFRESH,)
|
|
assert USER_TOKEN not in repr(identity) and USER_REFRESH not in repr(identity)
|
|
assert identity.metadata() == {"login": "alice", "mode": "oauth", "sudo": None}
|
|
|
|
|
|
async def test_static_bearer_and_no_token_go_to_the_admin_resolver(tmp_path: Path) -> None:
|
|
clock = FakeClock()
|
|
admin = FakeAdmin()
|
|
subject = resolver(make_store(tmp_path, clock), admin=admin, clock=clock)
|
|
|
|
assert (await subject.resolve(STATIC)).mode == "admin-token"
|
|
assert (await subject.resolve(None)).mode == "admin-token"
|
|
assert admin.seen == [STATIC, None]
|
|
|
|
|
|
async def test_oauth_only_deployment_fails_closed_for_admin_mode(tmp_path: Path) -> None:
|
|
clock = FakeClock()
|
|
subject = resolver(make_store(tmp_path, clock), admin=None, clock=clock)
|
|
|
|
for access in (STATIC, None):
|
|
with pytest.raises(ToolError) as caught:
|
|
await subject.resolve(access)
|
|
assert caught.value.detail == OAUTH_ONLY and caught.value.status_code == 401
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"access",
|
|
[
|
|
AccessToken(token="t", client_id="some-client", scopes=["gitea"]),
|
|
AccessToken(token="t", client_id="some-client", scopes=["gitea"], claims={"auth": "x"}),
|
|
oauth_access(session_id=None),
|
|
oauth_access(session_id=5),
|
|
],
|
|
)
|
|
async def test_unknown_token_kinds_never_reach_the_admin(
|
|
tmp_path: Path, access: AccessToken
|
|
) -> None:
|
|
clock = FakeClock()
|
|
admin = FakeAdmin()
|
|
|
|
with pytest.raises(ToolError, match="unknown identity"):
|
|
await resolver(make_store(tmp_path, clock), admin=admin, clock=clock).resolve(access)
|
|
assert admin.seen == []
|
|
|
|
|
|
async def test_oauth_marker_on_a_static_client_id_still_needs_a_session(tmp_path: Path) -> None:
|
|
"""A token claiming OAuth is never treated as the static bearer."""
|
|
clock = FakeClock()
|
|
admin = FakeAdmin()
|
|
forged = AccessToken(
|
|
token="t",
|
|
client_id=STATIC_CLIENT_ID,
|
|
scopes=["gitea"],
|
|
subject="alice",
|
|
claims={"auth": "oauth", "session_id": "missing"},
|
|
)
|
|
|
|
with pytest.raises(ToolError, match="sign in again"):
|
|
await resolver(make_store(tmp_path, clock), admin=admin, clock=clock).resolve(forged)
|
|
assert admin.seen == []
|
|
|
|
|
|
@pytest.mark.parametrize(("session_id", "subject"), [("gone", "alice"), ("s1", "mallory")])
|
|
async def test_ended_session_or_subject_mismatch_means_sign_in_again(
|
|
tmp_path: Path, session_id: str, subject: str
|
|
) -> None:
|
|
clock = FakeClock()
|
|
admin = FakeAdmin()
|
|
|
|
with pytest.raises(ToolError) as caught:
|
|
await resolver(make_store(tmp_path, clock), admin=admin, clock=clock).resolve(
|
|
oauth_access(session_id, subject)
|
|
)
|
|
|
|
assert caught.value.detail == SIGN_IN_AGAIN and caught.value.status_code == 401
|
|
assert admin.seen == []
|
|
|
|
|
|
async def test_token_without_expiry_is_never_refreshed(tmp_path: Path) -> None:
|
|
clock = FakeClock()
|
|
upstream = FakeUpstream()
|
|
store = make_store(tmp_path, clock, gitea_expires_at=None)
|
|
|
|
await resolver(store, upstream, clock=clock).resolve(oauth_access())
|
|
|
|
assert upstream.calls == []
|
|
|
|
|
|
async def test_refresh_keeps_the_old_refresh_token_when_gitea_sends_none(tmp_path: Path) -> None:
|
|
clock = FakeClock()
|
|
upstream = FakeUpstream(UpstreamTokens("new-access", None, clock.now + 3600))
|
|
store = make_store(tmp_path, clock, gitea_expires_at=clock.now + 10)
|
|
|
|
identity = await resolver(store, upstream, clock=clock).resolve(oauth_access())
|
|
|
|
assert identity.token == "new-access"
|
|
assert store.sessions["s1"].gitea_refresh_token == USER_REFRESH
|
|
reloaded = OAuthStore.load(tmp_path, ENCRYPTION_KEY, clock)
|
|
assert reloaded.sessions["s1"].gitea_access_token == "new-access"
|
|
|
|
|
|
async def test_concurrent_calls_refresh_once(tmp_path: Path) -> None:
|
|
clock = FakeClock()
|
|
upstream = FakeUpstream(
|
|
UpstreamTokens("new-access", "new-refresh", clock.now + 3600), delay=0.01
|
|
)
|
|
store = make_store(tmp_path, clock, gitea_expires_at=clock.now + 10)
|
|
subject = resolver(store, upstream, clock=clock)
|
|
|
|
identities = await asyncio.gather(*(subject.resolve(oauth_access()) for _ in range(5)))
|
|
|
|
assert upstream.calls == [USER_REFRESH]
|
|
assert {identity.token for identity in identities} == {"new-access"}
|
|
|
|
|
|
async def test_transient_failure_keeps_a_still_valid_token(tmp_path: Path) -> None:
|
|
clock = FakeClock()
|
|
upstream = FakeUpstream(UpstreamError("token request failed: HTTP 502"))
|
|
store = make_store(tmp_path, clock, gitea_expires_at=clock.now + 10)
|
|
|
|
identity = await resolver(store, upstream, clock=clock).resolve(oauth_access())
|
|
|
|
assert identity.token == USER_TOKEN
|
|
assert "s1" in store.sessions
|
|
|
|
|
|
async def test_transient_failure_with_an_expired_token_is_retryable(tmp_path: Path) -> None:
|
|
clock = FakeClock()
|
|
upstream = FakeUpstream(UpstreamError("token request failed: HTTP 502"))
|
|
store = make_store(tmp_path, clock, gitea_expires_at=clock.now - 10)
|
|
|
|
with pytest.raises(ToolError) as caught:
|
|
await resolver(store, upstream, clock=clock).resolve(oauth_access())
|
|
|
|
assert caught.value.status_code == 503
|
|
assert "s1" in store.sessions
|
|
|
|
|
|
async def test_invalid_grant_ends_the_session_and_its_tokens(tmp_path: Path) -> None:
|
|
clock = FakeClock()
|
|
upstream = FakeUpstream(UpstreamInvalidGrant("Gitea rejected the grant"))
|
|
store = make_store(tmp_path, clock, gitea_expires_at=clock.now + 10)
|
|
admin = FakeAdmin()
|
|
|
|
with pytest.raises(ToolError, match="sign in again"):
|
|
await resolver(store, upstream, admin=admin, clock=clock).resolve(oauth_access())
|
|
|
|
assert store.sessions == {} and store.tokens == {}
|
|
assert admin.seen == []
|
|
assert OAuthStore.load(tmp_path, ENCRYPTION_KEY, clock).sessions == {}
|
|
|
|
|
|
@pytest.mark.parametrize(("offset", "survives"), [(30.0, True), (-30.0, False)])
|
|
async def test_no_refresh_token(tmp_path: Path, offset: float, survives: bool) -> None:
|
|
clock = FakeClock()
|
|
store = make_store(
|
|
tmp_path, clock, gitea_refresh_token=None, gitea_expires_at=clock.now + offset
|
|
)
|
|
subject = resolver(store, FakeUpstream(), clock=clock)
|
|
|
|
if survives:
|
|
assert (await subject.resolve(oauth_access())).token == USER_TOKEN
|
|
else:
|
|
with pytest.raises(ToolError, match="sign in again"):
|
|
await subject.resolve(oauth_access())
|
|
assert ("s1" in store.sessions) is survives
|
|
|
|
|
|
async def test_session_ended_while_waiting_for_the_refresh_lock(tmp_path: Path) -> None:
|
|
clock = FakeClock()
|
|
store = make_store(tmp_path, clock, gitea_expires_at=clock.now + 10)
|
|
subject = resolver(store, FakeUpstream(), clock=clock)
|
|
lock = subject._locks.setdefault("s1", asyncio.Lock())
|
|
|
|
async with lock:
|
|
task = asyncio.ensure_future(subject.resolve(oauth_access()))
|
|
await asyncio.sleep(0)
|
|
store.end_session("s1")
|
|
with pytest.raises(ToolError, match="sign in again"):
|
|
await task
|
|
|
|
|
|
async def test_save_failures_are_logged_not_raised(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
|
) -> None:
|
|
clock = FakeClock()
|
|
store = make_store(tmp_path, clock, gitea_expires_at=clock.now + 10)
|
|
|
|
async def broken_save() -> None:
|
|
raise OSError("disk full")
|
|
|
|
monkeypatch.setattr(store, "save", broken_save)
|
|
upstream = FakeUpstream(
|
|
UpstreamTokens("new-access", None, clock.now + 3600), UpstreamInvalidGrant("x")
|
|
)
|
|
subject = resolver(store, upstream, clock=clock)
|
|
|
|
assert (await subject.resolve(oauth_access())).token == "new-access"
|
|
store.sessions["s1"].gitea_expires_at = clock.now + 10
|
|
with pytest.raises(ToolError):
|
|
await subject.resolve(oauth_access())
|
|
|
|
assert "could not save the refreshed Gitea token" in caplog.text
|
|
assert "could not save the OAuth store" in caplog.text
|
|
|
|
|
|
async def test_audit_failure_does_not_block_ending_a_session(
|
|
tmp_path: Path, caplog: pytest.LogCaptureFixture
|
|
) -> None:
|
|
clock = FakeClock()
|
|
store = make_store(tmp_path, clock, gitea_expires_at=clock.now + 10)
|
|
|
|
class BrokenAudit:
|
|
def log_security_event(self, *args: Any, **kwargs: Any) -> None:
|
|
raise RuntimeError("audit down")
|
|
|
|
subject = OAuthIdentityResolver(
|
|
store,
|
|
FakeUpstream(UpstreamInvalidGrant("x")), # type: ignore[arg-type]
|
|
None,
|
|
clock=clock,
|
|
audit=BrokenAudit(), # type: ignore[arg-type]
|
|
)
|
|
|
|
with pytest.raises(ToolError):
|
|
await subject.resolve(oauth_access())
|
|
assert store.sessions == {}
|
|
assert "audit log write failed: RuntimeError" in caplog.text
|
|
|
|
|
|
async def test_executor_redacts_session_secrets(httpx_mock: Any) -> None:
|
|
executor = OAuthGiteaExecutor(oauth_settings(), extra_secrets=[ENCRYPTION_KEY])
|
|
httpx_mock.add_response(
|
|
url="https://gitea.example.com/api/v1/version",
|
|
json={"a": USER_REFRESH, "b": ENCRYPTION_KEY, "c": USER_TOKEN},
|
|
)
|
|
identity = SessionIdentity(
|
|
login="alice", mode="oauth", token=USER_TOKEN, redact=(USER_REFRESH,)
|
|
)
|
|
|
|
outcome = await executor.request("GET", "/version", identity)
|
|
await executor.aclose()
|
|
|
|
assert outcome.payload["data"] == {"a": "[REDACTED]", "b": "[REDACTED]", "c": "[REDACTED]"}
|