Compare commits
7
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ea82ac44e9
|
||
|
|
e73c4b2c18
|
||
|
|
fa8bb0487d
|
||
|
|
15dc4b6dad
|
||
|
|
3af2e7a975
|
||
|
|
540a7b111e
|
||
|
|
7c46badebb
|
No files matched your search
@@ -4,7 +4,7 @@ from ._http_context import HttpContext
|
||||
from ._mixin import KayaMixin
|
||||
from ._tree import Tree, PathIterator
|
||||
from ._path_handler import PathHandler, Matches
|
||||
from ._websocket import WebSocket, WebSocketMessage
|
||||
from ._websocket import WebSocket, WebSocketDisconnect, WebSocketMessage
|
||||
|
||||
|
||||
__all__ = [
|
||||
@@ -18,5 +18,6 @@ __all__ = [
|
||||
'Matches',
|
||||
'PathIterator',
|
||||
'WebSocket',
|
||||
'WebSocketDisconnect',
|
||||
'WebSocketMessage'
|
||||
]
|
||||
@@ -9,7 +9,7 @@ from ._http_method import HttpMethod
|
||||
from ._mixin import KayaMixin
|
||||
from ._path_handler import Context
|
||||
from ._types import StrOrStrings
|
||||
from ._websocket import WebSocket, WebSocketMessage
|
||||
from ._websocket import WebSocket, WebSocketDisconnect, WebSocketMessage
|
||||
from ._asgi import AsgiContext, AsgiWebSocket
|
||||
from ._tree import Tree
|
||||
from ._types.asgi import LifespanScope, HTTPScope as ASGIHTTPScope, WebSocketScope as ASGIWebSocketScope
|
||||
@@ -164,6 +164,12 @@ class KayaApp(AbstractKayaApp):
|
||||
await handler.handle_request(ws, captured)
|
||||
else:
|
||||
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:
|
||||
ws.exception = exc
|
||||
raise
|
||||
|
||||
@@ -17,7 +17,7 @@ from pwo import Maybe
|
||||
from pathlib import Path
|
||||
from ._http_method import HttpMethod
|
||||
from ._http_context import HttpContext
|
||||
from ._websocket import WebSocket, WebSocketMessage
|
||||
from ._websocket import WebSocket, WebSocketDisconnect, WebSocketMessage
|
||||
from ._types import StrOrStrings
|
||||
from ._types.asgi import HTTPScope, WebSocketScope
|
||||
|
||||
@@ -179,7 +179,16 @@ class AsgiWebSocket(WebSocket):
|
||||
# Emit a list rather than a tuple: wsproto's AcceptConnection (used
|
||||
# by httpx_ws and others) requires list concatenation.
|
||||
message['headers'] = list(encode_headers(headers))
|
||||
await self._send(message)
|
||||
await self._send_message(message)
|
||||
|
||||
async def _send_message(self, message: Mapping[str, Any]) -> None:
|
||||
try:
|
||||
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:
|
||||
message = await self._receive()
|
||||
@@ -200,13 +209,13 @@ class AsgiWebSocket(WebSocket):
|
||||
return WebSocketMessage(kind='close')
|
||||
|
||||
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:
|
||||
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:
|
||||
await self._send({'type': 'websocket.close', 'code': code})
|
||||
await self._send_message({'type': 'websocket.close', 'code': code})
|
||||
|
||||
async def __anext__(self) -> WebSocketMessage:
|
||||
message = await self.receive()
|
||||
|
||||
@@ -13,6 +13,41 @@ class WebSocketMessage:
|
||||
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):
|
||||
path: str
|
||||
query_string: str
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
import unittest
|
||||
import httpx
|
||||
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 kaya.core import KayaApp, WebSocket
|
||||
from kaya.core import KayaApp, WebSocket, WebSocketDisconnect, WebSocketMessage
|
||||
from kaya.core._asgi import AsgiWebSocket
|
||||
|
||||
|
||||
@@ -66,7 +66,7 @@ class WebSocketTest(unittest.TestCase):
|
||||
transport = ASGIWebSocketTransport(app=self.app)
|
||||
|
||||
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:
|
||||
pass
|
||||
|
||||
@@ -151,3 +151,113 @@ class WebSocketTest(unittest.TestCase):
|
||||
}
|
||||
ws = AsgiWebSocket(scope, receive, send)
|
||||
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': [],
|
||||
}
|
||||
@@ -40,20 +40,174 @@ async def profile(ctx: HttpContext):
|
||||
`OIDCMixin` depends on `SessionMixin`; passing only `oidc` to `KayaApp(mixins=...)`
|
||||
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
|
||||
|
||||
- Generic OIDC discovery
|
||||
- Authorization Code Flow with PKCE (S256)
|
||||
- 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
|
||||
- Session fixation defense via `regenerate_id()` after login
|
||||
- Optional userinfo endpoint fetch
|
||||
- Refresh token support
|
||||
- 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.)
|
||||
|
||||
## 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
|
||||
|
||||
- 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.
|
||||
- Only algorithms listed in `OIDCConfig.allowed_id_token_algorithms` are accepted.
|
||||
- Always use HTTPS for `redirect_uri` in production.
|
||||
|
||||
@@ -1,11 +1,15 @@
|
||||
from ._client import OIDCClient
|
||||
from ._config import OIDCConfig
|
||||
from ._mixin import OIDCMixin, OIDCUser
|
||||
from ._config import OIDCConfig, PARPolicy
|
||||
from ._mixin import IdentityStrategy, OIDCMixin, OIDCUser
|
||||
from ._providers import OIDCProviders
|
||||
|
||||
|
||||
__all__ = [
|
||||
'IdentityStrategy',
|
||||
'OIDCClient',
|
||||
'OIDCConfig',
|
||||
'OIDCMixin',
|
||||
'OIDCProviders',
|
||||
'OIDCUser',
|
||||
'PARPolicy',
|
||||
]
|
||||
@@ -1,3 +1,4 @@
|
||||
import re
|
||||
from typing import Any, Mapping, Optional, Sequence, cast
|
||||
from urllib.parse import urlencode
|
||||
|
||||
@@ -5,9 +6,11 @@ import httpx
|
||||
import jwt
|
||||
from jwt import PyJWK
|
||||
|
||||
from ._config import OIDCConfig
|
||||
from ._config import PARPolicy, OIDCConfig
|
||||
from ._utils import generate_nonce, generate_pkce, generate_state
|
||||
|
||||
_TENANT_ID_PLACEHOLDER = re.compile(r'\{tenantid\}', re.IGNORECASE)
|
||||
|
||||
|
||||
class OIDCClient:
|
||||
"""Low-level OIDC client implementing discovery, PKCE, and token validation."""
|
||||
@@ -42,8 +45,14 @@ class OIDCClient:
|
||||
async with httpx.AsyncClient() as client:
|
||||
return await client.request(method, url, **kwargs)
|
||||
|
||||
async def build_authorization_url(self) -> tuple[str, str, str, str]:
|
||||
"""Return (authorization_url, state, nonce, code_verifier)."""
|
||||
async def build_authorization_url(self, par: PARPolicy = 'never') -> tuple[str, str, str, str]:
|
||||
"""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()
|
||||
authorization_endpoint = metadata.get('authorization_endpoint')
|
||||
if not isinstance(authorization_endpoint, str):
|
||||
@@ -54,7 +63,6 @@ class OIDCClient:
|
||||
code_verifier, code_challenge = generate_pkce()
|
||||
|
||||
params = {
|
||||
'client_id': self._config.client_id,
|
||||
'response_type': 'code',
|
||||
'scope': ' '.join(self._config.scopes),
|
||||
'redirect_uri': self._config.redirect_uri,
|
||||
@@ -63,9 +71,67 @@ class OIDCClient:
|
||||
'code_challenge': code_challenge,
|
||||
'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
|
||||
|
||||
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]:
|
||||
metadata = await self._ensure_metadata()
|
||||
token_endpoint = metadata.get('token_endpoint')
|
||||
@@ -79,11 +145,8 @@ class OIDCClient:
|
||||
'client_id': self._config.client_id,
|
||||
'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()
|
||||
return cast(Mapping[str, Any], resp.json())
|
||||
|
||||
@@ -98,15 +161,35 @@ class OIDCClient:
|
||||
'refresh_token': refresh_token,
|
||||
'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()
|
||||
return cast(Mapping[str, Any], resp.json())
|
||||
|
||||
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)
|
||||
alg = header.get('alg')
|
||||
if not isinstance(alg, str):
|
||||
@@ -115,29 +198,59 @@ class OIDCClient:
|
||||
if alg == 'none':
|
||||
if not self._config.allow_unsigned_id_tokens:
|
||||
raise ValueError('Unsigned ID tokens are not allowed')
|
||||
metadata_issuer = await self._get_metadata_issuer()
|
||||
payload = jwt.decode(
|
||||
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:
|
||||
if alg not in self._config.allowed_id_token_algorithms:
|
||||
raise ValueError(f'ID token algorithm {alg} is not allowed')
|
||||
metadata_issuer = await self._get_metadata_issuer()
|
||||
jwks = await self._ensure_jwks()
|
||||
signing_key = self._get_signing_key(jwks, header.get('kid'), alg)
|
||||
payload = jwt.decode(
|
||||
id_token,
|
||||
signing_key.key,
|
||||
algorithms=[alg],
|
||||
issuer=self._config.issuer,
|
||||
audience=self._config.client_id,
|
||||
)
|
||||
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError('ID token payload is not a JSON object')
|
||||
if payload.get('nonce') != nonce:
|
||||
raise ValueError('ID token nonce mismatch')
|
||||
self._validate_issuer(payload, metadata_issuer)
|
||||
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:
|
||||
keys = jwks.get('keys')
|
||||
if not isinstance(keys, Sequence):
|
||||
|
||||
@@ -1,13 +1,31 @@
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Optional, Sequence
|
||||
from typing import Any, Literal, Optional, Sequence
|
||||
|
||||
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
|
||||
class OIDCConfig:
|
||||
"""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.
|
||||
``fetch_userinfo`` controls whether the userinfo endpoint is queried after
|
||||
token exchange.
|
||||
@@ -35,8 +53,21 @@ class OIDCConfig:
|
||||
'EdDSA',
|
||||
))
|
||||
http_client: Optional[httpx.AsyncClient] = None
|
||||
provider: str = 'default'
|
||||
par: PARPolicy = 'never'
|
||||
|
||||
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)
|
||||
if 'openid' not in scopes_list:
|
||||
self.scopes = ('openid', *scopes_list)
|
||||
@@ -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.session import Session, SessionMixin
|
||||
from urllib.parse import parse_qs
|
||||
@@ -10,6 +11,13 @@ from ._config import OIDCConfig
|
||||
|
||||
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]):
|
||||
"""Read-only view of the OIDC userinfo / ID token claims."""
|
||||
@@ -28,7 +36,7 @@ class OIDCUser(Mapping[str, Any]):
|
||||
|
||||
@property
|
||||
def sub(self) -> str:
|
||||
return self._data['sub']
|
||||
return cast(str, self._data['sub'])
|
||||
|
||||
@property
|
||||
def email(self) -> Optional[str]:
|
||||
@@ -56,6 +64,17 @@ class OIDCMixin(KayaMixin):
|
||||
- ``callback_path`` (default ``/auth/callback``): handles the provider callback.
|
||||
- ``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::
|
||||
|
||||
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._session = session
|
||||
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
|
||||
def dependencies(self) -> Sequence[KayaMixin]:
|
||||
@@ -81,11 +123,13 @@ class OIDCMixin(KayaMixin):
|
||||
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()
|
||||
auth_url, state, nonce, code_verifier = await self._client.build_authorization_url(
|
||||
par=self._config.par
|
||||
)
|
||||
session = self._session_of(ctx)
|
||||
session['oidc_state'] = state
|
||||
session['oidc_nonce'] = nonce
|
||||
session['oidc_code_verifier'] = code_verifier
|
||||
session[self._key('state')] = state
|
||||
session[self._key('nonce')] = nonce
|
||||
session[self._key('code_verifier')] = code_verifier
|
||||
await ctx.send_empty(302, {'Location': auth_url})
|
||||
|
||||
@app.GET(self._config.callback_path)
|
||||
@@ -108,13 +152,13 @@ class OIDCMixin(KayaMixin):
|
||||
return
|
||||
|
||||
session = self._session_of(ctx)
|
||||
expected_state = session.get('oidc_state')
|
||||
expected_state = session.get(self._key('state'))
|
||||
if state != expected_state:
|
||||
await ctx.send_str(400, 'Invalid state')
|
||||
return
|
||||
|
||||
nonce = session.get('oidc_nonce', '')
|
||||
code_verifier = session.get('oidc_code_verifier', '')
|
||||
nonce = session.get(self._key('nonce'), '')
|
||||
code_verifier = session.get(self._key('code_verifier'), '')
|
||||
|
||||
try:
|
||||
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)
|
||||
user.update(userinfo)
|
||||
|
||||
session.pop('oidc_state', None)
|
||||
session.pop('oidc_nonce', None)
|
||||
session.pop('oidc_code_verifier', None)
|
||||
session.pop(self._key('state'), None)
|
||||
session.pop(self._key('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['oidc_id_token'] = id_token
|
||||
session[self._key('user')] = user
|
||||
session[self._key('id_token')] = id_token
|
||||
if isinstance(access_token, str):
|
||||
session['oidc_access_token'] = access_token
|
||||
session[self._key('access_token')] = access_token
|
||||
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()
|
||||
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}')
|
||||
|
||||
@app.GET(self._config.logout_path)
|
||||
async def logout(ctx: HttpContext) -> None:
|
||||
session = self._session_of(ctx)
|
||||
id_token = session.get('oidc_id_token')
|
||||
session.invalidate()
|
||||
id_token = session.get(self._key('id_token'))
|
||||
self._logout_session(session)
|
||||
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})
|
||||
|
||||
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
|
||||
def _session_of(ctx: HttpContext) -> Session:
|
||||
session = ctx.session
|
||||
@@ -169,14 +241,31 @@ class OIDCMixin(KayaMixin):
|
||||
return None
|
||||
|
||||
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]:
|
||||
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):
|
||||
return None
|
||||
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:
|
||||
async def wrapper(ctx: HttpContext, *args: Any, **kwargs: Any) -> None:
|
||||
if not self.is_authenticated(ctx):
|
||||
@@ -186,18 +275,44 @@ class OIDCMixin(KayaMixin):
|
||||
return wrapper
|
||||
|
||||
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):
|
||||
return None
|
||||
try:
|
||||
token_response = await self._client.refresh_token(refresh_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):
|
||||
session['oidc_access_token'] = access_token
|
||||
if 'id_token' in token_response:
|
||||
session['oidc_id_token'] = token_response['id_token']
|
||||
session[self._key('access_token')] = access_token
|
||||
if isinstance(id_token, str):
|
||||
session[self._key('id_token')] = id_token
|
||||
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
|
||||
except Exception:
|
||||
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
@@ -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
|
||||
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
|
||||
|
||||
```python
|
||||
@@ -81,6 +123,7 @@ OpenAPIMixin(
|
||||
servers=[{'url': 'https://api.example.com'}],
|
||||
openapi_version='3.1.0',
|
||||
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 ._mixin import OpenAPIMixin
|
||||
from ._spec import generate_spec
|
||||
from ._spec import generate_spec, schema_ref
|
||||
|
||||
|
||||
__all__ = [
|
||||
'OpenAPIMixin',
|
||||
'generate_spec',
|
||||
'operation',
|
||||
'schema_ref',
|
||||
]
|
||||
@@ -25,7 +25,9 @@ def operation(summary: Optional[str] = None,
|
||||
:class:`~kaya.openapi.OpenAPIMixin` when generating the specification.
|
||||
|
||||
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::
|
||||
|
||||
|
||||
@@ -55,6 +55,11 @@ class OpenAPIMixin(KayaMixin):
|
||||
bodies and response schemas to individual handlers; the handler docstring
|
||||
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::
|
||||
|
||||
openapi = OpenAPIMixin(title='My API', version='1.0.0')
|
||||
@@ -74,7 +79,8 @@ class OpenAPIMixin(KayaMixin):
|
||||
docs_path: str = '/docs',
|
||||
servers: Optional[Sequence[Mapping[str, Any]]] = None,
|
||||
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._version = version
|
||||
self._description = description
|
||||
@@ -83,6 +89,7 @@ class OpenAPIMixin(KayaMixin):
|
||||
self._servers = servers
|
||||
self._openapi_version = openapi_version
|
||||
self._include_self = include_self
|
||||
self._components = components
|
||||
|
||||
def apply(self, app: KayaApp) -> None:
|
||||
@app.GET(self._spec_path)
|
||||
@@ -107,4 +114,5 @@ class OpenAPIMixin(KayaMixin):
|
||||
description=self._description,
|
||||
servers=self._servers,
|
||||
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 dataclasses import dataclass
|
||||
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._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
|
||||
class _Route:
|
||||
raw_path: str
|
||||
@@ -122,13 +133,67 @@ def _build_operation(params: Sequence[Mapping[str, Any]],
|
||||
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,
|
||||
title: str,
|
||||
version: str,
|
||||
description: str = '',
|
||||
servers: Optional[Sequence[Mapping[str, Any]]] = None,
|
||||
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`.
|
||||
|
||||
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 exclude_paths: Kaya paths (e.g. ``/openapi.json``) to omit
|
||||
: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
|
||||
"""
|
||||
routes: list[_Route] = []
|
||||
@@ -174,4 +242,7 @@ def generate_spec(app: KayaApp,
|
||||
}
|
||||
if servers:
|
||||
spec['servers'] = [dict(server) for server in servers]
|
||||
if components:
|
||||
spec['components'] = deepcopy(dict(components))
|
||||
_validate_component_refs(spec)
|
||||
return spec
|
||||
@@ -6,7 +6,7 @@ import httpx
|
||||
from pwo import async_test
|
||||
|
||||
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):
|
||||
@@ -163,5 +163,127 @@ class OpenAPITest(unittest.TestCase):
|
||||
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__':
|
||||
unittest.main()
|
||||
@@ -69,9 +69,13 @@ async def home(ctx: HttpContext):
|
||||
`http.response.status_code` when the handler sends the response; 5xx
|
||||
statuses mark the span as failed. Exceptions escaping the handler are
|
||||
recorded on the span and also mark it as failed.
|
||||
- Every WebSocket connection gets one span for its whole lifetime. The close
|
||||
code is recorded as `kaya.websocket.close_code`; exceptions and configured
|
||||
error close codes mark the span as failed.
|
||||
- Every WebSocket connection gets one span for its whole lifetime, carrying
|
||||
the applicable standard attributes: `url.scheme`, `url.path`, `url.query`,
|
||||
`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,
|
||||
so traces propagate from upstream services.
|
||||
- Metrics:
|
||||
|
||||
@@ -84,6 +84,30 @@ def _server_attributes(ctx: HttpContext) -> Dict[str, Any]:
|
||||
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:
|
||||
return name.lower().replace('-', '_')
|
||||
|
||||
@@ -201,7 +225,11 @@ class OTelMixin(KayaMixin):
|
||||
usual HTTP semantic attributes; the response status code is recorded when
|
||||
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.
|
||||
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
|
||||
honored, so traces propagate from upstream services.
|
||||
|
||||
@@ -455,14 +483,7 @@ class OTelMixin(KayaMixin):
|
||||
async def _before_websocket(self, ws: WebSocket) -> Optional[WebSocket]:
|
||||
if self._is_excluded(ws.path):
|
||||
return None
|
||||
holder = self._start_span(
|
||||
f'WS {ws.path}',
|
||||
{
|
||||
'url.scheme': ws.scheme,
|
||||
'url.path': ws.path,
|
||||
},
|
||||
ws.headers,
|
||||
)
|
||||
holder = self._start_span(f'WS {ws.path}', _websocket_attributes(ws), ws.headers)
|
||||
self._spans[id(ws)] = holder
|
||||
self._call_hook(self._websocket_connect_hook, holder.span, ws)
|
||||
return _TracedWebSocket(ws, holder)
|
||||
@@ -475,10 +496,16 @@ class OTelMixin(KayaMixin):
|
||||
error = isinstance(exception, BaseException)
|
||||
if isinstance(exception, BaseException):
|
||||
holder.span.record_exception(exception)
|
||||
holder.span.set_attribute('error.type', type(exception).__qualname__)
|
||||
close_code = holder.close_code
|
||||
if close_code is not None:
|
||||
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._finish_span(holder, error)
|
||||
|
||||
|
||||
@@ -279,7 +279,7 @@ class MetricsTest(unittest.TestCase):
|
||||
class WebSocketTracingTest(unittest.TestCase):
|
||||
|
||||
@staticmethod
|
||||
def _make_ws(headers=()):
|
||||
def _make_ws(headers=(), query_string: bytes = b''):
|
||||
async def send(message):
|
||||
pass
|
||||
|
||||
@@ -289,7 +289,7 @@ class WebSocketTracingTest(unittest.TestCase):
|
||||
scope = {
|
||||
'type': 'websocket',
|
||||
'path': '/ws/games/abc',
|
||||
'query_string': b'',
|
||||
'query_string': query_string,
|
||||
'scheme': 'ws',
|
||||
'client': ('127.0.0.1', 12345),
|
||||
'server': ('127.0.0.1', 80),
|
||||
@@ -308,7 +308,29 @@ class WebSocketTracingTest(unittest.TestCase):
|
||||
self.assertEqual(0, len(mixin._spans))
|
||||
spans = span_exporter.get_finished_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 def test_websocket_traceparent_propagates(self):
|
||||
@@ -344,6 +366,7 @@ class WebSocketTracingTest(unittest.TestCase):
|
||||
self.assertEqual(StatusCode.ERROR, spans[0].status.status_code)
|
||||
assert spans[0].attributes is not None
|
||||
self.assertEqual(1011, spans[0].attributes['kaya.websocket.close_code'])
|
||||
self.assertEqual('1011', spans[0].attributes['error.type'])
|
||||
|
||||
@async_test
|
||||
async def test_websocket_exception_marks_span_error(self):
|
||||
@@ -358,6 +381,8 @@ class WebSocketTracingTest(unittest.TestCase):
|
||||
self.assertEqual(1, len(spans))
|
||||
self.assertEqual(StatusCode.ERROR, spans[0].status.status_code)
|
||||
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 def test_websocket_hooks_are_called(self):
|
||||
@@ -386,6 +411,7 @@ class WebSocketTracingTest(unittest.TestCase):
|
||||
self.assertEqual(1, len(spans))
|
||||
self.assertEqual(StatusCode.UNSET, spans[0].status.status_code)
|
||||
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_close_hook'])
|
||||
|
||||
|
||||
@@ -17,13 +17,21 @@ from typing import (
|
||||
from granian._granian import ( # type: ignore[attr-defined]
|
||||
RSGIHTTPProtocol,
|
||||
RSGIHTTPScope,
|
||||
RSGIProtocolClosed,
|
||||
RSGIWebsocketProtocol,
|
||||
RSGIWebsocketScope,
|
||||
RSGIWebsocketTransport,
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
@@ -109,6 +117,7 @@ class RsgiContext(HttpContext):
|
||||
class RsgiWebSocket(WebSocket):
|
||||
_protocol: RSGIWebsocketProtocol
|
||||
_transport: Optional[RSGIWebsocketTransport]
|
||||
_closed: bool
|
||||
scheme: str
|
||||
path: str
|
||||
query_string: str
|
||||
@@ -124,6 +133,7 @@ class RsgiWebSocket(WebSocket):
|
||||
)
|
||||
self._protocol = protocol
|
||||
self._transport = None
|
||||
self._closed = False
|
||||
self.scheme = scope.scheme
|
||||
self.path = scope.path
|
||||
self.query_string = scope.query_string
|
||||
@@ -144,13 +154,28 @@ class RsgiWebSocket(WebSocket):
|
||||
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),
|
||||
# so handshake headers (e.g. Set-Cookie) cannot be sent on RSGI; they are silently ignored here.
|
||||
self._transport = await self._protocol.accept()
|
||||
try:
|
||||
self._transport = await self._protocol.accept()
|
||||
except RSGIProtocolClosed as exc:
|
||||
self._closed = True
|
||||
raise WebSocketDisconnect(code=1006) from exc
|
||||
|
||||
async def receive(self) -> WebSocketMessage:
|
||||
if self._transport is None:
|
||||
raise RuntimeError('WebSocket connection has not been accepted yet')
|
||||
message = await self._transport.receive()
|
||||
if self._closed:
|
||||
return WebSocketMessage(kind='close')
|
||||
try:
|
||||
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:
|
||||
self._closed = True
|
||||
return WebSocketMessage(kind='close')
|
||||
elif message.kind == 1:
|
||||
return WebSocketMessage(kind='binary', data=message.data)
|
||||
@@ -162,15 +187,32 @@ class RsgiWebSocket(WebSocket):
|
||||
async def send_text(self, data: str) -> None:
|
||||
if self._transport is None:
|
||||
raise RuntimeError('WebSocket connection has not been accepted yet')
|
||||
await self._transport.send_str(data)
|
||||
if self._closed:
|
||||
raise WebSocketDisconnect(code=1006)
|
||||
try:
|
||||
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:
|
||||
if self._transport is None:
|
||||
raise RuntimeError('WebSocket connection has not been accepted yet')
|
||||
await self._transport.send_bytes(data)
|
||||
if self._closed:
|
||||
raise WebSocketDisconnect(code=1006)
|
||||
try:
|
||||
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:
|
||||
self._protocol.close(code)
|
||||
if self._closed:
|
||||
return
|
||||
try:
|
||||
self._protocol.close(code)
|
||||
except RSGIProtocolClosed:
|
||||
self._closed = True
|
||||
|
||||
async def __anext__(self) -> WebSocketMessage:
|
||||
message = await self.receive()
|
||||
|
||||
@@ -1,22 +1,150 @@
|
||||
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
|
||||
|
||||
|
||||
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):
|
||||
|
||||
def test_misconfigured_granian(self):
|
||||
class FakeScope:
|
||||
scheme = 'ws'
|
||||
path = '/ws'
|
||||
query_string = ''
|
||||
headers = {}
|
||||
client = '127.0.0.1:12345'
|
||||
server = '127.0.0.1:80'
|
||||
|
||||
class FakeProtocol:
|
||||
class FakeProtocolWithoutWebsockets:
|
||||
pass
|
||||
|
||||
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))
|
||||
|
||||
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)
|
||||
Reference in new issue
Block a user