Refactor to composable KayaMixin architecture

Replace wrapper-based SessionMiddleware/OIDCApp with KayaMixin subclasses
applied via KayaApp(mixins=[...]). Mixins hook into handle_request and
handle_websocket via before/after hooks, so both ASGI and RSGI keep working.
Mixin dependencies are applied automatically and deduplicated.
This commit is contained in:
2026-07-23 22:09:59 +08:00
parent 24a797e3d2
commit 3ebf079533
14 changed files with 480 additions and 359 deletions
+13 -10
View File
@@ -10,34 +10,36 @@ Flow with PKCE**.
```python
import os
from kaya.core import HttpContext, KayaApp
from kaya.session import SessionMiddleware, InMemorySessionStore
from kaya.oidc import OIDCConfig, OIDCApp
from kaya.session import SessionMixin, InMemorySessionStore
from kaya.oidc import OIDCConfig, OIDCMixin
app = KayaApp()
session_app = SessionMiddleware(app, InMemorySessionStore())
oidc = OIDCApp(
session_app,
session = SessionMixin(InMemorySessionStore())
oidc = OIDCMixin(
OIDCConfig(
issuer=os.environ['OIDC_ISSUER'],
client_id=os.environ['OIDC_CLIENT_ID'],
client_secret=os.environ.get('OIDC_CLIENT_SECRET'),
redirect_uri='http://localhost:8000/auth/callback',
fetch_userinfo=True,
)
),
session=session,
)
app = KayaApp(mixins=[session, oidc])
@oidc.GET('/')
@app.GET('/')
async def home(ctx: HttpContext):
await ctx.send_str(200, 'public home')
@oidc.GET('/profile')
@app.GET('/profile')
@oidc.require_auth
async def profile(ctx: HttpContext):
user = oidc.get_user(ctx)
await ctx.send_str(200, f'Hello {user.email or user.sub}')
```
`OIDCMixin` depends on `SessionMixin`; passing only `oidc` to `KayaApp(mixins=...)`
also works because the app applies mixin dependencies automatically.
## Features
- Generic OIDC discovery
@@ -48,6 +50,7 @@ async def profile(ctx: HttpContext):
- Optional userinfo endpoint fetch
- Refresh token support
- RP-initiated logout (when provider advertises `end_session_endpoint`)
- Composable with any other `KayaMixin` (RSGI, MCP, etc.)
## Security notes
+2 -2
View File
@@ -1,11 +1,11 @@
from ._app import OIDCApp, OIDCUser
from ._client import OIDCClient
from ._config import OIDCConfig
from ._mixin import OIDCMixin, OIDCUser
__all__ = [
'OIDCApp',
'OIDCClient',
'OIDCConfig',
'OIDCMixin',
'OIDCUser',
]
@@ -1,21 +1,14 @@
from typing import Any, Awaitable, Callable, Mapping, MutableMapping, Optional, Sequence, cast
from urllib.parse import parse_qs, urlencode
from typing import Any, Awaitable, Callable, Mapping, Optional, Sequence
from kaya.core import HttpContext, HttpMethod, KayaApp
from kaya.session import Session, SessionMiddleware
from kaya.core import HttpContext, KayaApp, KayaMixin
from kaya.session import Session, SessionMixin
from urllib.parse import parse_qs
from ._client import OIDCClient
from ._config import OIDCConfig
type HttpHandler = Callable[..., Awaitable[None]]
type WebSocketHandler = Callable[..., Awaitable[None]]
type RouteDecorator = Callable[[HttpHandler], HttpHandler]
type WebSocketDecorator = Callable[[WebSocketHandler], WebSocketHandler]
type ASGIApp = Callable[
[MutableMapping[str, Any], Callable[[], Awaitable[Any]], Callable[[MutableMapping[str, Any]], Awaitable[None]]],
Awaitable[None],
]
class OIDCUser(Mapping[str, Any]):
@@ -50,42 +43,52 @@ class OIDCUser(Mapping[str, Any]):
return self._data.get('picture')
class OIDCApp:
"""ASGI app wrapper that adds OIDC authentication routes to a Kaya app.
class OIDCMixin(KayaMixin):
"""Kaya mixin adding OpenID Connect authentication.
The wrapped app must be a ``SessionMiddleware`` instance so that OIDC state,
nonce, and user data can be stored in ``ctx.session``.
Depends on :class:`~kaya.session.SessionMixin` so that OIDC state, nonce,
and user data can be stored in ``ctx.session``. The dependency is applied
automatically by ``KayaApp``.
Built-in routes:
Registers three routes on the app:
- ``login_path`` (default ``/auth/login``): redirects to the OIDC provider.
- ``callback_path`` (default ``/auth/callback``): handles the provider callback.
- ``logout_path`` (default ``/auth/logout``): logs the user out.
Example::
session = SessionMixin(InMemorySessionStore())
oidc = OIDCMixin(config, session=session)
app = KayaApp(mixins=[session, oidc])
@app.GET('/profile')
@oidc.require_auth
async def profile(ctx: HttpContext):
user = oidc.get_user(ctx)
...
"""
def __init__(self, app: SessionMiddleware, config: OIDCConfig) -> None:
self._app = app
def __init__(self, config: OIDCConfig, session: SessionMixin) -> None:
self._config = config
self._session = session
self._client = OIDCClient(config)
self._register_routes()
@staticmethod
def _session(ctx: HttpContext) -> Session:
session = ctx.session
assert isinstance(session, Session)
return session
@property
def dependencies(self) -> Sequence[KayaMixin]:
return [self._session]
def _register_routes(self) -> None:
@self._app.GET(self._config.login_path)
def apply(self, app: KayaApp) -> None:
@app.GET(self._config.login_path)
async def login(ctx: HttpContext) -> None:
auth_url, state, nonce, code_verifier = await self._client.build_authorization_url()
session = self._session(ctx)
session = self._session_of(ctx)
session['oidc_state'] = state
session['oidc_nonce'] = nonce
session['oidc_code_verifier'] = code_verifier
await ctx.send_empty(302, {'Location': auth_url})
@self._app.GET(self._config.callback_path)
@app.GET(self._config.callback_path)
async def callback(ctx: HttpContext) -> None:
query = parse_qs(ctx.query_string)
code = self._first_value(query.get('code'))
@@ -104,7 +107,7 @@ class OIDCApp:
await ctx.send_str(400, 'Missing code or state')
return
session = self._session(ctx)
session = self._session_of(ctx)
expected_state = session.get('oidc_state')
if state != expected_state:
await ctx.send_str(400, 'Invalid state')
@@ -144,66 +147,32 @@ class OIDCApp:
except ValueError as exc:
await ctx.send_str(400, f'Authentication failed: {exc}')
@self._app.GET(self._config.logout_path)
@app.GET(self._config.logout_path)
async def logout(ctx: HttpContext) -> None:
session = self._session(ctx)
session = self._session_of(ctx)
id_token = session.get('oidc_id_token')
session.invalidate()
logout_url = await self._client.build_logout_url(id_token if isinstance(id_token, str) else None)
location = logout_url if logout_url is not None else self._config.post_logout_redirect
await ctx.send_empty(302, {'Location': location})
@staticmethod
def _session_of(ctx: HttpContext) -> Session:
session = ctx.session
assert isinstance(session, Session)
return session
@staticmethod
def _first_value(values: Optional[Sequence[str]]) -> Optional[str]:
if values and len(values) > 0:
return values[0]
return None
def route(
self,
paths: str | Sequence[str],
methods: Optional[HttpMethod | Sequence[HttpMethod]] = None,
recursive: bool = False,
) -> RouteDecorator:
return self._app.route(paths, methods, recursive)
def GET(self, path: str, recursive: bool = False) -> RouteDecorator:
return self._app.GET(path, recursive)
def POST(self, path: str, recursive: bool = False) -> RouteDecorator:
return self._app.POST(path, recursive)
def PUT(self, path: str, recursive: bool = False) -> RouteDecorator:
return self._app.PUT(path, recursive)
def DELETE(self, path: str, recursive: bool = False) -> RouteDecorator:
return self._app.DELETE(path, recursive)
def OPTIONS(self, path: str, recursive: bool = False) -> RouteDecorator:
return self._app.OPTIONS(path, recursive)
def HEAD(self, path: str, recursive: bool = False) -> RouteDecorator:
return self._app.HEAD(path, recursive)
def PATCH(self, path: str, recursive: bool = False) -> RouteDecorator:
return self._app.PATCH(path, recursive)
def websocket(self, path: str, recursive: bool = False) -> WebSocketDecorator:
return self._app.websocket(path, recursive)
async def __call__(
self,
scope: MutableMapping[str, Any],
receive: Callable[[], Awaitable[Any]],
send: Callable[[MutableMapping[str, Any]], Awaitable[None]],
) -> None:
await cast(ASGIApp, self._app)(scope, receive, send)
def is_authenticated(self, ctx: HttpContext) -> bool:
return 'oidc_user' in self._session(ctx)
return 'oidc_user' in self._session_of(ctx)
def get_user(self, ctx: HttpContext) -> Optional[OIDCUser]:
user = self._session(ctx).get('oidc_user')
user = self._session_of(ctx).get('oidc_user')
if user is None or not isinstance(user, Mapping):
return None
return OIDCUser(user)
+63 -37
View File
@@ -1,19 +1,17 @@
import base64
import json
import unittest
from time import time
from typing import Any, Mapping
from typing import Mapping
from urllib.parse import parse_qs, urlparse
import httpx
import jwt
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from pwo import async_test
from kaya.core import HttpContext, KayaApp
from kaya.session import InMemorySessionStore, SessionMiddleware
from kaya.session import InMemorySessionStore, SessionMixin
from kaya.oidc import OIDCApp, OIDCConfig
from kaya.oidc import OIDCConfig, OIDCMixin
from kaya.oidc._client import OIDCClient
@@ -197,14 +195,13 @@ class OIDCClientTest(unittest.TestCase):
self.assertEqual(['/'], query['post_logout_redirect_uri'])
class OIDCAppTest(unittest.TestCase):
class OIDCMixinTest(unittest.TestCase):
def _build_app(self, fetch_userinfo: bool = False) -> tuple[OIDCApp, MockOIDCProvider, InMemorySessionStore]:
def _build_app(self, fetch_userinfo: bool = False) -> tuple[KayaApp, OIDCMixin, MockOIDCProvider, InMemorySessionStore]:
provider = MockOIDCProvider()
http_client = httpx.AsyncClient(transport=MockTransport(provider))
store = InMemorySessionStore()
app = KayaApp()
session_app = SessionMiddleware(app, store)
session = SessionMixin(store)
config = OIDCConfig(
issuer=provider.issuer,
client_id='client',
@@ -213,34 +210,35 @@ class OIDCAppTest(unittest.TestCase):
http_client=http_client,
fetch_userinfo=fetch_userinfo,
)
oidc_app = OIDCApp(session_app, config)
return oidc_app, provider, store
oidc = OIDCMixin(config, session=session)
app = KayaApp(mixins=[session, oidc])
return app, oidc, provider, store
def _setup_routes(self, oidc_app: OIDCApp) -> None:
@oidc_app.GET('/')
def _setup_routes(self, app: KayaApp, oidc: OIDCMixin) -> None:
@app.GET('/')
async def home(ctx: HttpContext) -> None:
await ctx.send_str(200, 'home')
@oidc_app.GET('/profile')
@oidc_app.require_auth
@app.GET('/profile')
@oidc.require_auth
async def profile(ctx: HttpContext) -> None:
user = oidc_app.get_user(ctx)
user = oidc.get_user(ctx)
if user is None:
await ctx.send_empty(401)
return
await ctx.send_str(200, f'Hello {user.email}')
@oidc_app.GET('/refresh')
@oidc_app.require_auth
@app.GET('/refresh')
@oidc.require_auth
async def refresh(ctx: HttpContext) -> None:
new_token = await oidc_app.refresh_access_token(ctx.session)
new_token = await oidc.refresh_access_token(ctx.session)
await ctx.send_str(200, new_token or 'no-token')
@async_test
async def test_login_redirect(self) -> None:
oidc_app, provider, store = self._build_app()
self._setup_routes(oidc_app)
transport = httpx.ASGITransport(app=oidc_app)
app, oidc, provider, store = self._build_app()
self._setup_routes(app, oidc)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
r = await client.get('/auth/login', follow_redirects=False)
self.assertEqual(302, r.status_code)
@@ -250,9 +248,9 @@ class OIDCAppTest(unittest.TestCase):
@async_test
async def test_callback_success(self) -> None:
oidc_app, provider, store = self._build_app(fetch_userinfo=True)
self._setup_routes(oidc_app)
transport = httpx.ASGITransport(app=oidc_app)
app, oidc, provider, store = self._build_app(fetch_userinfo=True)
self._setup_routes(app, oidc)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
r = await client.get('/auth/login', follow_redirects=False)
self.assertEqual(302, r.status_code)
@@ -274,18 +272,18 @@ class OIDCAppTest(unittest.TestCase):
@async_test
async def test_callback_invalid_state(self) -> None:
oidc_app, provider, store = self._build_app()
self._setup_routes(oidc_app)
transport = httpx.ASGITransport(app=oidc_app)
app, oidc, provider, store = self._build_app()
self._setup_routes(app, oidc)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
r = await client.get('/auth/callback', params={'code': 'mock-code', 'state': 'wrong'}, follow_redirects=False)
self.assertEqual(400, r.status_code)
@async_test
async def test_logout(self) -> None:
oidc_app, provider, store = self._build_app()
self._setup_routes(oidc_app)
transport = httpx.ASGITransport(app=oidc_app)
app, oidc, provider, store = self._build_app()
self._setup_routes(app, oidc)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
r = await client.get('/auth/login', follow_redirects=False)
parsed = urlparse(r.headers['Location'])
@@ -304,9 +302,9 @@ class OIDCAppTest(unittest.TestCase):
@async_test
async def test_require_auth_redirect(self) -> None:
oidc_app, provider, store = self._build_app()
self._setup_routes(oidc_app)
transport = httpx.ASGITransport(app=oidc_app)
app, oidc, provider, store = self._build_app()
self._setup_routes(app, oidc)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
r = await client.get('/profile', follow_redirects=False)
self.assertEqual(302, r.status_code)
@@ -314,9 +312,9 @@ class OIDCAppTest(unittest.TestCase):
@async_test
async def test_refresh_access_token(self) -> None:
oidc_app, provider, store = self._build_app()
self._setup_routes(oidc_app)
transport = httpx.ASGITransport(app=oidc_app)
app, oidc, provider, store = self._build_app()
self._setup_routes(app, oidc)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
r = await client.get('/auth/login', follow_redirects=False)
parsed = urlparse(r.headers['Location'])
@@ -329,6 +327,34 @@ class OIDCAppTest(unittest.TestCase):
self.assertEqual(200, r.status_code)
self.assertEqual('new-access-token', r.text)
@async_test
async def test_dependency_applied_automatically(self) -> None:
# Only pass oidc to KayaApp; SessionMixin should be applied via dependencies.
provider = MockOIDCProvider()
http_client = httpx.AsyncClient(transport=MockTransport(provider))
store = InMemorySessionStore()
session = SessionMixin(store)
config = OIDCConfig(
issuer=provider.issuer,
client_id='client',
client_secret='secret',
redirect_uri='http://localhost:8000/auth/callback',
http_client=http_client,
)
oidc = OIDCMixin(config, session=session)
app = KayaApp(mixins=[oidc])
@app.GET('/')
async def home(ctx: HttpContext) -> None:
ctx.session['x'] = 1
await ctx.send_str(200, 'ok')
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
r = await client.get('/')
self.assertEqual(200, r.status_code)
self.assertIn('Set-Cookie', r.headers)
if __name__ == '__main__':
unittest.main()