241 lines
8.7 KiB
Python
241 lines
8.7 KiB
Python
"""OIDC provider 单元测试 (U4 SSO 集成)。
|
||
|
||
验证：
|
||
- redirect URL 生成（含 state）
|
||
- callback 用户 JIT 创建按 sub 关联
|
||
- 禁止邮箱合并 (R1a)
|
||
- 相同 sub 复用现有用户
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import uuid
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import aiosqlite
|
||
import pytest
|
||
|
||
from agentkit.server.auth.idp_registry import (
|
||
IDPConfig,
|
||
IDPRegistry,
|
||
OIDCConfig,
|
||
)
|
||
from agentkit.server.auth.models import init_auth_db
|
||
from agentkit.server.auth.providers.oidc import OIDCProvider, get_state_cache
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Fixtures
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.fixture
|
||
def oidc_idp(monkeypatch: pytest.MonkeyPatch) -> IDPRegistry:
|
||
"""注册一个测试 OIDC IDP 并替换全局 registry。"""
|
||
registry = IDPRegistry()
|
||
registry.register(
|
||
IDPConfig(
|
||
name="test_oidc",
|
||
type="oidc",
|
||
display_name="Test OIDC",
|
||
oidc=OIDCConfig(
|
||
client_id="test-client-id",
|
||
client_secret="test-client-secret",
|
||
server_metadata_url="https://idp.test/authorize",
|
||
redirect_uri="https://app.test/api/v1/auth/oauth/test_oidc/callback",
|
||
scope="openid email profile",
|
||
),
|
||
)
|
||
)
|
||
monkeypatch.setattr("agentkit.server.auth.providers.oidc.get_idp_registry", lambda: registry)
|
||
return registry
|
||
|
||
|
||
@pytest.fixture
|
||
async def auth_db(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
|
||
"""创建临时 auth DB 并指向 AGENTKIT_AUTH_DB。"""
|
||
db_path = tmp_path / "auth.db"
|
||
await init_auth_db(db_path)
|
||
monkeypatch.setenv("AGENTKIT_AUTH_DB", str(db_path))
|
||
return db_path
|
||
|
||
|
||
class _FakeOAuth2Client:
|
||
"""替代 AsyncOAuth2Client 的假 client — 不发 HTTP。"""
|
||
|
||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||
self.token = {"access_token": "fake-access-token", "id_token": None}
|
||
|
||
async def __aenter__(self) -> "_FakeOAuth2Client":
|
||
return self
|
||
|
||
async def __aexit__(self, *args: Any) -> None:
|
||
pass
|
||
|
||
def create_authorization_url(self, url: str, **kwargs: Any) -> tuple[str, dict[str, Any]]:
|
||
"""构造假 authorize URL（含 state / scope 参数）。"""
|
||
state = kwargs.get("state", "")
|
||
scope = kwargs.get("scope", "")
|
||
client_id = kwargs.get("client_id", "test-client-id")
|
||
redirect_uri = kwargs.get("redirect_uri", "")
|
||
return (
|
||
f"{url}?response_type=code&client_id={client_id}"
|
||
f"&redirect_uri={redirect_uri}&scope={scope}&state={state}",
|
||
{"state": state, "url": url},
|
||
)
|
||
|
||
async def fetch_token(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
|
||
return self.token
|
||
|
||
async def get(self, url: str) -> Any:
|
||
class _Resp:
|
||
status_code = 200
|
||
|
||
def json(self) -> dict[str, Any]:
|
||
return {
|
||
"sub": "oidc-sub-123",
|
||
"email": "sso_user@example.com",
|
||
"name": "SSO User",
|
||
}
|
||
|
||
return _Resp()
|
||
|
||
|
||
@pytest.fixture
|
||
def fake_oauth_client(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
"""替换 oidc 模块的 AsyncOAuth2Client 为假实现。"""
|
||
monkeypatch.setattr("agentkit.server.auth.providers.oidc.AsyncOAuth2Client", _FakeOAuth2Client)
|
||
|
||
|
||
@pytest.fixture(autouse=True)
|
||
def _reset_state_cache() -> None:
|
||
"""每个测试前清空 state 缓存。"""
|
||
cache = get_state_cache()
|
||
cache._store.clear()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tests
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestGetRedirectUrl:
|
||
async def test_generates_url_with_state(self, oidc_idp: IDPRegistry) -> None:
|
||
url = await OIDCProvider().get_redirect_url("test_oidc")
|
||
assert "https://idp.test/authorize" in url
|
||
assert "state=" in url
|
||
assert "client_id=test-client-id" in url
|
||
|
||
async def test_unknown_provider_raises(self, oidc_idp: IDPRegistry) -> None:
|
||
with pytest.raises(ValueError, match="未配置"):
|
||
await OIDCProvider().get_redirect_url("nonexistent")
|
||
|
||
async def test_state_cached_for_csrf(self, oidc_idp: IDPRegistry) -> None:
|
||
url = await OIDCProvider().get_redirect_url("test_oidc")
|
||
state = url.split("state=", 1)[1].split("&", 1)[0]
|
||
# state 应在缓存中（consume 验证）
|
||
assert get_state_cache().consume(state) is True
|
||
|
||
|
||
class TestHandleCallback:
|
||
async def test_creates_user_by_sub(
|
||
self, oidc_idp: IDPRegistry, auth_db: Path, fake_oauth_client: None
|
||
) -> None:
|
||
"""首次登录 JIT 创建用户，按 sub 关联到 user_idp_links。"""
|
||
# 预置 state
|
||
provider = OIDCProvider()
|
||
url = await provider.get_redirect_url("test_oidc")
|
||
state = url.split("state=", 1)[1].split("&", 1)[0]
|
||
|
||
user = await provider.handle_callback("test_oidc", "fake-code", state)
|
||
assert user["username"] == "SSO User"
|
||
assert user["email"] == "sso_user@example.com"
|
||
|
||
# 验证 user_idp_links 按 sub 关联
|
||
async with aiosqlite.connect(str(auth_db)) as db:
|
||
db.row_factory = aiosqlite.Row
|
||
cursor = await db.execute(
|
||
"SELECT * FROM user_idp_links WHERE idp_name = ? AND idp_subject = ?",
|
||
("test_oidc", "oidc-sub-123"),
|
||
)
|
||
link = await cursor.fetchone()
|
||
assert link is not None
|
||
assert link["user_id"] == user["id"]
|
||
|
||
async def test_same_sub_reuses_user(
|
||
self, oidc_idp: IDPRegistry, auth_db: Path, fake_oauth_client: None
|
||
) -> None:
|
||
"""相同 sub 第二次登录复用现有用户，不重复创建。"""
|
||
provider = OIDCProvider()
|
||
# 第一次登录
|
||
url1 = await provider.get_redirect_url("test_oidc")
|
||
state1 = url1.split("state=", 1)[1].split("&", 1)[0]
|
||
user1 = await provider.handle_callback("test_oidc", "code1", state1)
|
||
|
||
# 第二次登录（相同 sub）
|
||
url2 = await provider.get_redirect_url("test_oidc")
|
||
state2 = url2.split("state=", 1)[1].split("&", 1)[0]
|
||
user2 = await provider.handle_callback("test_oidc", "code2", state2)
|
||
|
||
assert user1["id"] == user2["id"]
|
||
|
||
async def test_no_email_merge(
|
||
self,
|
||
oidc_idp: IDPRegistry,
|
||
auth_db: Path,
|
||
fake_oauth_client: None,
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
"""R1a: 不同 sub 但相同 email — 禁止合并，创建新用户。"""
|
||
# 预创建一个本地用户，email 与 SSO userinfo 相同
|
||
now = datetime.now(timezone.utc).isoformat()
|
||
local_id = str(uuid.uuid4())
|
||
async with aiosqlite.connect(str(auth_db)) as db:
|
||
await db.execute(
|
||
"INSERT INTO users (id, username, email, password_hash, role, is_active, "
|
||
"is_terminal_authorized, is_server_terminal_authorized, created_at, updated_at) "
|
||
"VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||
(
|
||
local_id,
|
||
"local_user",
|
||
"sso_user@example.com",
|
||
"$2b$12$dummyhash",
|
||
"member",
|
||
1,
|
||
0,
|
||
0,
|
||
now,
|
||
now,
|
||
),
|
||
)
|
||
await db.commit()
|
||
|
||
# SSO 用不同 sub（_FakeOAuth2Client 返回 sub=oidc-sub-123）
|
||
provider = OIDCProvider()
|
||
url = await provider.get_redirect_url("test_oidc")
|
||
state = url.split("state=", 1)[1].split("&", 1)[0]
|
||
sso_user = await provider.handle_callback("test_oidc", "code", state)
|
||
|
||
# 应创建新用户，而非复用 local_user
|
||
assert sso_user["id"] != local_id
|
||
# R1a: 邮箱冲突时不合并本地用户，改用 fallback 邮箱（保留 UNIQUE 约束）
|
||
assert sso_user["email"] != "sso_user@example.com"
|
||
assert sso_user["email"].endswith("@sso.test_oidc.local")
|
||
|
||
async def test_invalid_state_raises(
|
||
self, oidc_idp: IDPRegistry, auth_db: Path, fake_oauth_client: None
|
||
) -> None:
|
||
"""无效 state（CSRF 校验失败）应拒绝。"""
|
||
with pytest.raises(ValueError, match="state"):
|
||
await OIDCProvider().handle_callback("test_oidc", "code", "invalid-state")
|
||
|
||
|
||
class TestAuthenticateNotImplemented:
|
||
async def test_raises(self) -> None:
|
||
from agentkit.server.auth.providers.exceptions import ProviderNotImplemented
|
||
|
||
with pytest.raises(ProviderNotImplemented):
|
||
await OIDCProvider().authenticate(username="x", password="y")
|