Author SHA1 Message Date
woggioni ea82ac44e9 Add reusable component schemas to kaya-openapi
CI / Build Pip package (push) Successful in 7m57s
Operations could only inline schemas: the generated document had no
components section, so #/components/... references dangled. OpenAPIMixin
and generate_spec now accept a components mapping (merged verbatim into
the document) and kaya.openapi exports schema_ref() to build
#/components/schemas/<name> references. Component references are
validated at generation time and unknown ones raise ValueError.
2026-10-14 20:37:14 +08:00
woggioni e73c4b2c18 Validate ID tokens against discovered issuer and refreshed tokens
CI / Build Pip package (push) Successful in 8m49s
OIDCClient.validate_id_token compared the iss claim against
OIDCConfig.issuer (the discovery URL) instead of the issuer advertised by
the provider metadata. This broke Microsoft Entra ID, whose v2 metadata
advertises the templated issuer
https://login.microsoftonline.com/{tenantid}/v2.0 while tokens carry the
concrete tenant GUID, and it also affected tenant-domain issuers.

- Validate iss against the discovered metadata issuer and substitute
  {tenantid} with the token's tid claim.
- Catch jwt.PyJWTError in the OIDC callback so invalid tokens return 400
  instead of 500 (InvalidIssuerError is not a ValueError).
- Re-enable aud/exp/iat/nbf checks for unsigned (alg: none) ID tokens.
- Validate ID tokens returned by the refresh grant before storing them,
  including nonce continuity when present and sub continuity with the
  session user.
- Document Entra ID support and update tests.
2026-10-01 23:22:00 +08:00
woggioni fa8bb0487d Add Pushed Authorization Request (RFC 9126) support
CI / Build Pip package (push) Successful in 14m8s
OIDCConfig.par selects never (default), auto, or always. When enabled, build_authorization_url pushes the authorization parameters to the provider's PAR endpoint and redirects the browser with only client_id and request_uri. Providers that require PAR raise a clear error when configured with par='never'. Client authentication follows the token endpoint (HTTP Basic or public client).
2026-10-01 08:23:40 +08:00
woggioni 15dc4b6dad Add single/multi identity strategy to OIDC providers
OIDCProviders and OIDCMixin accept strategy='single-identity' (default) or 'multi-identity'. Single-identity evicts other providers' identities and tokens after a successful login so the session holds exactly one user, matching Spring Security's last-login-wins behavior. Ephemeral flow keys are preserved so in-flight logins can still complete.
2026-10-01 08:23:40 +08:00
woggioni 3af2e7a975 Add support for multiple OIDC providers
Namespace per-provider session state as oidc:{provider}:... so several providers can share one session, and add an OIDCProviders mixin that wires multiple configs onto distinct /auth/{provider}/... routes. Logging out of one provider no longer signs the user out of the others.
2026-10-01 08:23:40 +08:00
woggioni 540a7b111e Add standard attributes to kaya-otel websocket spans 2026-09-30 08:58:06 +08:00
woggioni 7c46badebb Surface vanished websocket peers as WebSocketDisconnect
CI / Build Pip package (push) Successful in 4m55s
Granian's RSGI transport raises RSGIProtocolClosed when the peer is gone
without a close handshake, which escaped receive() and was logged as an
application error, and sends to the dead transport raised the raw Granian
exception as well.

Treat a closed RSGI transport like the ASGI websocket.disconnect message:
receive() reports a close message so async-for loops end normally, while
writes (accept/send/close) raise a new kaya.core.WebSocketDisconnect
carrying close code 1006. KayaApp.handle_websocket swallows that exception
as normal termination, so after hooks still run and disconnects are no
longer reported as errors. The ASGI adapter maps its server's OSError
disconnect (e.g. uvicorn's ClientDisconnected) to the same exception for
consistent semantics across backends.
2026-09-29 22:10:16 +08:00
23 changed files with 2213 additions and 115 deletions

No files matched your search

+2 -1
View File
@@ -4,7 +4,7 @@ from ._http_context import HttpContext
from ._mixin import KayaMixin from ._mixin import KayaMixin
from ._tree import Tree, PathIterator from ._tree import Tree, PathIterator
from ._path_handler import PathHandler, Matches from ._path_handler import PathHandler, Matches
from ._websocket import WebSocket, WebSocketMessage from ._websocket import WebSocket, WebSocketDisconnect, WebSocketMessage
__all__ = [ __all__ = [
@@ -18,5 +18,6 @@ __all__ = [
'Matches', 'Matches',
'PathIterator', 'PathIterator',
'WebSocket', 'WebSocket',
'WebSocketDisconnect',
'WebSocketMessage' 'WebSocketMessage'
] ]
+7 -1
View File
@@ -9,7 +9,7 @@ from ._http_method import HttpMethod
from ._mixin import KayaMixin from ._mixin import KayaMixin
from ._path_handler import Context from ._path_handler import Context
from ._types import StrOrStrings from ._types import StrOrStrings
from ._websocket import WebSocket, WebSocketMessage from ._websocket import WebSocket, WebSocketDisconnect, WebSocketMessage
from ._asgi import AsgiContext, AsgiWebSocket from ._asgi import AsgiContext, AsgiWebSocket
from ._tree import Tree from ._tree import Tree
from ._types.asgi import LifespanScope, HTTPScope as ASGIHTTPScope, WebSocketScope as ASGIWebSocketScope from ._types.asgi import LifespanScope, HTTPScope as ASGIHTTPScope, WebSocketScope as ASGIWebSocketScope
@@ -164,6 +164,12 @@ class KayaApp(AbstractKayaApp):
await handler.handle_request(ws, captured) await handler.handle_request(ws, captured)
else: else:
await ws.close(1000) await ws.close(1000)
except WebSocketDisconnect:
# The peer went away while the handler was using the websocket
# (e.g. sending on a vanished transport). This is normal
# termination, not an application error, so it is not recorded
# on the socket nor re-raised.
pass
except Exception as exc: except Exception as exc:
ws.exception = exc ws.exception = exc
raise raise
+13 -4
View File
@@ -17,7 +17,7 @@ from pwo import Maybe
from pathlib import Path from pathlib import Path
from ._http_method import HttpMethod from ._http_method import HttpMethod
from ._http_context import HttpContext from ._http_context import HttpContext
from ._websocket import WebSocket, WebSocketMessage from ._websocket import WebSocket, WebSocketDisconnect, WebSocketMessage
from ._types import StrOrStrings from ._types import StrOrStrings
from ._types.asgi import HTTPScope, WebSocketScope from ._types.asgi import HTTPScope, WebSocketScope
@@ -179,7 +179,16 @@ class AsgiWebSocket(WebSocket):
# Emit a list rather than a tuple: wsproto's AcceptConnection (used # Emit a list rather than a tuple: wsproto's AcceptConnection (used
# by httpx_ws and others) requires list concatenation. # by httpx_ws and others) requires list concatenation.
message['headers'] = list(encode_headers(headers)) message['headers'] = list(encode_headers(headers))
await self._send_message(message)
async def _send_message(self, message: Mapping[str, Any]) -> None:
try:
await self._send(message) await self._send(message)
except OSError as exc:
# ASGI servers (e.g. uvicorn) report a vanished peer with an
# OSError subclass; surface it as a kaya-level disconnect so
# handlers do not need to know about server internals.
raise WebSocketDisconnect(code=1006) from exc
async def receive(self) -> WebSocketMessage: async def receive(self) -> WebSocketMessage:
message = await self._receive() message = await self._receive()
@@ -200,13 +209,13 @@ class AsgiWebSocket(WebSocket):
return WebSocketMessage(kind='close') return WebSocketMessage(kind='close')
async def send_text(self, data: str) -> None: async def send_text(self, data: str) -> None:
await self._send({'type': 'websocket.send', 'text': data}) await self._send_message({'type': 'websocket.send', 'text': data})
async def send_bytes(self, data: bytes) -> None: async def send_bytes(self, data: bytes) -> None:
await self._send({'type': 'websocket.send', 'bytes': data}) await self._send_message({'type': 'websocket.send', 'bytes': data})
async def close(self, code: int = 1000) -> None: async def close(self, code: int = 1000) -> None:
await self._send({'type': 'websocket.close', 'code': code}) await self._send_message({'type': 'websocket.close', 'code': code})
async def __anext__(self) -> WebSocketMessage: async def __anext__(self) -> WebSocketMessage:
message = await self.receive() message = await self.receive()
@@ -13,6 +13,41 @@ class WebSocketMessage:
data: WebSocketData = None data: WebSocketData = None
class WebSocketDisconnect(Exception):
"""Raised when using a websocket whose peer has disconnected.
Concrete websocket implementations raise this from ``send_text()`` and
``send_bytes()`` when the transport reports that the peer is gone (for
example when a browser tab is closed without a close handshake), so
handlers can catch it and stop producing output:
.. code-block:: python
try:
await ws.send_text('hello')
except WebSocketDisconnect:
pass
Reading from a disconnected websocket is not an error: ``receive()``
reports a ``close`` message and iterating with ``async for`` simply ends,
which is what most handlers should rely on instead.
``code`` is a websocket close code; ``1006`` means the connection was
closed abnormally, without a close handshake.
"""
code: int
reason: str
def __init__(self, code: int = 1006, reason: str = '') -> None:
self.code = code
self.reason = reason
message = f'WebSocket disconnected (code={code})'
if reason:
message = f'{message}: {reason}'
super().__init__(message)
class WebSocket(ABC): class WebSocket(ABC):
path: str path: str
query_string: str query_string: str
+113 -3
View File
@@ -1,9 +1,9 @@
import unittest import unittest
import httpx import httpx
from pwo import async_test from pwo import async_test
from httpx_ws import aconnect_ws, WebSocketDisconnect from httpx_ws import aconnect_ws, WebSocketDisconnect as HttpxWebSocketDisconnect
from httpx_ws.transport import ASGIWebSocketTransport from httpx_ws.transport import ASGIWebSocketTransport
from kaya.core import KayaApp, WebSocket from kaya.core import KayaApp, WebSocket, WebSocketDisconnect, WebSocketMessage
from kaya.core._asgi import AsgiWebSocket from kaya.core._asgi import AsgiWebSocket
@@ -66,7 +66,7 @@ class WebSocketTest(unittest.TestCase):
transport = ASGIWebSocketTransport(app=self.app) transport = ASGIWebSocketTransport(app=self.app)
async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
with self.assertRaises(WebSocketDisconnect) as cm: with self.assertRaises(HttpxWebSocketDisconnect) as cm:
async with aconnect_ws("/unknown", client) as ws: async with aconnect_ws("/unknown", client) as ws:
pass pass
@@ -151,3 +151,113 @@ class WebSocketTest(unittest.TestCase):
} }
ws = AsgiWebSocket(scope, receive, send) ws = AsgiWebSocket(scope, receive, send)
self.assertEqual('ws', ws.scheme) self.assertEqual('ws', ws.scheme)
class FakeWebSocket(WebSocket):
"""Minimal in-memory websocket for exercising ``KayaApp`` directly."""
def __init__(self, path: str) -> None:
self.path = path
self.query_string = ''
self.scheme = 'ws'
self.headers: dict[str, list[str]] = {}
self.client = ('127.0.0.1', 12345)
self.server = ('127.0.0.1', 80)
self.closed: list[int] = []
async def accept(self, headers=None) -> None:
pass
async def receive(self) -> WebSocketMessage:
return WebSocketMessage(kind='close')
async def send_text(self, data: str) -> None:
pass
async def send_bytes(self, data: bytes) -> None:
pass
async def close(self, code: int = 1000) -> None:
self.closed.append(code)
async def __anext__(self) -> WebSocketMessage:
raise StopAsyncIteration
class WebSocketDisconnectTest(unittest.TestCase):
def setUp(self):
self.app = KayaApp()
@self.app.websocket('/disconnect')
async def disconnect(ws: WebSocket) -> None:
await ws.accept()
await ws.send_text('before')
raise WebSocketDisconnect()
@self.app.websocket('/failure')
async def failure(ws: WebSocket) -> None:
await ws.accept()
raise RuntimeError('boom')
@async_test
async def test_disconnect_terminates_quietly(self):
after_hook_calls = []
async def after(ws: WebSocket) -> None:
after_hook_calls.append(ws)
self.app.add_after_websocket_hook(after)
ws = FakeWebSocket('/disconnect')
await self.app.handle_websocket(ws)
self.assertEqual([ws], after_hook_calls)
self.assertIsNone(ws.exception)
@async_test
async def test_other_errors_are_recorded_and_reraised(self):
ws = FakeWebSocket('/failure')
with self.assertRaises(RuntimeError):
await self.app.handle_websocket(ws)
self.assertIsInstance(ws.exception, RuntimeError)
@async_test
async def test_asgi_send_after_disconnect_raises_websocket_disconnect(self):
async def send(message):
raise OSError('client gone')
async def receive():
return {'type': 'websocket.disconnect', 'code': 1006}
ws = AsgiWebSocket(self._scope(), receive, send)
with self.assertRaises(WebSocketDisconnect) as ctx:
await ws.send_text('hello')
self.assertEqual(1006, ctx.exception.code)
with self.assertRaises(WebSocketDisconnect):
await ws.send_bytes(b'hello')
with self.assertRaises(WebSocketDisconnect):
await ws.close()
@async_test
async def test_asgi_accept_after_disconnect_raises_websocket_disconnect(self):
async def send(message):
raise OSError('client gone')
async def receive():
return {'type': 'websocket.connect'}
ws = AsgiWebSocket(self._scope(), receive, send)
with self.assertRaises(WebSocketDisconnect) as ctx:
await ws.accept()
self.assertEqual(1006, ctx.exception.code)
@staticmethod
def _scope() -> dict:
return {
'type': 'websocket',
'path': '/echo',
'query_string': b'',
'scheme': 'ws',
'client': ('127.0.0.1', 12345),
'server': ('127.0.0.1', 80),
'headers': [],
}
+154
View File
@@ -40,20 +40,174 @@ async def profile(ctx: HttpContext):
`OIDCMixin` depends on `SessionMixin`; passing only `oidc` to `KayaApp(mixins=...)` `OIDCMixin` depends on `SessionMixin`; passing only `oidc` to `KayaApp(mixins=...)`
also works because the app applies mixin dependencies automatically. also works because the app applies mixin dependencies automatically.
## Multiple providers
Use `OIDCProviders` to run several OIDC providers in one application. Each
config's `provider` becomes a route segment under `base_path`
(`/auth/{provider}/login`, `/auth/{provider}/callback` and
`/auth/{provider}/logout`), and all providers share a single `SessionMixin`:
```python
from kaya.oidc import OIDCConfig, OIDCProviders
session = SessionMixin(InMemorySessionStore())
oidc = OIDCProviders(
[
OIDCConfig(
issuer='https://accounts.google.com',
client_id=os.environ['GOOGLE_CLIENT_ID'],
client_secret=os.environ['GOOGLE_CLIENT_SECRET'],
redirect_uri='http://localhost:8000/auth/google/callback',
provider='google',
),
OIDCConfig(
issuer='https://github.com/login/oauth', # or any OIDC issuer
client_id=os.environ['GITHUB_CLIENT_ID'],
client_secret=os.environ['GITHUB_CLIENT_SECRET'],
redirect_uri='http://localhost:8000/auth/github/callback',
provider='github',
),
],
session=session,
base_path='/auth', # optional
)
app = KayaApp(mixins=[session, oidc])
@app.GET('/google-profile')
@oidc.require_auth('google')
async def google_profile(ctx: HttpContext):
user = oidc.get_user(ctx, 'google')
await ctx.send_str(200, f"Hello {user.email or user.sub}")
@app.GET('/who')
async def who(ctx: HttpContext):
users = oidc.get_users(ctx) # {provider: OIDCUser}
providers = oidc.logged_in_providers(ctx) # ['google', ...]
await ctx.send_str(200, ', '.join(users))
```
Each provider's session state is namespaced as `oidc:{provider}:...`. Make sure
each config's `redirect_uri` points at that provider's derived callback path.
### Identity strategy
`strategy` controls how many identities may coexist in the shared session:
- `'single-identity'` (default): completing a login evicts every other
provider's identity and tokens, so the session holds exactly one user — last
login wins, like Spring Security. After a GitHub login, `get_user(ctx,
'google')` returns `None` and `require_auth('google')` redirects to the Google
login.
- `'multi-identity'`: each provider stays logged in independently. Use it for
linked accounts, aggregating several provider APIs, or step-up flows. Logging
out of one provider keeps the shared session alive while another provider is
still logged in.
```python
oidc = OIDCProviders([google_config, github_config], session=session,
strategy='multi-identity')
```
Ephemeral flow keys (`state`, `nonce`, `code_verifier`) are never evicted, so
logins already in flight can still complete; under `'single-identity'` the last
callback to finish wins. Under either strategy, `get_users(ctx)` only reports
providers actually present in the session.
The composable alternative is to create one `OIDCMixin` per provider, pointing
each at distinct `login_path`/`callback_path`/`logout_path` values and sharing
the same `SessionMixin`. The same `strategy` keyword is accepted and defaults to
`'single-identity'`:
```python
google = OIDCMixin(google_config, session=session)
github = OIDCMixin(github_config, session=session, strategy='multi-identity')
app = KayaApp(mixins=[session, google, github])
```
> **Upgrading:** session keys are now prefixed with `oidc:{provider}:` (the
> default provider is `default`). Existing logged-in users are signed out once
> and simply re-authenticate.
## Pushed authorization requests (PAR)
`OIDCConfig.par` controls RFC 9126 Pushed Authorization Requests:
- `'never'` (default): the authorization request travels as query parameters,
as before.
- `'auto'`: push the request whenever the provider advertises a
`pushed_authorization_request_endpoint`; otherwise fall back to the plain
redirect.
- `'always'`: always push; login raises a `RuntimeError` if the provider does
not support PAR.
```python
OIDCConfig(..., par='auto')
```
A provider that advertises `require_pushed_authorization_requests: true` must be
configured with `par='auto'` or `par='always'`. With `par='never'` the login
fails fast with a clear error instead of issuing a redirect the provider would
reject.
The same parameters as the plain authorization request (PKCE challenge,
`state`, `nonce`, scopes, ...) are POSTed directly to the provider, and the
browser is redirected with only `client_id` and the single-use `request_uri`.
Request parameters therefore no longer travel through the user agent. PKCE,
`state`, and `nonce` are still used and protect against request URI swapping.
Client authentication follows the token endpoint: HTTP Basic when
`client_secret` is set, or none for public clients. PAR with other client
authentication methods (`private_key_jwt`, mTLS) is not supported.
## Features ## Features
- Generic OIDC discovery - Generic OIDC discovery
- Authorization Code Flow with PKCE (S256) - Authorization Code Flow with PKCE (S256)
- ID token signature validation with JWKS - ID token signature validation with JWKS
- Issuer validation against discovery metadata, including support for
Microsoft Entra ID `{tenantid}` issuer templates
- Refreshed ID token validation on token refresh
- `state` and `nonce` protection - `state` and `nonce` protection
- Session fixation defense via `regenerate_id()` after login - Session fixation defense via `regenerate_id()` after login
- Optional userinfo endpoint fetch - Optional userinfo endpoint fetch
- Refresh token support - Refresh token support
- RP-initiated logout (when provider advertises `end_session_endpoint`) - RP-initiated logout (when provider advertises `end_session_endpoint`)
- Pushed Authorization Requests (RFC 9126) via `OIDCConfig.par`
- Multiple providers in one app via `OIDCProviders` (or several `OIDCMixin` instances)
- `single-identity` (default) or `multi-identity` session strategy
- Composable with any other `KayaMixin` (RSGI, MCP, etc.) - Composable with any other `KayaMixin` (RSGI, MCP, etc.)
## Microsoft Entra ID
Entra ID v2 ID tokens carry an `iss` claim of the form
`https://login.microsoftonline.com/<tenant-id>/v2.0`, while the discovery
metadata for the multi-tenant endpoints advertises the templated issuer
`https://login.microsoftonline.com/{tenantid}/v2.0`. `kaya-oidc` validates the
`iss` claim against the discovered metadata and substitutes the token's `tid`
claim into `{tenantid}`, so both tenant-specific and multi-tenant issuers work:
```python
OIDCConfig(
issuer='https://login.microsoftonline.com/<tenant-id-or-domain>/v2.0',
client_id=os.environ['OIDC_CLIENT_ID'],
client_secret=os.environ['OIDC_CLIENT_SECRET'],
redirect_uri='https://example.com/auth/callback',
)
```
Use a tenant-specific issuer (the tenant GUID or a verified domain name) to
restrict logins to that tenant. The `common` and `organizations` aliases are
multi-tenant: any tenant in that cloud can obtain an ID token that passes
validation. If you build a multi-tenant application, check the `tid` claim
against the tenants you trust.
## Security notes ## Security notes
- The `iss` claim is validated against the `issuer` value of the provider's
discovery document, not against `OIDCConfig.issuer` (they can differ, for
example for Entra ID multi-tenant endpoints).
- Refreshed ID tokens are validated before being stored; the `nonce` is also
checked when the refreshed token carries one.
- The `none` signing algorithm is rejected by default. - The `none` signing algorithm is rejected by default.
- Only algorithms listed in `OIDCConfig.allowed_id_token_algorithms` are accepted. - Only algorithms listed in `OIDCConfig.allowed_id_token_algorithms` are accepted.
- Always use HTTPS for `redirect_uri` in production. - Always use HTTPS for `redirect_uri` in production.
+6 -2
View File
@@ -1,11 +1,15 @@
from ._client import OIDCClient from ._client import OIDCClient
from ._config import OIDCConfig from ._config import OIDCConfig, PARPolicy
from ._mixin import OIDCMixin, OIDCUser from ._mixin import IdentityStrategy, OIDCMixin, OIDCUser
from ._providers import OIDCProviders
__all__ = [ __all__ = [
'IdentityStrategy',
'OIDCClient', 'OIDCClient',
'OIDCConfig', 'OIDCConfig',
'OIDCMixin', 'OIDCMixin',
'OIDCProviders',
'OIDCUser', 'OIDCUser',
'PARPolicy',
] ]
+130 -17
View File
@@ -1,3 +1,4 @@
import re
from typing import Any, Mapping, Optional, Sequence, cast from typing import Any, Mapping, Optional, Sequence, cast
from urllib.parse import urlencode from urllib.parse import urlencode
@@ -5,9 +6,11 @@ import httpx
import jwt import jwt
from jwt import PyJWK from jwt import PyJWK
from ._config import OIDCConfig from ._config import PARPolicy, OIDCConfig
from ._utils import generate_nonce, generate_pkce, generate_state from ._utils import generate_nonce, generate_pkce, generate_state
_TENANT_ID_PLACEHOLDER = re.compile(r'\{tenantid\}', re.IGNORECASE)
class OIDCClient: class OIDCClient:
"""Low-level OIDC client implementing discovery, PKCE, and token validation.""" """Low-level OIDC client implementing discovery, PKCE, and token validation."""
@@ -42,8 +45,14 @@ class OIDCClient:
async with httpx.AsyncClient() as client: async with httpx.AsyncClient() as client:
return await client.request(method, url, **kwargs) return await client.request(method, url, **kwargs)
async def build_authorization_url(self) -> tuple[str, str, str, str]: async def build_authorization_url(self, par: PARPolicy = 'never') -> tuple[str, str, str, str]:
"""Return (authorization_url, state, nonce, code_verifier).""" """Return (authorization_url, state, nonce, code_verifier).
When ``par`` selects Pushed Authorization Requests (RFC 9126), the
authorization request parameters are POSTed to the provider's PAR
endpoint first and the returned URL carries only ``client_id`` and
``request_uri``.
"""
metadata = await self._ensure_metadata() metadata = await self._ensure_metadata()
authorization_endpoint = metadata.get('authorization_endpoint') authorization_endpoint = metadata.get('authorization_endpoint')
if not isinstance(authorization_endpoint, str): if not isinstance(authorization_endpoint, str):
@@ -54,7 +63,6 @@ class OIDCClient:
code_verifier, code_challenge = generate_pkce() code_verifier, code_challenge = generate_pkce()
params = { params = {
'client_id': self._config.client_id,
'response_type': 'code', 'response_type': 'code',
'scope': ' '.join(self._config.scopes), 'scope': ' '.join(self._config.scopes),
'redirect_uri': self._config.redirect_uri, 'redirect_uri': self._config.redirect_uri,
@@ -63,9 +71,67 @@ class OIDCClient:
'code_challenge': code_challenge, 'code_challenge': code_challenge,
'code_challenge_method': 'S256', 'code_challenge_method': 'S256',
} }
url = authorization_endpoint + '?' + urlencode(params)
if self._use_par(metadata, par):
request_uri = await self._push_authorization_request(metadata, params)
redirect_params = {'client_id': self._config.client_id, 'request_uri': request_uri}
else:
redirect_params = {'client_id': self._config.client_id, **params}
url = authorization_endpoint + '?' + urlencode(redirect_params)
return url, state, nonce, code_verifier return url, state, nonce, code_verifier
def _use_par(self, metadata: Mapping[str, Any], policy: PARPolicy) -> bool:
endpoint = metadata.get('pushed_authorization_request_endpoint')
required = metadata.get('require_pushed_authorization_requests') is True
if required:
if not isinstance(endpoint, str):
raise RuntimeError(
'OIDC provider requires pushed authorization requests but does not '
'advertise a pushed_authorization_request_endpoint'
)
if policy == 'never':
raise RuntimeError(
"OIDC provider requires pushed authorization requests; set par to "
"'auto' or 'always'"
)
return True
if policy == 'always':
if not isinstance(endpoint, str):
raise RuntimeError(
'OIDC provider does not advertise a pushed_authorization_request_endpoint'
)
return True
return policy == 'auto' and isinstance(endpoint, str)
async def _push_authorization_request(
self, metadata: Mapping[str, Any], params: Mapping[str, str]
) -> str:
endpoint = metadata.get('pushed_authorization_request_endpoint')
if not isinstance(endpoint, str):
raise RuntimeError(
'OIDC provider does not advertise a pushed_authorization_request_endpoint'
)
# RFC 9126 section 2.1: client_id is required in the body even when the
# client authenticates with the Authorization header.
data = {'client_id': self._config.client_id, **params}
resp = await self._request('POST', endpoint, data=data, auth=self._client_auth())
resp.raise_for_status()
body = resp.json()
request_uri = body.get('request_uri') if isinstance(body, Mapping) else None
if not isinstance(request_uri, str):
raise RuntimeError('PAR endpoint response is missing request_uri')
return request_uri
def _client_auth(self) -> Optional[tuple[str, str]]:
if self._config.client_secret is None:
return None
return (self._config.client_id, self._config.client_secret)
async def fetch_token(self, code: str, code_verifier: str) -> Mapping[str, Any]: async def fetch_token(self, code: str, code_verifier: str) -> Mapping[str, Any]:
metadata = await self._ensure_metadata() metadata = await self._ensure_metadata()
token_endpoint = metadata.get('token_endpoint') token_endpoint = metadata.get('token_endpoint')
@@ -79,11 +145,8 @@ class OIDCClient:
'client_id': self._config.client_id, 'client_id': self._config.client_id,
'code_verifier': code_verifier, 'code_verifier': code_verifier,
} }
auth: Optional[tuple[str, str]] = None
if self._config.client_secret is not None:
auth = (self._config.client_id, self._config.client_secret)
resp = await self._request('POST', token_endpoint, data=data, auth=auth) resp = await self._request('POST', token_endpoint, data=data, auth=self._client_auth())
resp.raise_for_status() resp.raise_for_status()
return cast(Mapping[str, Any], resp.json()) return cast(Mapping[str, Any], resp.json())
@@ -98,15 +161,35 @@ class OIDCClient:
'refresh_token': refresh_token, 'refresh_token': refresh_token,
'client_id': self._config.client_id, 'client_id': self._config.client_id,
} }
auth: Optional[tuple[str, str]] = None
if self._config.client_secret is not None:
auth = (self._config.client_id, self._config.client_secret)
resp = await self._request('POST', token_endpoint, data=data, auth=auth) resp = await self._request('POST', token_endpoint, data=data, auth=self._client_auth())
resp.raise_for_status() resp.raise_for_status()
return cast(Mapping[str, Any], resp.json()) return cast(Mapping[str, Any], resp.json())
async def validate_id_token(self, id_token: str, nonce: str) -> Mapping[str, Any]: async def validate_id_token(self, id_token: str, nonce: str) -> Mapping[str, Any]:
payload = await self._decode_id_token(id_token)
if payload.get('nonce') != nonce:
raise ValueError('ID token nonce mismatch')
return payload
async def validate_refreshed_id_token(
self,
id_token: str,
expected_nonce: Optional[str] = None,
) -> Mapping[str, Any]:
"""Validate an ID token returned by the refresh token grant.
Unlike the authorization code flow, the provider may omit the ``nonce``
claim; when it is present and the original nonce is known, the two must
match (OIDC Core 12.2).
"""
payload = await self._decode_id_token(id_token)
token_nonce = payload.get('nonce')
if token_nonce is not None and expected_nonce is not None and token_nonce != expected_nonce:
raise ValueError('ID token nonce mismatch')
return payload
async def _decode_id_token(self, id_token: str) -> Mapping[str, Any]:
header = jwt.get_unverified_header(id_token) header = jwt.get_unverified_header(id_token)
alg = header.get('alg') alg = header.get('alg')
if not isinstance(alg, str): if not isinstance(alg, str):
@@ -115,29 +198,59 @@ class OIDCClient:
if alg == 'none': if alg == 'none':
if not self._config.allow_unsigned_id_tokens: if not self._config.allow_unsigned_id_tokens:
raise ValueError('Unsigned ID tokens are not allowed') raise ValueError('Unsigned ID tokens are not allowed')
metadata_issuer = await self._get_metadata_issuer()
payload = jwt.decode( payload = jwt.decode(
id_token, id_token,
options={'verify_signature': False}, audience=self._config.client_id,
options={
'verify_signature': False,
'verify_aud': True,
'verify_exp': True,
'verify_iat': True,
'verify_nbf': True,
},
) )
else: else:
if alg not in self._config.allowed_id_token_algorithms: if alg not in self._config.allowed_id_token_algorithms:
raise ValueError(f'ID token algorithm {alg} is not allowed') raise ValueError(f'ID token algorithm {alg} is not allowed')
metadata_issuer = await self._get_metadata_issuer()
jwks = await self._ensure_jwks() jwks = await self._ensure_jwks()
signing_key = self._get_signing_key(jwks, header.get('kid'), alg) signing_key = self._get_signing_key(jwks, header.get('kid'), alg)
payload = jwt.decode( payload = jwt.decode(
id_token, id_token,
signing_key.key, signing_key.key,
algorithms=[alg], algorithms=[alg],
issuer=self._config.issuer,
audience=self._config.client_id, audience=self._config.client_id,
) )
if not isinstance(payload, dict): if not isinstance(payload, dict):
raise ValueError('ID token payload is not a JSON object') raise ValueError('ID token payload is not a JSON object')
if payload.get('nonce') != nonce: self._validate_issuer(payload, metadata_issuer)
raise ValueError('ID token nonce mismatch')
return payload return payload
async def _get_metadata_issuer(self) -> str:
metadata = await self._ensure_metadata()
issuer = metadata.get('issuer')
if not isinstance(issuer, str):
raise RuntimeError('OIDC provider does not advertise an issuer')
return issuer
@staticmethod
def _validate_issuer(payload: Mapping[str, Any], metadata_issuer: str) -> None:
iss = payload.get('iss')
if not isinstance(iss, str):
raise jwt.MissingRequiredClaimError('iss')
expected = metadata_issuer
if _TENANT_ID_PLACEHOLDER.search(expected):
tenant_id = payload.get('tid')
if not isinstance(tenant_id, str):
raise jwt.InvalidIssuerError('Invalid issuer')
expected = _TENANT_ID_PLACEHOLDER.sub(tenant_id, expected, count=1)
if iss != expected:
raise jwt.InvalidIssuerError('Invalid issuer')
def _get_signing_key(self, jwks: Mapping[str, Any], kid: Optional[str], alg: str) -> PyJWK: def _get_signing_key(self, jwks: Mapping[str, Any], kid: Optional[str], alg: str) -> PyJWK:
keys = jwks.get('keys') keys = jwks.get('keys')
if not isinstance(keys, Sequence): if not isinstance(keys, Sequence):
+32 -1
View File
@@ -1,13 +1,31 @@
import re
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Any, Optional, Sequence from typing import Any, Literal, Optional, Sequence
import httpx import httpx
_PROVIDER_PATTERN = re.compile(r'^[A-Za-z0-9_-]+$')
type PARPolicy = Literal['never', 'auto', 'always']
_PAR_POLICIES: tuple[PARPolicy, ...] = ('never', 'auto', 'always')
@dataclass @dataclass
class OIDCConfig: class OIDCConfig:
"""Configuration for an OpenID Connect provider. """Configuration for an OpenID Connect provider.
``provider`` is a stable identifier for this provider. It namespaces the
session state (``oidc:{provider}:...``) so several providers can be used in
the same application without interfering with each other.
``par`` selects when to use Pushed Authorization Requests (RFC 9126):
``'never'`` uses a plain query-string authorization request, ``'auto'``
pushes the request whenever the provider advertises a
``pushed_authorization_request_endpoint``, and ``'always'`` requires PAR.
A provider that sets ``require_pushed_authorization_requests`` is always
treated as requiring PAR.
``client_secret`` is optional; omit it for public clients. ``client_secret`` is optional; omit it for public clients.
``fetch_userinfo`` controls whether the userinfo endpoint is queried after ``fetch_userinfo`` controls whether the userinfo endpoint is queried after
token exchange. token exchange.
@@ -35,8 +53,21 @@ class OIDCConfig:
'EdDSA', 'EdDSA',
)) ))
http_client: Optional[httpx.AsyncClient] = None http_client: Optional[httpx.AsyncClient] = None
provider: str = 'default'
par: PARPolicy = 'never'
def __post_init__(self) -> None: def __post_init__(self) -> None:
if not self.provider or _PROVIDER_PATTERN.match(self.provider) is None:
raise ValueError(
'provider must be a non-empty string of letters, digits, '
"underscores or hyphens (got "
f'{self.provider!r})'
)
if self.par not in _PAR_POLICIES:
raise ValueError(
f'par must be one of {", ".join(map(repr, _PAR_POLICIES))} '
f'(got {self.par!r})'
)
scopes_list = list(self.scopes) scopes_list = list(self.scopes)
if 'openid' not in scopes_list: if 'openid' not in scopes_list:
self.scopes = ('openid', *scopes_list) self.scopes = ('openid', *scopes_list)
+142 -27
View File
@@ -1,5 +1,6 @@
from typing import Any, Awaitable, Callable, Mapping, Optional, Sequence from typing import Any, Awaitable, Callable, Literal, Mapping, Optional, Sequence, cast
import jwt
from kaya.core import HttpContext, KayaApp, KayaMixin from kaya.core import HttpContext, KayaApp, KayaMixin
from kaya.session import Session, SessionMixin from kaya.session import Session, SessionMixin
from urllib.parse import parse_qs from urllib.parse import parse_qs
@@ -10,6 +11,13 @@ from ._config import OIDCConfig
type HttpHandler = Callable[..., Awaitable[None]] type HttpHandler = Callable[..., Awaitable[None]]
type IdentityStrategy = Literal['single-identity', 'multi-identity']
_SESSION_NAMESPACE = 'oidc'
_USER_SUFFIX = 'user'
_IDENTITY_FIELDS = frozenset({'user', 'id_token', 'access_token', 'refresh_token'})
_IDENTITY_STRATEGIES: tuple[IdentityStrategy, ...] = ('single-identity', 'multi-identity')
class OIDCUser(Mapping[str, Any]): class OIDCUser(Mapping[str, Any]):
"""Read-only view of the OIDC userinfo / ID token claims.""" """Read-only view of the OIDC userinfo / ID token claims."""
@@ -28,7 +36,7 @@ class OIDCUser(Mapping[str, Any]):
@property @property
def sub(self) -> str: def sub(self) -> str:
return self._data['sub'] return cast(str, self._data['sub'])
@property @property
def email(self) -> Optional[str]: def email(self) -> Optional[str]:
@@ -56,6 +64,17 @@ class OIDCMixin(KayaMixin):
- ``callback_path`` (default ``/auth/callback``): handles the provider callback. - ``callback_path`` (default ``/auth/callback``): handles the provider callback.
- ``logout_path`` (default ``/auth/logout``): logs the user out. - ``logout_path`` (default ``/auth/logout``): logs the user out.
Session state is namespaced by ``OIDCConfig.provider`` as
``oidc:{provider}:...`` so several providers can share one session. Use one
``OIDCMixin`` per provider (with distinct route paths), or
:class:`~kaya.oidc.OIDCProviders` to wire several at once.
``strategy`` controls how many identities may coexist in the session. The
default ``'single-identity'`` evicts every other provider's identity after a
successful login, so the session has exactly one user (last login wins).
``'multi-identity'`` keeps each provider logged in independently, which is
useful for linked accounts, provider API aggregation, and step-up flows.
Example:: Example::
session = SessionMixin(InMemorySessionStore()) session = SessionMixin(InMemorySessionStore())
@@ -69,10 +88,33 @@ class OIDCMixin(KayaMixin):
... ...
""" """
def __init__(self, config: OIDCConfig, session: SessionMixin) -> None: def __init__(
self,
config: OIDCConfig,
session: SessionMixin,
strategy: IdentityStrategy = 'single-identity',
) -> None:
if strategy not in _IDENTITY_STRATEGIES:
raise ValueError(
'strategy must be one of '
f'{", ".join(map(repr, _IDENTITY_STRATEGIES))} (got {strategy!r})'
)
self._config = config self._config = config
self._session = session self._session = session
self._client = OIDCClient(config) self._client = OIDCClient(config)
self._strategy = strategy
self._prefix = f'{_SESSION_NAMESPACE}:{config.provider}:'
@property
def provider(self) -> str:
return self._config.provider
@property
def strategy(self) -> IdentityStrategy:
return self._strategy
def _key(self, name: str) -> str:
return f'{self._prefix}{name}'
@property @property
def dependencies(self) -> Sequence[KayaMixin]: def dependencies(self) -> Sequence[KayaMixin]:
@@ -81,11 +123,13 @@ class OIDCMixin(KayaMixin):
def apply(self, app: KayaApp) -> None: def apply(self, app: KayaApp) -> None:
@app.GET(self._config.login_path) @app.GET(self._config.login_path)
async def login(ctx: HttpContext) -> None: async def login(ctx: HttpContext) -> None:
auth_url, state, nonce, code_verifier = await self._client.build_authorization_url() auth_url, state, nonce, code_verifier = await self._client.build_authorization_url(
par=self._config.par
)
session = self._session_of(ctx) session = self._session_of(ctx)
session['oidc_state'] = state session[self._key('state')] = state
session['oidc_nonce'] = nonce session[self._key('nonce')] = nonce
session['oidc_code_verifier'] = code_verifier session[self._key('code_verifier')] = code_verifier
await ctx.send_empty(302, {'Location': auth_url}) await ctx.send_empty(302, {'Location': auth_url})
@app.GET(self._config.callback_path) @app.GET(self._config.callback_path)
@@ -108,13 +152,13 @@ class OIDCMixin(KayaMixin):
return return
session = self._session_of(ctx) session = self._session_of(ctx)
expected_state = session.get('oidc_state') expected_state = session.get(self._key('state'))
if state != expected_state: if state != expected_state:
await ctx.send_str(400, 'Invalid state') await ctx.send_str(400, 'Invalid state')
return return
nonce = session.get('oidc_nonce', '') nonce = session.get(self._key('nonce'), '')
code_verifier = session.get('oidc_code_verifier', '') code_verifier = session.get(self._key('code_verifier'), '')
try: try:
token_response = await self._client.fetch_token(code, code_verifier) token_response = await self._client.fetch_token(code, code_verifier)
@@ -131,31 +175,59 @@ class OIDCMixin(KayaMixin):
userinfo = await self._client.fetch_userinfo(access_token) userinfo = await self._client.fetch_userinfo(access_token)
user.update(userinfo) user.update(userinfo)
session.pop('oidc_state', None) session.pop(self._key('state'), None)
session.pop('oidc_nonce', None) session.pop(self._key('code_verifier'), None)
session.pop('oidc_code_verifier', None) # Keep this provider's nonce so refreshed ID tokens can be
# checked for nonce continuity (OIDC Core 12.2).
session['oidc_user'] = user session[self._key('user')] = user
session['oidc_id_token'] = id_token session[self._key('id_token')] = id_token
if isinstance(access_token, str): if isinstance(access_token, str):
session['oidc_access_token'] = access_token session[self._key('access_token')] = access_token
if 'refresh_token' in token_response: if 'refresh_token' in token_response:
session['oidc_refresh_token'] = token_response['refresh_token'] session[self._key('refresh_token')] = token_response['refresh_token']
if self._strategy == 'single-identity':
self._strip_other_identities(session)
session.regenerate_id() session.regenerate_id()
await ctx.send_empty(302, {'Location': self._config.post_login_redirect}) await ctx.send_empty(302, {'Location': self._config.post_login_redirect})
except ValueError as exc: except (ValueError, jwt.PyJWTError) as exc:
await ctx.send_str(400, f'Authentication failed: {exc}') await ctx.send_str(400, f'Authentication failed: {exc}')
@app.GET(self._config.logout_path) @app.GET(self._config.logout_path)
async def logout(ctx: HttpContext) -> None: async def logout(ctx: HttpContext) -> None:
session = self._session_of(ctx) session = self._session_of(ctx)
id_token = session.get('oidc_id_token') id_token = session.get(self._key('id_token'))
session.invalidate() self._logout_session(session)
logout_url = await self._client.build_logout_url(id_token if isinstance(id_token, str) else None) 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 location = logout_url if logout_url is not None else self._config.post_logout_redirect
await ctx.send_empty(302, {'Location': location}) await ctx.send_empty(302, {'Location': location})
def _logout_session(self, session: Session) -> None:
"""Drop this provider's session state.
The shared session is only invalidated when no other provider remains
logged in, so logging out of one provider does not sign the user out of
the others.
"""
for key in [key for key in session if key.startswith(self._prefix)]:
del session[key]
if not any(_is_user_key(key) for key in session):
session.invalidate()
def _strip_other_identities(self, session: Session) -> None:
"""Remove every other provider's identity and tokens from the session.
Only identity keys are dropped; ephemeral flow keys (``state``,
``nonce``, ``code_verifier``) are left alone so logins already in flight
can still complete. Under this strategy the last callback to finish
wins.
"""
for key in [key for key in session
if _is_identity_key(key) and not key.startswith(self._prefix)]:
del session[key]
@staticmethod @staticmethod
def _session_of(ctx: HttpContext) -> Session: def _session_of(ctx: HttpContext) -> Session:
session = ctx.session session = ctx.session
@@ -169,14 +241,31 @@ class OIDCMixin(KayaMixin):
return None return None
def is_authenticated(self, ctx: HttpContext) -> bool: def is_authenticated(self, ctx: HttpContext) -> bool:
return 'oidc_user' in self._session_of(ctx) return self._key('user') in self._session_of(ctx)
def get_user(self, ctx: HttpContext) -> Optional[OIDCUser]: def get_user(self, ctx: HttpContext) -> Optional[OIDCUser]:
user = self._session_of(ctx).get('oidc_user') user = self._session_of(ctx).get(self._key('user'))
if user is None or not isinstance(user, Mapping): if user is None or not isinstance(user, Mapping):
return None return None
return OIDCUser(user) return OIDCUser(user)
def logged_in_providers(self, ctx: HttpContext) -> list[str]:
"""Return the names of every provider with a user in this session."""
return [_provider_of_user_key(key)
for key in self._session_of(ctx) if _is_user_key(key)]
def get_users(self, ctx: HttpContext) -> dict[str, OIDCUser]:
"""Return the users of every provider logged into this session."""
session = self._session_of(ctx)
users: dict[str, OIDCUser] = {}
for key in session:
if not _is_user_key(key):
continue
data = session.get(key)
if isinstance(data, Mapping):
users[_provider_of_user_key(key)] = OIDCUser(data)
return users
def require_auth(self, handler: HttpHandler) -> HttpHandler: def require_auth(self, handler: HttpHandler) -> HttpHandler:
async def wrapper(ctx: HttpContext, *args: Any, **kwargs: Any) -> None: async def wrapper(ctx: HttpContext, *args: Any, **kwargs: Any) -> None:
if not self.is_authenticated(ctx): if not self.is_authenticated(ctx):
@@ -186,18 +275,44 @@ class OIDCMixin(KayaMixin):
return wrapper return wrapper
async def refresh_access_token(self, session: Session) -> Optional[str]: async def refresh_access_token(self, session: Session) -> Optional[str]:
refresh_token = session.get('oidc_refresh_token') refresh_token = session.get(self._key('refresh_token'))
if not isinstance(refresh_token, str): if not isinstance(refresh_token, str):
return None return None
try: try:
token_response = await self._client.refresh_token(refresh_token) token_response = await self._client.refresh_token(refresh_token)
access_token = token_response.get('access_token') access_token = token_response.get('access_token')
id_token = token_response.get('id_token')
if isinstance(id_token, str):
stored_nonce = session.get(self._key('nonce'))
claims = await self._client.validate_refreshed_id_token(
id_token,
stored_nonce if isinstance(stored_nonce, str) else None,
)
user = session.get(self._key('user'))
original_sub = user.get('sub') if isinstance(user, Mapping) else None
if isinstance(original_sub, str) and claims.get('sub') != original_sub:
raise ValueError('ID token subject mismatch')
if isinstance(access_token, str): if isinstance(access_token, str):
session['oidc_access_token'] = access_token session[self._key('access_token')] = access_token
if 'id_token' in token_response: if isinstance(id_token, str):
session['oidc_id_token'] = token_response['id_token'] session[self._key('id_token')] = id_token
if 'refresh_token' in token_response: if 'refresh_token' in token_response:
session['oidc_refresh_token'] = token_response['refresh_token'] session[self._key('refresh_token')] = token_response['refresh_token']
return access_token if isinstance(access_token, str) else None return access_token if isinstance(access_token, str) else None
except Exception: except Exception:
return None return None
def _is_user_key(key: str) -> bool:
return key.startswith(f'{_SESSION_NAMESPACE}:') and key.endswith(f':{_USER_SUFFIX}')
def _is_identity_key(key: str) -> bool:
return (key.startswith(f'{_SESSION_NAMESPACE}:')
and key.rsplit(':', 1)[-1] in _IDENTITY_FIELDS)
def _provider_of_user_key(key: str) -> str:
return key[len(_SESSION_NAMESPACE) + 1:-(len(_USER_SUFFIX) + 1)]
@@ -0,0 +1,134 @@
from dataclasses import replace
from types import MappingProxyType
from typing import Callable, Mapping, Optional, Sequence
from kaya.core import HttpContext, KayaApp, KayaMixin
from kaya.session import Session, SessionMixin
from ._config import OIDCConfig
from ._mixin import (
_IDENTITY_STRATEGIES,
HttpHandler,
IdentityStrategy,
OIDCMixin,
OIDCUser,
)
class OIDCProviders(KayaMixin):
"""Wire several OpenID Connect providers into one application.
Each config's ``provider`` becomes a route segment under ``base_path``:
``{base_path}/{provider}/login``, ``{base_path}/{provider}/callback`` and
``{base_path}/{provider}/logout``. All providers share the supplied
:class:`~kaya.session.SessionMixin` and namespace their session state by
provider.
``strategy`` controls how many identities may coexist in the session. The
default ``'single-identity'`` evicts every other provider's identity after a
successful login (last login wins). ``'multi-identity'`` keeps each provider
logged in independently, so logging out of one provider does not sign the
user out of the others.
Example::
session = SessionMixin(InMemorySessionStore())
oidc = OIDCProviders([google_config, github_config], session=session)
app = KayaApp(mixins=[session, oidc])
@app.GET('/profile')
@oidc.require_auth('google')
async def profile(ctx: HttpContext):
user = oidc.get_user(ctx, 'google')
...
"""
def __init__(
self,
configs: Sequence[OIDCConfig],
session: SessionMixin,
base_path: str = '/auth',
strategy: IdentityStrategy = 'single-identity',
) -> None:
if not configs:
raise ValueError('At least one OIDCConfig is required')
if not base_path.startswith('/'):
raise ValueError('base_path must start with "/"')
if strategy not in _IDENTITY_STRATEGIES:
raise ValueError(
'strategy must be one of '
f'{", ".join(map(repr, _IDENTITY_STRATEGIES))} (got {strategy!r})'
)
self._session = session
self._strategy = strategy
self._base_path = base_path.rstrip('/')
self._providers: dict[str, OIDCMixin] = {}
for config in configs:
name = config.provider
if name in self._providers:
raise ValueError(f'Duplicate OIDC provider {name!r}')
self._providers[name] = OIDCMixin(
replace(
config,
login_path=f'{self._base_path}/{name}/login',
callback_path=f'{self._base_path}/{name}/callback',
logout_path=f'{self._base_path}/{name}/logout',
),
session=session,
strategy=strategy,
)
@property
def base_path(self) -> str:
return self._base_path
@property
def strategy(self) -> IdentityStrategy:
return self._strategy
@property
def providers(self) -> Mapping[str, OIDCMixin]:
return MappingProxyType(self._providers)
def provider(self, name: str) -> OIDCMixin:
try:
return self._providers[name]
except KeyError:
raise KeyError(f'Unknown OIDC provider {name!r}') from None
def __getitem__(self, name: str) -> OIDCMixin:
return self.provider(name)
@property
def dependencies(self) -> Sequence[KayaMixin]:
return [self._session]
def apply(self, app: KayaApp) -> None:
for mixin in self._providers.values():
mixin.apply(app)
def is_authenticated(self, ctx: HttpContext, provider: str) -> bool:
return self.provider(provider).is_authenticated(ctx)
def get_user(self, ctx: HttpContext, provider: str) -> Optional[OIDCUser]:
return self.provider(provider).get_user(ctx)
def get_users(self, ctx: HttpContext) -> dict[str, OIDCUser]:
"""Return the user of every provider with a session on ``ctx``."""
users: dict[str, OIDCUser] = {}
for name, mixin in self._providers.items():
user = mixin.get_user(ctx)
if user is not None:
users[name] = user
return users
def logged_in_providers(self, ctx: HttpContext) -> list[str]:
"""Return the names of the providers logged into ``ctx``."""
return [name for name, mixin in self._providers.items() if mixin.is_authenticated(ctx)]
def require_auth(self, provider: str) -> Callable[[HttpHandler], HttpHandler]:
return self.provider(provider).require_auth
async def refresh_access_token(self, session: Session, provider: str) -> Optional[str]:
return await self.provider(provider).refresh_access_token(session)
File diff suppressed because it is too large. Load diff
+43
View File
@@ -69,6 +69,48 @@ object, so any valid OpenAPI 3.1 construct can be used:
`parameters` entries whose `name` and `in` match an auto-generated path `parameters` entries whose `name` and `in` match an auto-generated path
parameter override it; all others are appended. parameter override it; all others are appended.
## Reusable schemas (`components`)
Schemas can be defined once and referenced from operation fragments: pass a
`components` mapping to the mixin and use `schema_ref` in `@operation`.
```python
from kaya.openapi import OpenAPIMixin, operation, schema_ref
USER = {
'type': 'object',
'required': ['id'],
'properties': {'id': {'type': 'integer'}, 'name': {'type': 'string'}},
}
USER_PAGE = {
'type': 'object',
'properties': {'results': {'type': 'array', 'items': schema_ref('User')}},
}
app = KayaApp(mixins=[OpenAPIMixin(
title='My API',
version='1.0.0',
components={'schemas': {'User': USER, 'UserPage': USER_PAGE}},
)])
@app.GET('/users')
@operation(summary='List users',
responses={200: {
'description': 'A page of users',
'content': {'application/json': {'schema': schema_ref('UserPage')}},
}})
async def list_users(ctx: HttpContext) -> None:
...
```
`components` is merged verbatim (as a copy) into the generated document under
`components`, so Swagger UI lists the named schemas. Every
`#/components/<section>/<name>` reference — in operations or between
components — is validated when the document is generated: a missing section or
entry raises `ValueError`. Reflected references (`./other.json#/...`, absolute
URLs) are left alone.
## Configuration ## Configuration
```python ```python
@@ -81,6 +123,7 @@ OpenAPIMixin(
servers=[{'url': 'https://api.example.com'}], servers=[{'url': 'https://api.example.com'}],
openapi_version='3.1.0', openapi_version='3.1.0',
include_self=False, # include spec/docs endpoints in the document include_self=False, # include spec/docs endpoints in the document
components=None, # reusable OpenAPI components (e.g. schemas)
) )
``` ```
@@ -1,10 +1,11 @@
from ._metadata import operation from ._metadata import operation
from ._mixin import OpenAPIMixin from ._mixin import OpenAPIMixin
from ._spec import generate_spec from ._spec import generate_spec, schema_ref
__all__ = [ __all__ = [
'OpenAPIMixin', 'OpenAPIMixin',
'generate_spec', 'generate_spec',
'operation', 'operation',
'schema_ref',
] ]
@@ -25,7 +25,9 @@ def operation(summary: Optional[str] = None,
:class:`~kaya.openapi.OpenAPIMixin` when generating the specification. :class:`~kaya.openapi.OpenAPIMixin` when generating the specification.
All schema fragments are plain dicts inserted verbatim into the generated All schema fragments are plain dicts inserted verbatim into the generated
OpenAPI document, so any valid OpenAPI 3.1 construct can be used. OpenAPI document, so any valid OpenAPI 3.1 construct can be used. ``$ref``
fragments may point at the reusable schemas registered through the
``components`` argument of :class:`~kaya.openapi.OpenAPIMixin`.
Example:: Example::
@@ -55,6 +55,11 @@ class OpenAPIMixin(KayaMixin):
bodies and response schemas to individual handlers; the handler docstring bodies and response schemas to individual handlers; the handler docstring
is used as the operation description when no explicit one is given. is used as the operation description when no explicit one is given.
Reusable schemas are supplied through ``components`` and referenced from
operation fragments with :func:`~kaya.openapi.schema_ref`; every
``#/components/...`` reference is validated when the document is
generated.
Example:: Example::
openapi = OpenAPIMixin(title='My API', version='1.0.0') openapi = OpenAPIMixin(title='My API', version='1.0.0')
@@ -74,7 +79,8 @@ class OpenAPIMixin(KayaMixin):
docs_path: str = '/docs', docs_path: str = '/docs',
servers: Optional[Sequence[Mapping[str, Any]]] = None, servers: Optional[Sequence[Mapping[str, Any]]] = None,
openapi_version: str = '3.1.0', openapi_version: str = '3.1.0',
include_self: bool = False) -> None: include_self: bool = False,
components: Optional[Mapping[str, Any]] = None) -> None:
self._title = title self._title = title
self._version = version self._version = version
self._description = description self._description = description
@@ -83,6 +89,7 @@ class OpenAPIMixin(KayaMixin):
self._servers = servers self._servers = servers
self._openapi_version = openapi_version self._openapi_version = openapi_version
self._include_self = include_self self._include_self = include_self
self._components = components
def apply(self, app: KayaApp) -> None: def apply(self, app: KayaApp) -> None:
@app.GET(self._spec_path) @app.GET(self._spec_path)
@@ -107,4 +114,5 @@ class OpenAPIMixin(KayaMixin):
description=self._description, description=self._description,
servers=self._servers, servers=self._servers,
exclude_paths=exclude_paths, exclude_paths=exclude_paths,
openapi_version=self._openapi_version) openapi_version=self._openapi_version,
components=self._components)
@@ -1,7 +1,7 @@
from copy import deepcopy from copy import deepcopy
from dataclasses import dataclass from dataclasses import dataclass
from inspect import getdoc from inspect import getdoc
from typing import AbstractSet, Any, Mapping, Optional, Sequence from typing import AbstractSet, Any, Dict, Iterator, Mapping, Optional, Sequence
from kaya.core import HttpMethod, KayaApp from kaya.core import HttpMethod, KayaApp
from kaya.core._path_handler import PathHandler from kaya.core._path_handler import PathHandler
@@ -22,6 +22,17 @@ STANDARD_METHODS: Sequence[HttpMethod] = (
) )
def schema_ref(name: str) -> Dict[str, str]:
"""Return a reference to the component schema ``name``.
The reference resolves to ``#/components/schemas/<name>`` in the
generated document; the schema itself is supplied through the
``components`` argument of :func:`generate_spec` (or
:class:`~kaya.openapi.OpenAPIMixin`).
"""
return {'$ref': f'#/components/schemas/{name}'}
@dataclass @dataclass
class _Route: class _Route:
raw_path: str raw_path: str
@@ -122,13 +133,67 @@ def _build_operation(params: Sequence[Mapping[str, Any]],
return operation return operation
def _unescape_pointer_token(token: str) -> str:
"""Decode one JSON Pointer reference token (``~1`` → ``/``, ``~0`` → ``~``)."""
return token.replace('~1', '/').replace('~0', '~')
def _iter_refs(node: Any) -> Iterator[str]:
"""Yield every ``$ref`` string found anywhere under ``node``."""
if isinstance(node, Mapping):
ref = node.get('$ref')
if isinstance(ref, str):
yield ref
for value in node.values():
yield from _iter_refs(value)
elif isinstance(node, (list, tuple)):
for value in node:
yield from _iter_refs(value)
def _validate_component_refs(spec: Mapping[str, Any]) -> None:
"""Fail fast on references to missing ``components`` entries.
Every ``#/components/<section>/<name>`` reference found in the operations
or in the components themselves must resolve against the document;
references in any other form (relative files, absolute URLs, other
pointers) are ignored. Raises :class:`ValueError` naming the offending
reference and the operation it appears in.
"""
components = spec.get('components')
known: Mapping[str, Any] = components if isinstance(components, Mapping) else {}
def check(node: Any, location: str) -> None:
for ref in _iter_refs(node):
if not ref.startswith('#/components/'):
continue
parts = [_unescape_pointer_token(part) for part in ref[len('#/components/'):].split('/')]
if len(parts) < 2:
raise ValueError(f'invalid component reference {ref!r} in {location}')
section, name = parts[0], parts[1]
bucket = known.get(section)
if not isinstance(bucket, Mapping) or name not in bucket:
raise ValueError(f'unknown component reference {ref!r} in {location}')
paths = spec.get('paths')
if isinstance(paths, Mapping):
for path, item in paths.items():
if not isinstance(item, Mapping):
continue
for method, operation in item.items():
check(operation, f'{method.upper()} {path}')
if isinstance(components, Mapping):
check(components, 'components')
def generate_spec(app: KayaApp, def generate_spec(app: KayaApp,
title: str, title: str,
version: str, version: str,
description: str = '', description: str = '',
servers: Optional[Sequence[Mapping[str, Any]]] = None, servers: Optional[Sequence[Mapping[str, Any]]] = None,
exclude_paths: AbstractSet[str] = frozenset(), exclude_paths: AbstractSet[str] = frozenset(),
openapi_version: str = '3.1.0') -> Mapping[str, Any]: openapi_version: str = '3.1.0',
components: Optional[Mapping[str, Any]] = None) -> Mapping[str, Any]:
"""Generate an OpenAPI specification document from a :class:`KayaApp`. """Generate an OpenAPI specification document from a :class:`KayaApp`.
The application's routing tree is walked and every route is converted to The application's routing tree is walked and every route is converted to
@@ -147,6 +212,9 @@ def generate_spec(app: KayaApp,
:param servers: list of OpenAPI server objects :param servers: list of OpenAPI server objects
:param exclude_paths: Kaya paths (e.g. ``/openapi.json``) to omit :param exclude_paths: Kaya paths (e.g. ``/openapi.json``) to omit
:param openapi_version: OpenAPI version to declare :param openapi_version: OpenAPI version to declare
:param components: OpenAPI ``components`` object merged verbatim into the
document; references to ``#/components/<section>/<name>`` must resolve
against it (e.g. :func:`schema_ref` for reusable schemas)
:return: the OpenAPI document as a JSON-serializable mapping :return: the OpenAPI document as a JSON-serializable mapping
""" """
routes: list[_Route] = [] routes: list[_Route] = []
@@ -174,4 +242,7 @@ def generate_spec(app: KayaApp,
} }
if servers: if servers:
spec['servers'] = [dict(server) for server in servers] spec['servers'] = [dict(server) for server in servers]
if components:
spec['components'] = deepcopy(dict(components))
_validate_component_refs(spec)
return spec return spec
+123 -1
View File
@@ -6,7 +6,7 @@ import httpx
from pwo import async_test from pwo import async_test
from kaya.core import HttpContext, HttpMethod, KayaApp, WebSocket from kaya.core import HttpContext, HttpMethod, KayaApp, WebSocket
from kaya.openapi import OpenAPIMixin, operation from kaya.openapi import OpenAPIMixin, generate_spec, operation, schema_ref
class OpenAPITest(unittest.TestCase): class OpenAPITest(unittest.TestCase):
@@ -163,5 +163,127 @@ class OpenAPITest(unittest.TestCase):
self.assertIn('/spec.json', response.text) self.assertIn('/spec.json', response.text)
class ComponentsTest(unittest.TestCase):
"""Reusable schemas through ``components`` and ``schema_ref``."""
async def _spec_of(self, app: KayaApp) -> Mapping[str, Any]:
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
response = await client.get('/openapi.json')
self.assertEqual(200, response.status_code)
return json.loads(response.text)
def test_schema_ref_helper(self) -> None:
self.assertEqual({'$ref': '#/components/schemas/User'}, schema_ref('User'))
@async_test
async def test_components_merged_and_refs_preserved(self) -> None:
user = {'type': 'object', 'properties': {'id': {'type': 'integer'}}}
page = {'type': 'object',
'properties': {'results': {'type': 'array', 'items': schema_ref('User')}}}
app = KayaApp(mixins=[OpenAPIMixin(
title='Components API',
version='1.0.0',
components={'schemas': {'User': user, 'UserPage': page}},
)])
@app.GET('/users')
@operation(responses={200: {
'description': 'A page of users',
'content': {'application/json': {'schema': schema_ref('UserPage')}},
}})
async def list_users(ctx: HttpContext) -> None:
await ctx.send_str(200, '[]')
spec = await self._spec_of(app)
self.assertEqual({'schemas': {'User': user, 'UserPage': page}}, spec['components'])
schema = spec['paths']['/users']['get']['responses']['200']['content']['application/json']['schema']
self.assertEqual(schema_ref('UserPage'), schema)
@async_test
async def test_request_body_ref_preserved(self) -> None:
app = KayaApp(mixins=[OpenAPIMixin(
title='Components API',
version='1.0.0',
components={'schemas': {'User': {'type': 'object'}}},
)])
@app.POST('/users')
@operation(request_body={'required': True,
'content': {'application/json': {'schema': schema_ref('User')}}},
responses={201: {'description': 'Created'}})
async def create_user(ctx: HttpContext) -> None:
await ctx.send_str(201, '{}')
spec = await self._spec_of(app)
schema = spec['paths']['/users']['post']['requestBody']['content']['application/json']['schema']
self.assertEqual(schema_ref('User'), schema)
@async_test
async def test_default_spec_has_no_components(self) -> None:
app = KayaApp(mixins=[OpenAPIMixin(title='Plain API', version='1.0.0')])
@app.GET('/plain')
async def plain(ctx: HttpContext) -> None:
await ctx.send_str(200, 'ok')
spec = await self._spec_of(app)
self.assertNotIn('components', spec)
def test_unknown_schema_ref_raises(self) -> None:
app = KayaApp()
@app.GET('/broken')
@operation(responses={200: {
'description': 'x',
'content': {'application/json': {'schema': schema_ref('Missing')}},
}})
async def broken(ctx: HttpContext) -> None:
await ctx.send_str(200, 'x')
with self.assertRaises(ValueError) as caught:
generate_spec(app, title='Broken', version='1.0.0',
components={'schemas': {'User': {}}})
message = str(caught.exception)
self.assertIn('#/components/schemas/Missing', message)
self.assertIn('GET /broken', message)
def test_component_to_component_ref_validates(self) -> None:
app = KayaApp()
spec = generate_spec(app, title='Nested', version='1.0.0', components={'schemas': {
'User': {'type': 'object'},
'Page': {'type': 'object',
'properties': {'results': {'type': 'array', 'items': schema_ref('User')}}},
}})
self.assertIn('components', spec)
def test_external_refs_are_ignored(self) -> None:
app = KayaApp()
@app.GET('/external')
@operation(responses={200: {
'description': 'x',
'content': {'application/json': {'schema': {'$ref': './other.json#/Thing'}}},
}})
async def external(ctx: HttpContext) -> None:
await ctx.send_str(200, 'x')
spec = generate_spec(app, title='External', version='1.0.0')
self.assertNotIn('components', spec)
def test_generic_component_section_validated(self) -> None:
app = KayaApp()
@app.GET('/p')
@operation(parameters=[{'name': 'q', 'in': 'query',
'schema': {'$ref': '#/components/parameters/Missing'}}])
async def p(ctx: HttpContext) -> None:
await ctx.send_str(200, 'x')
with self.assertRaises(ValueError) as caught:
generate_spec(app, title='P', version='1.0.0', components={'parameters': {}})
self.assertIn('#/components/parameters/Missing', str(caught.exception))
if __name__ == '__main__': if __name__ == '__main__':
unittest.main() unittest.main()
+7 -3
View File
@@ -69,9 +69,13 @@ async def home(ctx: HttpContext):
`http.response.status_code` when the handler sends the response; 5xx `http.response.status_code` when the handler sends the response; 5xx
statuses mark the span as failed. Exceptions escaping the handler are statuses mark the span as failed. Exceptions escaping the handler are
recorded on the span and also mark it as failed. recorded on the span and also mark it as failed.
- Every WebSocket connection gets one span for its whole lifetime. The close - Every WebSocket connection gets one span for its whole lifetime, carrying
code is recorded as `kaya.websocket.close_code`; exceptions and configured the applicable standard attributes: `url.scheme`, `url.path`, `url.query`,
error close codes mark the span as failed. `network.protocol.name` (`websocket`), `network.transport`,
`client.address`/`client.port` and `server.address`/`server.port`. The
close code is recorded as `kaya.websocket.close_code`; exceptions and
configured error close codes mark the span as failed and set `error.type`
(the exception class name, or the close code).
- W3C `traceparent`/`tracestate` headers on incoming requests are honored, - W3C `traceparent`/`tracestate` headers on incoming requests are honored,
so traces propagate from upstream services. so traces propagate from upstream services.
- Metrics: - Metrics:
+37 -10
View File
@@ -84,6 +84,30 @@ def _server_attributes(ctx: HttpContext) -> Dict[str, Any]:
return attributes return attributes
def _websocket_attributes(ws: WebSocket) -> Dict[str, Any]:
"""Standard attributes applicable to a websocket connection span.
OpenTelemetry defines no websocket-specific conventions; these are the
general URL, network and peer attributes that apply to the connection.
"""
attributes: Dict[str, Any] = {
'url.scheme': ws.scheme,
'url.path': ws.path,
'network.protocol.name': 'websocket',
'network.transport': 'tcp',
}
if ws.query_string:
attributes['url.query'] = ws.query_string
if ws.client is not None:
attributes['client.address'] = ws.client[0]
attributes['client.port'] = ws.client[1]
if ws.server is not None:
attributes['server.address'] = ws.server[0]
if ws.server[1] is not None:
attributes['server.port'] = ws.server[1]
return attributes
def _normalize_header_name(name: str) -> str: def _normalize_header_name(name: str) -> str:
return name.lower().replace('-', '_') return name.lower().replace('-', '_')
@@ -201,7 +225,11 @@ class OTelMixin(KayaMixin):
usual HTTP semantic attributes; the response status code is recorded when usual HTTP semantic attributes; the response status code is recorded when
the handler sends the response and 5xx statuses mark the span as failed. the handler sends the response and 5xx statuses mark the span as failed.
Exceptions escaping the handler are recorded and mark the span as failed. Exceptions escaping the handler are recorded and mark the span as failed.
WebSocket connections get one span for the whole connection lifetime. WebSocket connections get one span for the whole connection lifetime,
carrying the applicable standard websocket URL, network and peer
attributes (``url.scheme``, ``url.path``, ``network.protocol.name``,
``client.address``/``server.address``, ...); exceptions and configured
error close codes set ``error.type`` and mark the span as failed.
W3C ``traceparent``/``tracestate`` headers on incoming requests are W3C ``traceparent``/``tracestate`` headers on incoming requests are
honored, so traces propagate from upstream services. honored, so traces propagate from upstream services.
@@ -455,14 +483,7 @@ class OTelMixin(KayaMixin):
async def _before_websocket(self, ws: WebSocket) -> Optional[WebSocket]: async def _before_websocket(self, ws: WebSocket) -> Optional[WebSocket]:
if self._is_excluded(ws.path): if self._is_excluded(ws.path):
return None return None
holder = self._start_span( holder = self._start_span(f'WS {ws.path}', _websocket_attributes(ws), ws.headers)
f'WS {ws.path}',
{
'url.scheme': ws.scheme,
'url.path': ws.path,
},
ws.headers,
)
self._spans[id(ws)] = holder self._spans[id(ws)] = holder
self._call_hook(self._websocket_connect_hook, holder.span, ws) self._call_hook(self._websocket_connect_hook, holder.span, ws)
return _TracedWebSocket(ws, holder) return _TracedWebSocket(ws, holder)
@@ -475,10 +496,16 @@ class OTelMixin(KayaMixin):
error = isinstance(exception, BaseException) error = isinstance(exception, BaseException)
if isinstance(exception, BaseException): if isinstance(exception, BaseException):
holder.span.record_exception(exception) holder.span.record_exception(exception)
holder.span.set_attribute('error.type', type(exception).__qualname__)
close_code = holder.close_code close_code = holder.close_code
if close_code is not None: if close_code is not None:
holder.span.set_attribute('kaya.websocket.close_code', close_code) holder.span.set_attribute('kaya.websocket.close_code', close_code)
error = error or close_code in self._websocket_error_close_codes if close_code in self._websocket_error_close_codes:
error = True
if not isinstance(exception, BaseException):
# Mirror the HTTP conventions, where server-side error
# statuses are recorded as ``error.type`` values.
holder.span.set_attribute('error.type', str(close_code))
self._call_hook(self._websocket_close_hook, holder.span, ws) self._call_hook(self._websocket_close_hook, holder.span, ws)
self._finish_span(holder, error) self._finish_span(holder, error)
+29 -3
View File
@@ -279,7 +279,7 @@ class MetricsTest(unittest.TestCase):
class WebSocketTracingTest(unittest.TestCase): class WebSocketTracingTest(unittest.TestCase):
@staticmethod @staticmethod
def _make_ws(headers=()): def _make_ws(headers=(), query_string: bytes = b''):
async def send(message): async def send(message):
pass pass
@@ -289,7 +289,7 @@ class WebSocketTracingTest(unittest.TestCase):
scope = { scope = {
'type': 'websocket', 'type': 'websocket',
'path': '/ws/games/abc', 'path': '/ws/games/abc',
'query_string': b'', 'query_string': query_string,
'scheme': 'ws', 'scheme': 'ws',
'client': ('127.0.0.1', 12345), 'client': ('127.0.0.1', 12345),
'server': ('127.0.0.1', 80), 'server': ('127.0.0.1', 80),
@@ -308,7 +308,29 @@ class WebSocketTracingTest(unittest.TestCase):
self.assertEqual(0, len(mixin._spans)) self.assertEqual(0, len(mixin._spans))
spans = span_exporter.get_finished_spans() spans = span_exporter.get_finished_spans()
self.assertEqual(1, len(spans)) self.assertEqual(1, len(spans))
self.assertEqual('WS /ws/games/abc', spans[0].name) span = spans[0]
self.assertEqual('WS /ws/games/abc', span.name)
assert span.attributes is not None
self.assertEqual('ws', span.attributes['url.scheme'])
self.assertEqual('/ws/games/abc', span.attributes['url.path'])
self.assertNotIn('url.query', span.attributes)
self.assertEqual('websocket', span.attributes['network.protocol.name'])
self.assertEqual('tcp', span.attributes['network.transport'])
self.assertEqual('127.0.0.1', span.attributes['client.address'])
self.assertEqual(12345, span.attributes['client.port'])
self.assertEqual('127.0.0.1', span.attributes['server.address'])
self.assertEqual(80, span.attributes['server.port'])
@async_test
async def test_websocket_query_string_recorded(self):
mixin, span_exporter, _ = make_mixin()
ws = self._make_ws(query_string=b'room=1&token=abc')
wrapped = await mixin._before_websocket(ws)
await mixin._after_websocket(wrapped) # type: ignore[arg-type]
spans = span_exporter.get_finished_spans()
self.assertEqual(1, len(spans))
assert spans[0].attributes is not None
self.assertEqual('room=1&token=abc', spans[0].attributes['url.query'])
@async_test @async_test
async def test_websocket_traceparent_propagates(self): async def test_websocket_traceparent_propagates(self):
@@ -344,6 +366,7 @@ class WebSocketTracingTest(unittest.TestCase):
self.assertEqual(StatusCode.ERROR, spans[0].status.status_code) self.assertEqual(StatusCode.ERROR, spans[0].status.status_code)
assert spans[0].attributes is not None assert spans[0].attributes is not None
self.assertEqual(1011, spans[0].attributes['kaya.websocket.close_code']) self.assertEqual(1011, spans[0].attributes['kaya.websocket.close_code'])
self.assertEqual('1011', spans[0].attributes['error.type'])
@async_test @async_test
async def test_websocket_exception_marks_span_error(self): async def test_websocket_exception_marks_span_error(self):
@@ -358,6 +381,8 @@ class WebSocketTracingTest(unittest.TestCase):
self.assertEqual(1, len(spans)) self.assertEqual(1, len(spans))
self.assertEqual(StatusCode.ERROR, spans[0].status.status_code) self.assertEqual(StatusCode.ERROR, spans[0].status.status_code)
self.assertTrue(any(event.name == 'exception' for event in spans[0].events)) self.assertTrue(any(event.name == 'exception' for event in spans[0].events))
assert spans[0].attributes is not None
self.assertEqual('RuntimeError', spans[0].attributes['error.type'])
@async_test @async_test
async def test_websocket_hooks_are_called(self): async def test_websocket_hooks_are_called(self):
@@ -386,6 +411,7 @@ class WebSocketTracingTest(unittest.TestCase):
self.assertEqual(1, len(spans)) self.assertEqual(1, len(spans))
self.assertEqual(StatusCode.UNSET, spans[0].status.status_code) self.assertEqual(StatusCode.UNSET, spans[0].status.status_code)
assert spans[0].attributes is not None assert spans[0].attributes is not None
self.assertNotIn('error.type', spans[0].attributes)
self.assertTrue(spans[0].attributes['test.websocket_connect_hook']) self.assertTrue(spans[0].attributes['test.websocket_connect_hook'])
self.assertTrue(spans[0].attributes['test.websocket_close_hook']) self.assertTrue(spans[0].attributes['test.websocket_close_hook'])
+43 -1
View File
@@ -17,13 +17,21 @@ from typing import (
from granian._granian import ( # type: ignore[attr-defined] from granian._granian import ( # type: ignore[attr-defined]
RSGIHTTPProtocol, RSGIHTTPProtocol,
RSGIHTTPScope, RSGIHTTPScope,
RSGIProtocolClosed,
RSGIWebsocketProtocol, RSGIWebsocketProtocol,
RSGIWebsocketScope, RSGIWebsocketScope,
RSGIWebsocketTransport, RSGIWebsocketTransport,
) )
from pwo import Maybe from pwo import Maybe
from kaya.core import AbstractKayaApp, HttpContext, HttpMethod, WebSocket, WebSocketMessage from kaya.core import (
AbstractKayaApp,
HttpContext,
HttpMethod,
WebSocket,
WebSocketDisconnect,
WebSocketMessage,
)
from kaya.core._types import StrOrStrings from kaya.core._types import StrOrStrings
@@ -109,6 +117,7 @@ class RsgiContext(HttpContext):
class RsgiWebSocket(WebSocket): class RsgiWebSocket(WebSocket):
_protocol: RSGIWebsocketProtocol _protocol: RSGIWebsocketProtocol
_transport: Optional[RSGIWebsocketTransport] _transport: Optional[RSGIWebsocketTransport]
_closed: bool
scheme: str scheme: str
path: str path: str
query_string: str query_string: str
@@ -124,6 +133,7 @@ class RsgiWebSocket(WebSocket):
) )
self._protocol = protocol self._protocol = protocol
self._transport = None self._transport = None
self._closed = False
self.scheme = scope.scheme self.scheme = scope.scheme
self.path = scope.path self.path = scope.path
self.query_string = scope.query_string self.query_string = scope.query_string
@@ -144,13 +154,28 @@ class RsgiWebSocket(WebSocket):
async def accept(self, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None: async def accept(self, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
# RSGI's websocket accept() takes no arguments (https://github.com/emmett-framework/granian/blob/master/docs/spec/RSGI.md), # RSGI's websocket accept() takes no arguments (https://github.com/emmett-framework/granian/blob/master/docs/spec/RSGI.md),
# so handshake headers (e.g. Set-Cookie) cannot be sent on RSGI; they are silently ignored here. # so handshake headers (e.g. Set-Cookie) cannot be sent on RSGI; they are silently ignored here.
try:
self._transport = await self._protocol.accept() self._transport = await self._protocol.accept()
except RSGIProtocolClosed as exc:
self._closed = True
raise WebSocketDisconnect(code=1006) from exc
async def receive(self) -> WebSocketMessage: async def receive(self) -> WebSocketMessage:
if self._transport is None: if self._transport is None:
raise RuntimeError('WebSocket connection has not been accepted yet') raise RuntimeError('WebSocket connection has not been accepted yet')
if self._closed:
return WebSocketMessage(kind='close')
try:
message = await self._transport.receive() message = await self._transport.receive()
except RSGIProtocolClosed:
# Granian raises this when the peer is gone without a close
# handshake (browser tab closed, network dropped, ...). Surface it
# as a close message so receive loops terminate normally instead
# of seeing an exception.
self._closed = True
return WebSocketMessage(kind='close')
if message.kind == 0: if message.kind == 0:
self._closed = True
return WebSocketMessage(kind='close') return WebSocketMessage(kind='close')
elif message.kind == 1: elif message.kind == 1:
return WebSocketMessage(kind='binary', data=message.data) return WebSocketMessage(kind='binary', data=message.data)
@@ -162,15 +187,32 @@ class RsgiWebSocket(WebSocket):
async def send_text(self, data: str) -> None: async def send_text(self, data: str) -> None:
if self._transport is None: if self._transport is None:
raise RuntimeError('WebSocket connection has not been accepted yet') raise RuntimeError('WebSocket connection has not been accepted yet')
if self._closed:
raise WebSocketDisconnect(code=1006)
try:
await self._transport.send_str(data) await self._transport.send_str(data)
except RSGIProtocolClosed as exc:
self._closed = True
raise WebSocketDisconnect(code=1006) from exc
async def send_bytes(self, data: bytes) -> None: async def send_bytes(self, data: bytes) -> None:
if self._transport is None: if self._transport is None:
raise RuntimeError('WebSocket connection has not been accepted yet') raise RuntimeError('WebSocket connection has not been accepted yet')
if self._closed:
raise WebSocketDisconnect(code=1006)
try:
await self._transport.send_bytes(data) await self._transport.send_bytes(data)
except RSGIProtocolClosed as exc:
self._closed = True
raise WebSocketDisconnect(code=1006) from exc
async def close(self, code: int = 1000) -> None: async def close(self, code: int = 1000) -> None:
if self._closed:
return
try:
self._protocol.close(code) self._protocol.close(code)
except RSGIProtocolClosed:
self._closed = True
async def __anext__(self) -> WebSocketMessage: async def __anext__(self) -> WebSocketMessage:
message = await self.receive() message = await self.receive()
+138 -10
View File
@@ -1,22 +1,150 @@
import unittest import unittest
from types import SimpleNamespace
from typing import Any
from granian.rsgi import ProtocolClosed
from pwo import async_test
from kaya.core import WebSocketDisconnect
from kaya.rsgi import RsgiWebSocket from kaya.rsgi import RsgiWebSocket
class FakeScope:
scheme = 'ws'
path = '/ws'
query_string = ''
headers: dict[str, list[str]] = {}
client = '127.0.0.1:12345'
server = '127.0.0.1:80'
class FakeTransport:
"""A transport whose peer vanished without a close handshake."""
def __init__(self) -> None:
self.receive_calls = 0
self.sent: list[Any] = []
async def receive(self) -> Any:
self.receive_calls += 1
raise ProtocolClosed()
async def send_str(self, data: str) -> None:
self.sent.append(data)
raise ProtocolClosed()
async def send_bytes(self, data: bytes) -> None:
self.sent.append(data)
raise ProtocolClosed()
class CleanCloseTransport(FakeTransport):
"""A transport that reports a close message (kind 0)."""
async def receive(self) -> Any:
self.receive_calls += 1
return SimpleNamespace(kind=0, data=None)
class FakeProtocol:
def __init__(self, transport: FakeTransport) -> None:
self._transport = transport
self.closed: list[int] = []
async def accept(self) -> FakeTransport:
return self._transport
def close(self, code: int) -> None:
self.closed.append(code)
class ClosedOnAcceptProtocol(FakeProtocol):
"""A protocol whose peer vanished during the handshake."""
async def accept(self) -> FakeTransport:
raise ProtocolClosed()
class RsgiWebSocketTest(unittest.TestCase): class RsgiWebSocketTest(unittest.TestCase):
def test_misconfigured_granian(self): def test_misconfigured_granian(self):
class FakeScope: class FakeProtocolWithoutWebsockets:
scheme = 'ws'
path = '/ws'
query_string = ''
headers = {}
client = '127.0.0.1:12345'
server = '127.0.0.1:80'
class FakeProtocol:
pass pass
with self.assertRaises(RuntimeError) as ctx: with self.assertRaises(RuntimeError) as ctx:
RsgiWebSocket(FakeScope(), FakeProtocol()) # type: ignore[arg-type] RsgiWebSocket(FakeScope(), FakeProtocolWithoutWebsockets()) # type: ignore[arg-type]
self.assertIn('Granian was not configured for websockets', str(ctx.exception)) self.assertIn('Granian was not configured for websockets', str(ctx.exception))
async def _socket(self, transport: FakeTransport) -> tuple[RsgiWebSocket, FakeProtocol]:
protocol = FakeProtocol(transport)
ws = RsgiWebSocket(FakeScope(), protocol) # type: ignore[arg-type]
await ws.accept()
return ws, protocol
@async_test
async def test_receive_returns_close_when_transport_is_gone(self):
ws, _ = await self._socket(FakeTransport())
message = await ws.receive()
self.assertEqual('close', message.kind)
@async_test
async def test_iteration_ends_when_transport_is_gone(self):
ws, _ = await self._socket(FakeTransport())
messages = [message async for message in ws]
self.assertEqual([], messages)
@async_test
async def test_receive_remembers_the_close(self):
transport = FakeTransport()
ws, _ = await self._socket(transport)
self.assertEqual('close', (await ws.receive()).kind)
self.assertEqual('close', (await ws.receive()).kind)
self.assertEqual(1, transport.receive_calls)
@async_test
async def test_send_text_raises_disconnect_when_transport_is_gone(self):
ws, _ = await self._socket(FakeTransport())
with self.assertRaises(WebSocketDisconnect) as ctx:
await ws.send_text('hello')
self.assertEqual(1006, ctx.exception.code)
@async_test
async def test_send_bytes_raises_disconnect_when_transport_is_gone(self):
ws, _ = await self._socket(FakeTransport())
with self.assertRaises(WebSocketDisconnect) as ctx:
await ws.send_bytes(b'hello')
self.assertEqual(1006, ctx.exception.code)
@async_test
async def test_sends_after_disconnect_do_not_touch_the_transport(self):
transport = FakeTransport()
ws, _ = await self._socket(transport)
with self.assertRaises(WebSocketDisconnect):
await ws.send_text('first')
with self.assertRaises(WebSocketDisconnect):
await ws.send_text('second')
self.assertEqual(['first'], transport.sent)
@async_test
async def test_send_after_a_clean_close_raises_disconnect(self):
transport = CleanCloseTransport()
ws, _ = await self._socket(transport)
self.assertEqual('close', (await ws.receive()).kind)
with self.assertRaises(WebSocketDisconnect):
await ws.send_text('hello')
self.assertEqual([], transport.sent)
@async_test
async def test_accept_raises_disconnect_when_transport_is_gone(self):
ws = RsgiWebSocket(FakeScope(), ClosedOnAcceptProtocol(FakeTransport())) # type: ignore[arg-type]
with self.assertRaises(WebSocketDisconnect) as ctx:
await ws.accept()
self.assertEqual(1006, ctx.exception.code)
@async_test
async def test_close_is_noop_once_the_transport_is_gone(self):
ws, protocol = await self._socket(FakeTransport())
await ws.receive()
await ws.close(1000)
self.assertEqual([], protocol.closed)