Author SHA1 Message Date
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
woggioni 148a35b71c Add kaya-otel package for OpenTelemetry tracing and metrics
CI / Build Pip package (push) Successful in 4m8s
Instrument HTTP requests and WebSocket connections via Kaya hooks, covering both ASGI and RSGI. Records handler exceptions, WebSocket close codes, optional header capture, exclusions and lifecycle hooks.

Add route-template resolution and exception visibility to kaya-core so trace/metric attributes can use low-cardinality routes and failed spans can record escaped exceptions.
2026-09-19 00:39:15 +00:00
woggioni 69762d93df Document that mixin setup/shutdown loops may not be running
Under RSGI __rsgi_init__ is called before the loop starts, so
asyncio.get_running_loop() raises RuntimeError inside setup(). Spell out
that mixins must use the loop they are given, since the ASGI lifespan
incidentally runs with a live loop.
2026-09-18 02:37:08 +00:00
woggioni aa35e2d30c Refactor forwarded header handling into opt-in kaya-forwarded package with trusted CIDRs
CI / Build Pip package (push) Successful in 3m59s
2026-09-05 15:11:05 +08:00
35 changed files with 3351 additions and 278 deletions

No files matched your search

+24
View File
@@ -56,6 +56,14 @@ jobs:
run: |
.venv/bin/python -m mypy -p kaya.cors
.venv/bin/python -m unittest discover -s packages/kaya-cors/tests
- name: Check kaya-forwarded
run: |
.venv/bin/python -m mypy -p kaya.forwarded
.venv/bin/python -m unittest discover -s packages/kaya-forwarded/tests
- name: Check kaya-otel
run: |
.venv/bin/python -m mypy -p kaya.otel
.venv/bin/python -m unittest discover -s packages/kaya-otel/tests
- name: Publish kaya-core artifacts
env:
TWINE_REPOSITORY_URL: ${{ vars.PYPI_REGISTRY_URL }}
@@ -120,3 +128,19 @@ jobs:
run: |
.venv/bin/pyproject-build packages/kaya-cors
.venv/bin/twine upload --repository gitea packages/kaya-cors/dist/*.whl packages/kaya-cors/dist/*.tar.gz
- name: Publish kaya-forwarded artifacts
env:
TWINE_REPOSITORY_URL: ${{ vars.PYPI_REGISTRY_URL }}
TWINE_USERNAME: ${{ vars.PUBLISHER_USERNAME }}
TWINE_PASSWORD: ${{ secrets.PUBLISHER_TOKEN }}
run: |
.venv/bin/pyproject-build packages/kaya-forwarded
.venv/bin/twine upload --repository gitea packages/kaya-forwarded/dist/*.whl packages/kaya-forwarded/dist/*.tar.gz
- name: Publish kaya-otel artifacts
env:
TWINE_REPOSITORY_URL: ${{ vars.PYPI_REGISTRY_URL }}
TWINE_USERNAME: ${{ vars.PUBLISHER_USERNAME }}
TWINE_PASSWORD: ${{ secrets.PUBLISHER_TOKEN }}
run: |
.venv/bin/pyproject-build packages/kaya-otel
.venv/bin/twine upload --repository gitea packages/kaya-otel/dist/*.whl packages/kaya-otel/dist/*.tar.gz
+9 -1
View File
@@ -14,6 +14,8 @@ This repository is a monorepo for the Kaya framework. The code is split into ind
- **kaya-oidc** — OpenID Connect authentication (`packages/kaya-oidc/`)
- **kaya-openapi** — automatic OpenAPI specification generation (`packages/kaya-openapi/`)
- **kaya-cors** — CORS (Cross-Origin Resource Sharing) support (`packages/kaya-cors/`)
- **kaya-forwarded** — trusted-proxy `Forwarded`/`X-Forwarded-*` client address resolution (`packages/kaya-forwarded/`)
- **kaya-otel** — OpenTelemetry tracing and metrics (`packages/kaya-otel/`)
Additional `kaya-*` packages can be added as new directories under `packages/`.
@@ -28,7 +30,7 @@ pip install --index-url https://gitea.woggioni.net/api/packages/woggioni/pypi/si
Install the packages in development mode:
```bash
pip install -e packages/kaya-core -e packages/kaya-rsgi -e packages/kaya-session -e packages/kaya-session-redis -e packages/kaya-session-memcache -e packages/kaya-oidc -e packages/kaya-openapi -e packages/kaya-cors
pip install -e packages/kaya-core -e packages/kaya-rsgi -e packages/kaya-session -e packages/kaya-session-redis -e packages/kaya-session-memcache -e packages/kaya-oidc -e packages/kaya-openapi -e packages/kaya-cors -e packages/kaya-forwarded -e packages/kaya-otel
```
Run the example:
@@ -48,6 +50,8 @@ python -m unittest discover -s packages/kaya-session-memcache/tests
python -m unittest discover -s packages/kaya-oidc/tests
python -m unittest discover -s packages/kaya-openapi/tests
python -m unittest discover -s packages/kaya-cors/tests
python -m unittest discover -s packages/kaya-forwarded/tests
python -m unittest discover -s packages/kaya-otel/tests
```
## Static analysis
@@ -61,6 +65,8 @@ mypy -p kaya.session.memcache
mypy -p kaya.oidc
mypy -p kaya.openapi
mypy -p kaya.cors
mypy -p kaya.forwarded
mypy -p kaya.otel
```
## Building packages
@@ -74,4 +80,6 @@ python -m build packages/kaya-session-memcache
python -m build packages/kaya-oidc
python -m build packages/kaya-openapi
python -m build packages/kaya-cors
python -m build packages/kaya-forwarded
python -m build packages/kaya-otel
```
+3 -3
View File
@@ -1,10 +1,10 @@
from ._app import AbstractKayaApp, KayaApp
from ._http_method import HttpMethod
from ._http_context import HttpContext, resolve_client
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__ = [
@@ -13,11 +13,11 @@ __all__ = [
'KayaApp',
'KayaMixin',
'HttpContext',
'resolve_client',
'Tree',
'PathHandler',
'Matches',
'PathIterator',
'WebSocket',
'WebSocketDisconnect',
'WebSocketMessage'
]
+17 -1
View File
@@ -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
@@ -145,6 +145,9 @@ class KayaApp(AbstractKayaApp):
await handler.handle_request(ctx, captured)
else:
await ctx.send_empty(404)
except Exception as exc:
ctx.exception = exc
raise
finally:
for hook in reversed(self._after_request_hooks):
await hook(ctx)
@@ -161,6 +164,15 @@ 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
finally:
for hook in reversed(self._after_websocket_hooks):
await hook(ws)
@@ -173,6 +185,10 @@ class KayaApp(AbstractKayaApp):
for mixin in self._mixins:
mixin.shutdown(loop)
def route_template(self, url: str, method: HttpMethod = HttpMethod.GET) -> Optional[str]:
"""Return the registered route template that would handle ``url``."""
return self._tree.route_template(url, method)
def route(self,
paths: StrOrStrings,
methods: Optional[HttpMethod | Sequence[HttpMethod]] = None,
+18 -9
View File
@@ -16,8 +16,8 @@ from typing import (
from pwo import Maybe
from pathlib import Path
from ._http_method import HttpMethod
from ._http_context import HttpContext, resolve_client
from ._websocket import WebSocket, WebSocketMessage
from ._http_context import HttpContext
from ._websocket import WebSocket, WebSocketDisconnect, WebSocketMessage
from ._types import StrOrStrings
from ._types.asgi import HTTPScope, WebSocketScope
@@ -84,9 +84,9 @@ class AsgiContext(HttpContext):
self.query_string = scope['query_string'].decode()
self.method = HttpMethod(scope['method'])
self.scheme = scope.get('scheme', 'http')
self.headers = decode_headers(scope['headers'])
self.client = resolve_client(self.headers, scope['client'])
self.client = scope['client']
self.server = scope['server']
self.headers = decode_headers(scope['headers'])
self.request_body = request_body_iterator
self.session = (scope.get('state') or {}).get('kaya_session')
@@ -169,9 +169,9 @@ class AsgiWebSocket(WebSocket):
self.path = scope['path']
self.query_string = scope['query_string'].decode()
self.scheme = scope.get('scheme', 'ws')
self.headers = decode_headers(scope['headers'])
self.client = resolve_client(self.headers, scope['client'])
self.client = scope['client']
self.server = scope['server']
self.headers = decode_headers(scope['headers'])
async def accept(self, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
message: Dict[str, Any] = {'type': 'websocket.accept'}
@@ -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(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()
@@ -16,83 +16,6 @@ from ._http_method import HttpMethod
from ._types.base import StrOrStrings
def _split_host_port(value: str) -> Tuple[str, Optional[int]]:
if value.startswith('['):
# bracketed IPv6 address, optionally followed by :port
closing = value.find(']')
if closing == -1:
return value, None
host = value[1:closing]
rest = value[closing + 1:]
if rest.startswith(':'):
try:
return host, int(rest[1:])
except ValueError:
return host, None
return host, None
if value.count(':') == 1:
host, _, port_str = value.rpartition(':')
try:
return host, int(port_str)
except ValueError:
return value, None
return value, None
def _first_header_value(headers: Mapping[str, Sequence[str]], name: str) -> Optional[str]:
values = headers.get(name)
if not values:
return None
first = values[0].split(',')[0].strip()
return first or None
def resolve_client(headers: Mapping[str, Sequence[str]],
client: Optional[Tuple[str, int]]) -> Optional[Tuple[str, int]]:
"""
Resolve the client (host, port) pair, honoring forwarded headers.
Resolution order:
1. the ``for=`` parameter of the first entry of the RFC 7239 ``Forwarded`` header
(a ``:port`` suffix, if present, also populates the port)
2. the first entry of ``X-Forwarded-For``
3. the first entry of ``X-Forwarded-Host``
4. the socket peer address (``client``), returned unchanged
In the ``X-Forwarded-*`` cases the port is taken from ``X-Forwarded-Port``
when present and valid, otherwise the socket port is kept.
"""
socket_port = client[1] if client is not None else 0
forwarded_values = headers.get('forwarded')
if forwarded_values:
for raw_value in forwarded_values:
first_entry = raw_value.split(',')[0]
for param in first_entry.split(';'):
key, sep, value = param.partition('=')
if sep and key.strip().lower() == 'for':
for_value = value.strip().strip('"')
if for_value and for_value.lower() != 'unknown':
forwarded_host, forwarded_port = _split_host_port(for_value)
if forwarded_host:
return (forwarded_host,
forwarded_port if forwarded_port is not None else socket_port)
for header_name in ('x-forwarded-for', 'x-forwarded-host'):
host = _first_header_value(headers, header_name)
if host is not None:
port = socket_port
x_forwarded_port = _first_header_value(headers, 'x-forwarded-port')
if x_forwarded_port is not None:
try:
port = int(x_forwarded_port)
except ValueError:
pass
return host, port
return client
class HttpContext(ABC):
pathsend: bool
receive: Callable[[], Awaitable[Any]]
@@ -106,6 +29,7 @@ class HttpContext(ABC):
server: Optional[Tuple[str, Optional[int]]]
request_body: AsyncIterator[bytes]
session: Optional[Any] = None
exception: Optional[BaseException] = None
@abstractmethod
async def stream_body(self,
+15 -2
View File
@@ -29,9 +29,22 @@ class KayaMixin(ABC):
pass
def setup(self, loop: AbstractEventLoop) -> None:
"""Called on lifespan startup (default: no-op)."""
"""Called on lifespan startup (default: no-op).
Under RSGI (granian) the loop is NOT yet running when this is
called: schedule work on the passed ``loop`` (e.g.
``loop.create_task``) instead of calling
``asyncio.get_running_loop()``, which raises ``RuntimeError``
there. Under ASGI the lifespan runs inside a coroutine, so a
running loop happens to be available — do not rely on that.
"""
pass
def shutdown(self, loop: AbstractEventLoop) -> None:
"""Called on lifespan shutdown (default: no-op)."""
"""Called on lifespan shutdown (default: no-op).
As with :meth:`setup`, use the passed ``loop``; it may already be
stopped under RSGI, so ``asyncio.get_running_loop()`` is not
reliable here either.
"""
pass
+31
View File
@@ -237,6 +237,37 @@ class Tree:
# return (handler, unmatched)
return None
def route_template(self, url: str, method: HttpMethod = HttpMethod.GET) -> Optional[str]:
"""Return the registered route template that would handle ``url``.
Static segments are returned verbatim, ``${name}``/``${name:int}``
parameter matchers are reconstructed from the matched nodes, and glob
matchers keep their pattern. Returns ``None`` when no route (including
recursive fallback routes) would handle the request.
"""
result = self.find_node((p for p in PathIterator(urlparse(url).path)), method)
if result is None:
return None
node, captured = result
if len(captured.unmatched_paths) > 0 and not any(handler.recursive for handler in node.handlers):
return None
parts: List[str] = []
current: Optional[Node | PathMatcher] = node
while current is not None:
if isinstance(current, Node):
key = current.key
if not isinstance(key, HttpMethod) and key != '/':
parts.append(str(key))
elif isinstance(current, IntMatcher):
parts.append('${%s:int}' % current.name)
elif isinstance(current, StrMatcher):
parts.append('${%s}' % current.name)
elif isinstance(current, GlobMatcher):
parts.append(current.pattern)
current = current.parent
return '/' + '/'.join(reversed(parts)) if parts else '/'
def parse(self, leaf: str, parent: Optional[Node | PathMatcher]) -> Node | PathMatcher:
start = 0
result = index_of_with_escape(leaf, '${', '\\', 0)
@@ -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
@@ -21,6 +56,7 @@ class WebSocket(ABC):
client: Optional[Tuple[str, int]]
server: Optional[Tuple[str, Optional[int]]]
session: Optional[Any] = None
exception: Optional[BaseException] = None
@abstractmethod
async def accept(self, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
+29 -59
View File
@@ -62,11 +62,6 @@ class AsgiTest(unittest.TestCase):
async def handle_request(ctx: HttpContext, _: List[str]) -> None:
await ctx.stream_body(200, (chunk async for chunk in ctx.request_body))
@self.app.GET('/client-ip')
async def handle_request(ctx: HttpContext) -> None:
host, port = ctx.client if ctx.client is not None else (None, None)
await ctx.send_str(200, json.dumps({'host': host, 'port': port}))
@async_test
async def test_hello(self):
transport = httpx.ASGITransport(app=self.app)
@@ -195,60 +190,6 @@ class AsgiTest(unittest.TestCase):
'employee_id': 101325
}, response)
@async_test
async def test_client_ip_forwarded(self):
transport = httpx.ASGITransport(app=self.app)
async with httpx.AsyncClient(transport=transport, base_url="http://127.0.0.1:80") as client:
# socket peer, no forwarded headers
r = await client.get("/client-ip")
socket_client = json.loads(r.text)
self.assertEqual('127.0.0.1', socket_client['host'])
# RFC 7239 Forwarded header, with port
r = await client.get("/client-ip", headers={'Forwarded': 'for=203.0.113.5:1234'})
self.assertEqual({'host': '203.0.113.5', 'port': 1234}, json.loads(r.text))
# RFC 7239 Forwarded header, bracketed IPv6 with port
r = await client.get("/client-ip", headers={'Forwarded': 'for="[2001:db8::1]:4711"'})
self.assertEqual({'host': '2001:db8::1', 'port': 4711}, json.loads(r.text))
# RFC 7239 Forwarded header without port keeps the socket port
r = await client.get("/client-ip", headers={'Forwarded': 'for=203.0.113.5'})
self.assertEqual({'host': '203.0.113.5', 'port': socket_client['port']}, json.loads(r.text))
# Forwarded with for=unknown falls through to X-Forwarded-For
r = await client.get("/client-ip", headers={
'Forwarded': 'for=unknown',
'X-Forwarded-For': '198.51.100.7',
})
self.assertEqual('198.51.100.7', json.loads(r.text)['host'])
# Forwarded takes precedence over X-Forwarded-For
r = await client.get("/client-ip", headers={
'Forwarded': 'for=203.0.113.5',
'X-Forwarded-For': '198.51.100.7',
})
self.assertEqual('203.0.113.5', json.loads(r.text)['host'])
# X-Forwarded-For: first entry of the chain, port from X-Forwarded-Port
r = await client.get("/client-ip", headers={
'X-Forwarded-For': '203.0.113.5, 70.41.3.18',
'X-Forwarded-Port': '8443',
})
self.assertEqual({'host': '203.0.113.5', 'port': 8443}, json.loads(r.text))
# X-Forwarded-Host fallback
r = await client.get("/client-ip", headers={'X-Forwarded-Host': '198.51.100.7'})
self.assertEqual('198.51.100.7', json.loads(r.text)['host'])
# invalid X-Forwarded-Port is ignored, socket port is kept
r = await client.get("/client-ip", headers={
'X-Forwarded-For': '203.0.113.5',
'X-Forwarded-Port': 'not-a-port',
})
self.assertEqual({'host': '203.0.113.5', 'port': socket_client['port']}, json.loads(r.text))
@async_test
async def test_nested_param_routes(self):
app = KayaApp()
@@ -309,3 +250,32 @@ class AsgiTest(unittest.TestCase):
r = await client.put("/restaurants/42/menu")
self.assertEqual(404, r.status_code)
def test_route_template(self):
self.assertEqual('/employee/${employee_id}',
self.app.route_template('/employee/101325', HttpMethod.GET))
self.assertEqual('/square/${x:int}', self.app.route_template('/square/30', HttpMethod.GET))
self.assertIsNone(self.app.route_template('/unknown', HttpMethod.GET))
@async_test
async def test_exception_exposed_to_after_hooks(self):
app = KayaApp()
seen = []
async def after_request(ctx: HttpContext) -> None:
seen.append(ctx.exception)
app.add_after_request_hook(after_request)
@app.GET('/raises')
async def raises(ctx: HttpContext) -> None:
raise RuntimeError('boom')
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://127.0.0.1:80") as client:
with self.assertRaises(RuntimeError):
await client.get('/raises')
self.assertEqual(1, len(seen))
self.assertIsInstance(seen[0], RuntimeError)
self.assertEqual('boom', str(seen[0]))
+21
View File
@@ -80,6 +80,27 @@ class TreeTest(unittest.TestCase):
self.assertIs(Maybe.of(handler_num).map(self.handlers.__getitem__).or_none(),
Maybe.of_nullable(res).map(lambda it: it[0]).or_none())
def test_route_template(self):
cases: Tuple[Tuple[str, HttpMethod, Optional[str]], ...] = (
('/home/something', HttpMethod.GET, '/home/something'),
('/home/something_else', HttpMethod.POST, '/home/something_else'),
('/home/README.md', HttpMethod.GET, '/home/*.md'),
('/home/something/ciao/blah/README.md', HttpMethod.GET, '/home/something/*/blah/*.md'),
('/home/bar/ciao/blah/README.md', HttpMethod.GET, '/home/bar/*'),
('/unknown', HttpMethod.GET, None),
)
for url, method, expected in cases:
with self.subTest(f'{method} {url}'):
self.assertEqual(expected, self.tree.route_template(url, method))
def test_route_template_with_params(self):
tree = Tree()
tree.add((p for p in ('foo', '${id:int}')), HttpMethod.PUT, self.handlers[0])
tree.add((p for p in ('foo', '${name}')), HttpMethod.GET, self.handlers[1])
self.assertEqual('/foo/${id:int}', tree.route_template('/foo/42', HttpMethod.PUT))
self.assertEqual('/foo/${name}', tree.route_template('/foo/bar', HttpMethod.GET))
self.assertIsNone(tree.route_template('/foo/bar', HttpMethod.DELETE))
def test_two_method_agnostic_matchers_raise(self):
tree = Tree()
tree.add((p for p in ('foo', '*')), None, self.handlers[0])
+113 -34
View File
@@ -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
@@ -132,37 +132,6 @@ class WebSocketTest(unittest.TestCase):
self.assertEqual(1, len(sent_messages))
self.assertEqual({'type': 'websocket.accept'}, sent_messages[0])
@async_test
async def test_client_forwarded_header(self):
async def send(message):
pass
async def receive():
return {'type': 'websocket.connect'}
scope = {
'type': 'websocket',
'path': '/echo',
'query_string': b'',
'scheme': 'ws',
'client': ('127.0.0.1', 12345),
'server': ('127.0.0.1', 80),
'headers': [(b'forwarded', b'for=203.0.113.5:1234')],
}
ws = AsgiWebSocket(scope, receive, send)
self.assertEqual(('203.0.113.5', 1234), ws.client)
scope['headers'] = [
(b'x-forwarded-for', b'203.0.113.5, 70.41.3.18'),
(b'x-forwarded-port', b'8443'),
]
ws = AsgiWebSocket(scope, receive, send)
self.assertEqual(('203.0.113.5', 8443), ws.client)
scope['headers'] = []
ws = AsgiWebSocket(scope, receive, send)
self.assertEqual(('127.0.0.1', 12345), ws.client)
@async_test
async def test_websocket_scope_without_scheme(self):
# Daphne omits the optional `scheme` key from websocket scopes.
@@ -182,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': [],
}
+61
View File
@@ -0,0 +1,61 @@
# kaya-forwarded
Trusted-proxy forwarded header handling for the Kaya web framework.
Without this package, Kaya exposes the raw socket peer address as
`ctx.client` / `ws.client` and ignores `Forwarded` / `X-Forwarded-*` headers
entirely (they are client-controllable and trivially spoofable when the app is
directly exposed).
`ForwardedHeadersMixin` opts the application into honoring those headers, but
only when the direct socket peer is a trusted proxy, identified by a list of
trusted CIDRs/IPs.
## Usage
```python
from kaya.core import KayaApp, HttpContext
from kaya.forwarded import ForwardedHeadersMixin
app = KayaApp(mixins=[
ForwardedHeadersMixin(trusted_proxies=['127.0.0.1', '::1', '10.0.0.0/8'])
])
@app.GET('/whoami')
async def whoami(ctx: HttpContext):
host, port = ctx.client
await ctx.send_str(200, f'{host}:{port}')
```
## How it works
When a request arrives:
1. If the socket peer IP does not belong to any trusted CIDR (or there is no
peer address), the mixin leaves the context untouched — `client` remains
the socket peer and all proxy headers are ignored.
2. Otherwise the client address is resolved from the headers, in order:
- `Forwarded` (RFC 7239): the `for=` entries are walked **from right to
left**, skipping entries that are themselves trusted proxies (and
`unknown`); the first untrusted entry is the client. This defeats
spoofing when the edge proxy *appends* to the header (e.g. nginx with
`$proxy_add_x_forwarded_for`), because attacker-supplied leftmost entries
are never selected. A `:port` in the selected `for=` value also
populates the port.
- `X-Forwarded-For`: same right-to-left trusted-proxy walk; the port comes
from `X-Forwarded-Port` when present and valid.
- `X-Forwarded-Host`: first entry; port from `X-Forwarded-Port` as above.
3. If none of the headers are present or usable, the socket peer is kept.
If every entry in the chain is a trusted proxy, the leftmost entry is used
(the whole chain is trusted, so the leftmost is the original client).
The resolved address is exposed by wrapping the request context /
websocket (the same pattern as `kaya-session`), so both ASGI and RSGI keep
working and `ctx.session` from other mixins is preserved.
## Note
Even with this mixin, the edge proxy should still strip or overwrite inbound
`Forwarded` / `X-Forwarded-*` headers from clients — the mixin protects the
application, the proxy protects the chain.
+56
View File
@@ -0,0 +1,56 @@
[build-system]
requires = ["setuptools>=61.0", "setuptools-scm>=8"]
build-backend = "setuptools.build_meta"
[project]
name = "kaya-forwarded"
dynamic = ["version"]
authors = [
{ name="Walter Oggioni", email="oggioni.walter@gmail.com" },
]
description = "Trusted-proxy forwarded header handling for the Kaya lightweight ASGI web framework"
readme = "README.md"
requires-python = ">=3.10"
license = "MIT"
classifiers = [
'Development Status :: 3 - Alpha',
'Topic :: Utilities',
'Intended Audience :: System Administrators',
'Intended Audience :: Developers',
'Environment :: Console',
'Programming Language :: Python :: 3',
]
dependencies = [
"kaya-core",
]
[project.optional-dependencies]
dev = [
"build", "mypy", "ipdb", "twine", "httpx", "httpx-ws", "kaya-rsgi"
]
[project.urls]
"Homepage" = "https://github.com/woggioni/kaya"
"Bug Tracker" = "https://github.com/woggioni/kaya/issues"
[tool.setuptools.packages.find]
where = ["src"]
namespaces = true
[tool.mypy]
python_version = "3.12"
disallow_untyped_defs = true
show_error_codes = true
no_implicit_optional = true
warn_return_any = true
warn_unused_ignores = true
exclude = ["scripts", "docs", "test"]
strict = true
[tool.setuptools_scm]
root = "../.."
version_file = "src/kaya/forwarded/_version.py"
[tool.setuptools_scm.tag]
prefix = "release/"
@@ -0,0 +1,10 @@
from pkgutil import extend_path
__path__ = extend_path(__path__, __name__)
from ._mixin import ForwardedHeadersMixin
__all__ = [
'ForwardedHeadersMixin',
]
@@ -0,0 +1,237 @@
from ipaddress import ip_address, ip_network, IPv4Address, IPv4Network, IPv6Address, IPv6Network
from pathlib import Path
from typing import (
Any,
AsyncGenerator,
List,
Mapping,
Optional,
Sequence,
Tuple,
Union,
)
from kaya.core import HttpContext, KayaApp, KayaMixin, WebSocket, WebSocketMessage
from kaya.core._types import StrOrStrings
_IPAddress = Union[IPv4Address, IPv6Address]
_IPNetwork = Union[IPv4Network, IPv6Network]
def _split_host_port(value: str) -> Tuple[str, Optional[int]]:
if value.startswith('['):
# bracketed IPv6 address, optionally followed by :port
closing = value.find(']')
if closing == -1:
return value, None
host = value[1:closing]
rest = value[closing + 1:]
if rest.startswith(':'):
try:
return host, int(rest[1:])
except ValueError:
return host, None
return host, None
if value.count(':') == 1:
host, _, port_str = value.rpartition(':')
try:
return host, int(port_str)
except ValueError:
return value, None
return value, None
def _parse_forwarded_entries(headers: Mapping[str, Sequence[str]]) -> List[Tuple[str, Optional[int]]]:
"""Extract the (host, port) `for=` entries of the RFC 7239 `Forwarded` header,
flattened across all header values, in chain order (leftmost = original client).
"""
entries: List[Tuple[str, Optional[int]]] = []
for raw_value in headers.get('forwarded', ()):
for element in raw_value.split(','):
for param in element.split(';'):
key, sep, value = param.partition('=')
if sep and key.strip().lower() == 'for':
for_value = value.strip().strip('"')
if for_value and for_value.lower() != 'unknown':
entries.append(_split_host_port(for_value))
break
return entries
def _first_header_value(headers: Mapping[str, Sequence[str]], name: str) -> Optional[str]:
values = headers.get(name)
if not values:
return None
first = values[0].split(',')[0].strip()
return first or None
class _ForwardedHttpContext(HttpContext):
"""HttpContext wrapper that exposes the forwarded client address.
Everything except ``client`` is delegated to the wrapped context via
``__getattr__``, so it works with any concrete ``HttpContext`` (ASGI or
RSGI) and preserves attributes set by other mixins (e.g. ``session``).
"""
def __init__(self, ctx: HttpContext, client: Tuple[str, int]) -> None:
object.__setattr__(self, '_ctx', ctx)
object.__setattr__(self, 'session', ctx.session)
object.__setattr__(self, 'client', client)
def __getattr__(self, name: str) -> Any:
if name == '_ctx':
raise AttributeError(name)
return getattr(self._ctx, name)
async def stream_body(self,
status: int,
body_generator: AsyncGenerator[bytes, None],
headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
await self._ctx.stream_body(status, body_generator, headers)
async def send_bytes(self, status: int, body: bytes, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
await self._ctx.send_bytes(status, body, headers)
async def send_str(self, status: int, body: str, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
await self._ctx.send_str(status, body, headers)
async def send_file(self, status: int, path: Path, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
await self._ctx.send_file(status, path, headers)
async def send_empty(self, status: int, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
await self._ctx.send_empty(status, headers)
class _ForwardedWebSocket(WebSocket):
"""WebSocket wrapper that exposes the forwarded client address.
Everything except ``client`` is delegated to the wrapped socket via
``__getattr__``, so it works with any concrete ``WebSocket`` (ASGI or
RSGI) and preserves attributes set by other mixins (e.g. ``session``).
"""
def __init__(self, ws: WebSocket, client: Tuple[str, int]) -> None:
object.__setattr__(self, '_ws', ws)
object.__setattr__(self, 'session', ws.session)
object.__setattr__(self, 'client', client)
def __getattr__(self, name: str) -> Any:
if name == '_ws':
raise AttributeError(name)
return getattr(self._ws, name)
async def accept(self, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
await self._ws.accept(headers)
async def receive(self) -> WebSocketMessage:
return await self._ws.receive()
async def send_text(self, data: str) -> None:
await self._ws.send_text(data)
async def send_bytes(self, data: bytes) -> None:
await self._ws.send_bytes(data)
async def close(self, code: int = 1000) -> None:
await self._ws.close(code)
async def __anext__(self) -> WebSocketMessage:
return await self._ws.__anext__()
class ForwardedHeadersMixin(KayaMixin):
"""Kaya mixin that resolves the client address from proxy headers
(``Forwarded``, ``X-Forwarded-For``, ``X-Forwarded-Host``), but only when
the direct socket peer is a trusted proxy.
Without this mixin, Kaya exposes the raw socket peer as ``ctx.client`` /
``ws.client`` and ignores forwarded headers entirely. With the mixin
applied, forwarded headers are honored only if the socket peer IP belongs
to one of the ``trusted_proxies`` CIDRs; otherwise the context is left
untouched.
When the peer is trusted, the address chain is walked from right to left
and entries that are themselves trusted proxies are skipped, so a client
that prepends a spoofed entry cannot fool the resolution when the edge
proxy appends to the header (e.g. nginx with ``$proxy_add_x_forwarded_for``).
Example::
app = KayaApp(mixins=[
ForwardedHeadersMixin(trusted_proxies=['127.0.0.1', '::1', '10.0.0.0/8'])
])
"""
def __init__(self, trusted_proxies: Sequence[str] = ()) -> None:
self._trusted_networks: Tuple[_IPNetwork, ...] = tuple(
ip_network(cidr, strict=False) for cidr in trusted_proxies
)
def apply(self, app: KayaApp) -> None:
app.add_before_request_hook(self._before_request)
app.add_before_websocket_hook(self._before_websocket)
def _is_trusted(self, host: str) -> bool:
try:
addr: _IPAddress = ip_address(host)
except ValueError:
return False
return any(addr.version == network.version and addr in network
for network in self._trusted_networks)
def _select_forwarded_entry(self, entries: List[Tuple[str, Optional[int]]]) -> Optional[Tuple[str, Optional[int]]]:
"""Walk the chain right-to-left skipping trusted proxies; the first
untrusted entry is the client. If every entry is trusted, the leftmost
(original client) is returned.
"""
for entry in reversed(entries):
if not self._is_trusted(entry[0]):
return entry
return entries[0] if entries else None
def _forwarded_port(self, headers: Mapping[str, Sequence[str]], socket_port: int) -> int:
forwarded_port = _first_header_value(headers, 'x-forwarded-port')
if forwarded_port is not None:
try:
return int(forwarded_port)
except ValueError:
pass
return socket_port
def _resolve(self,
headers: Mapping[str, Sequence[str]],
client: Optional[Tuple[str, int]]) -> Optional[Tuple[str, int]]:
if client is None or not self._is_trusted(client[0]):
return client
socket_port = client[1]
selected = self._select_forwarded_entry(_parse_forwarded_entries(headers))
if selected is not None:
host, port = selected
return host, port if port is not None else socket_port
xff_values = headers.get('x-forwarded-for')
if xff_values:
xff_entries = [entry.strip() for value in xff_values for entry in value.split(',') if entry.strip()]
selected_host = self._select_forwarded_entry([(entry, None) for entry in xff_entries])
if selected_host is not None:
return selected_host[0], self._forwarded_port(headers, socket_port)
xfh = _first_header_value(headers, 'x-forwarded-host')
if xfh is not None:
return xfh, self._forwarded_port(headers, socket_port)
return client
async def _before_request(self, ctx: HttpContext) -> Optional[HttpContext]:
resolved = self._resolve(ctx.headers, ctx.client)
if resolved is None or resolved is ctx.client:
return None
return _ForwardedHttpContext(ctx, resolved)
async def _before_websocket(self, ws: WebSocket) -> Optional[WebSocket]:
resolved = self._resolve(ws.headers, ws.client)
if resolved is None or resolved is ws.client:
return None
return _ForwardedWebSocket(ws, resolved)
Whitespace-only changes.
@@ -0,0 +1,244 @@
import asyncio
import json
import unittest
from typing import Optional
import httpx
from pwo import async_test
from kaya.core import HttpContext, KayaApp
from kaya.core._asgi import AsgiWebSocket
from kaya.forwarded import ForwardedHeadersMixin
TRUSTED = ['127.0.0.1', '::1', '10.0.0.0/8']
def make_app(trusted_proxies=TRUSTED) -> KayaApp:
mixins = [ForwardedHeadersMixin(trusted_proxies=trusted_proxies)] if trusted_proxies is not None else []
app = KayaApp(mixins=mixins)
@app.GET('/client')
async def client(ctx: HttpContext) -> None:
host, port = ctx.client if ctx.client is not None else (None, None)
await ctx.send_str(200, json.dumps({'host': host, 'port': port}))
return app
async def request(app: KayaApp,
headers: Optional[dict[str, str]] = None,
client: tuple[str, int] = ('127.0.0.1', 123)) -> dict:
transport = httpx.ASGITransport(app=app, client=client)
async with httpx.AsyncClient(transport=transport, base_url="http://127.0.0.1:80") as http_client:
r = await http_client.get('/client', headers=headers)
assert r.status_code == 200
return json.loads(r.text)
class TrustedPeerTest(unittest.TestCase):
@async_test
async def test_no_headers_returns_socket_peer(self):
app = make_app()
result = await request(app)
self.assertEqual({'host': '127.0.0.1', 'port': 123}, result)
@async_test
async def test_forwarded_header_with_port(self):
app = make_app()
result = await request(app, headers={'Forwarded': 'for=203.0.113.5:1234'})
self.assertEqual({'host': '203.0.113.5', 'port': 1234}, result)
@async_test
async def test_forwarded_header_bracketed_ipv6(self):
app = make_app()
result = await request(app, headers={'Forwarded': 'for="[2001:db8::1]:4711"'})
self.assertEqual({'host': '2001:db8::1', 'port': 4711}, result)
@async_test
async def test_forwarded_header_without_port_keeps_socket_port(self):
app = make_app()
result = await request(app, headers={'Forwarded': 'for=203.0.113.5'})
self.assertEqual({'host': '203.0.113.5', 'port': 123}, result)
@async_test
async def test_forwarded_unknown_entry_skipped(self):
app = make_app()
result = await request(app, headers={'Forwarded': 'for=unknown, for=203.0.113.5'})
self.assertEqual('203.0.113.5', result['host'])
@async_test
async def test_forwarded_rightmost_untrusted_wins(self):
# attacker-controlled leftmost entry is skipped: the rightmost
# untrusted entry (appended by the trusted edge proxy) is the client
app = make_app()
result = await request(app, headers={'Forwarded': 'for=1.2.3.4, for=5.6.7.8, for=10.0.0.2'})
self.assertEqual('5.6.7.8', result['host'])
@async_test
async def test_x_forwarded_for_spoofed_leftmost_entry_skipped(self):
# XFF = "<attacker-supplied>, <real client>" as appended by the proxy
app = make_app()
result = await request(app, headers={'X-Forwarded-For': '1.2.3.4, 5.6.7.8'})
self.assertEqual('5.6.7.8', result['host'])
@async_test
async def test_x_forwarded_for_all_trusted_chain_uses_leftmost(self):
app = make_app()
result = await request(app, headers={'X-Forwarded-For': '10.0.0.5, 10.0.0.2'})
self.assertEqual('10.0.0.5', result['host'])
@async_test
async def test_x_forwarded_port(self):
app = make_app()
result = await request(app, headers={
'X-Forwarded-For': '203.0.113.5',
'X-Forwarded-Port': '8443',
})
self.assertEqual({'host': '203.0.113.5', 'port': 8443}, result)
@async_test
async def test_invalid_x_forwarded_port_ignored(self):
app = make_app()
result = await request(app, headers={
'X-Forwarded-For': '203.0.113.5',
'X-Forwarded-Port': 'not-a-port',
})
self.assertEqual({'host': '203.0.113.5', 'port': 123}, result)
@async_test
async def test_forwarded_takes_precedence_over_x_forwarded_for(self):
app = make_app()
result = await request(app, headers={
'Forwarded': 'for=203.0.113.5',
'X-Forwarded-For': '198.51.100.7',
})
self.assertEqual('203.0.113.5', result['host'])
@async_test
async def test_x_forwarded_host_fallback(self):
app = make_app()
result = await request(app, headers={'X-Forwarded-Host': '198.51.100.7'})
self.assertEqual('198.51.100.7', result['host'])
@async_test
async def test_ipv6_cidr_trust(self):
app = make_app(trusted_proxies=['2001:db8::/32'])
result = await request(app,
headers={'X-Forwarded-For': '203.0.113.5'},
client=('2001:db8::10', 9999))
self.assertEqual({'host': '203.0.113.5', 'port': 9999}, result)
class UntrustedPeerTest(unittest.TestCase):
@async_test
async def test_untrusted_peer_ignores_forwarded_headers(self):
app = make_app()
result = await request(app,
headers={'Forwarded': 'for=1.2.3.4', 'X-Forwarded-For': '1.2.3.4'},
client=('203.0.113.99', 4567))
self.assertEqual({'host': '203.0.113.99', 'port': 4567}, result)
@async_test
async def test_empty_trusted_proxies_ignores_everything(self):
app = make_app(trusted_proxies=[])
result = await request(app, headers={'X-Forwarded-For': '1.2.3.4'})
self.assertEqual({'host': '127.0.0.1', 'port': 123}, result)
@async_test
async def test_invalid_cidr_fails_fast(self):
with self.assertRaises(ValueError):
ForwardedHeadersMixin(trusted_proxies=['not-a-cidr'])
class OptOutTest(unittest.TestCase):
@async_test
async def test_without_mixin_headers_are_ignored(self):
app = KayaApp()
@app.GET('/client')
async def client(ctx: HttpContext) -> None:
host, port = ctx.client if ctx.client is not None else (None, None)
await ctx.send_str(200, json.dumps({'host': host, 'port': port}))
result = await request(app, headers={'Forwarded': 'for=1.2.3.4', 'X-Forwarded-For': '1.2.3.4'})
self.assertEqual({'host': '127.0.0.1', 'port': 123}, result)
class WebSocketTest(unittest.TestCase):
@staticmethod
def _make_ws(headers):
async def send(message):
pass
async def receive():
return {'type': 'websocket.connect'}
scope = {
'type': 'websocket',
'path': '/ws',
'query_string': b'',
'scheme': 'ws',
'client': ('127.0.0.1', 12345),
'server': ('127.0.0.1', 80),
'headers': headers,
}
return AsgiWebSocket(scope, receive, send)
@async_test
async def test_websocket_trusted_peer(self):
mixin = ForwardedHeadersMixin(trusted_proxies=TRUSTED)
ws = self._make_ws([(b'x-forwarded-for', b'1.2.3.4, 5.6.7.8')])
wrapped = await mixin._before_websocket(ws)
assert wrapped is not None
self.assertEqual(('5.6.7.8', 12345), wrapped.client)
@async_test
async def test_websocket_untrusted_peer(self):
mixin = ForwardedHeadersMixin(trusted_proxies=['10.0.0.0/8'])
ws = self._make_ws([(b'x-forwarded-for', b'1.2.3.4')])
wrapped = await mixin._before_websocket(ws)
self.assertIsNone(wrapped)
self.assertEqual(('127.0.0.1', 12345), ws.client)
class RsgiTest(unittest.TestCase):
def test_rsgi_context(self):
from kaya.rsgi import RsgiContext
class FakeScope:
scheme = 'http'
method = 'GET'
path = '/'
query_string = ''
headers = {'x-forwarded-for': '1.2.3.4, 5.6.7.8'}
client = '127.0.0.1:12345'
server = '127.0.0.1:80'
mixin = ForwardedHeadersMixin(trusted_proxies=TRUSTED)
ctx = RsgiContext(FakeScope(), object()) # type: ignore[arg-type]
wrapped = asyncio.run(mixin._before_request(ctx))
assert wrapped is not None
self.assertEqual(('5.6.7.8', 12345), wrapped.client)
def test_rsgi_context_untrusted_peer(self):
from kaya.rsgi import RsgiContext
class FakeScope:
scheme = 'http'
method = 'GET'
path = '/'
query_string = ''
headers = {'x-forwarded-for': '1.2.3.4'}
client = '192.0.2.1:12345'
server = '127.0.0.1:80'
mixin = ForwardedHeadersMixin(trusted_proxies=TRUSTED)
ctx = RsgiContext(FakeScope(), object()) # type: ignore[arg-type]
wrapped = asyncio.run(mixin._before_request(ctx))
self.assertIsNone(wrapped)
self.assertEqual(('192.0.2.1', 12345), ctx.client)
+122
View File
@@ -40,6 +40,125 @@ 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
@@ -50,6 +169,9 @@ also works because the app applies mixin dependencies automatically.
- 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.)
## Security notes
+6 -2
View File
@@ -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',
]
+70 -13
View File
@@ -5,7 +5,7 @@ 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
@@ -42,8 +42,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 +60,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 +68,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 +142,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,11 +158,8 @@ 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())
+32 -1
View File
@@ -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)
+126 -25
View File
@@ -1,4 +1,4 @@
from typing import Any, Awaitable, Callable, Mapping, Optional, Sequence
from typing import Any, Awaitable, Callable, Literal, Mapping, Optional, Sequence, cast
from kaya.core import HttpContext, KayaApp, KayaMixin
from kaya.session import Session, SessionMixin
@@ -10,6 +10,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 +35,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 +63,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 +87,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 +122,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 +151,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,16 +174,19 @@ 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('nonce'), None)
session.pop(self._key('code_verifier'), None)
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})
@@ -150,12 +196,36 @@ class OIDCMixin(KayaMixin):
@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 +239,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 +273,32 @@ 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')
if isinstance(access_token, str):
session['oidc_access_token'] = access_token
session[self._key('access_token')] = access_token
if 'id_token' in token_response:
session['oidc_id_token'] = token_response['id_token']
session[self._key('id_token')] = token_response['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)
+628 -7
View File
@@ -1,7 +1,8 @@
import base64
import unittest
from dataclasses import replace
from time import time
from typing import Mapping
from typing import Mapping, Optional
from urllib.parse import parse_qs, urlparse
import httpx
@@ -11,17 +12,32 @@ from pwo import async_test
from kaya.core import HttpContext, KayaApp
from kaya.session import InMemorySessionStore, SessionMixin
from kaya.oidc import OIDCConfig, OIDCMixin
from kaya.oidc import OIDCConfig, OIDCMixin, OIDCProviders
from kaya.oidc._client import OIDCClient
class MockOIDCProvider:
def __init__(self) -> None:
def __init__(
self,
name: str = 'mock',
*,
issuer: Optional[str] = None,
sub: str = 'user123',
email: str = 'user@example.com',
supports_par: bool = False,
require_par: bool = False,
par_error: bool = False,
) -> None:
self.private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
self.public_key = self.private_key.public_key()
self.kid = 'mock-key'
self.issuer = 'https://mock-oidc.local'
self.nonce = 'nonce-value'
self.issuer = issuer or f'https://{name}-oidc.local'
self.sub = sub
self.email = email
self.nonce = f'{name}-nonce-value'
self.par_error = par_error
self.par_requests: list[dict[str, list[str]]] = []
self.par_auth_headers: list[Optional[str]] = []
self.discovery = {
'issuer': self.issuer,
'authorization_endpoint': f'{self.issuer}/auth',
@@ -30,6 +46,10 @@ class MockOIDCProvider:
'end_session_endpoint': f'{self.issuer}/logout',
'jwks_uri': f'{self.issuer}/jwks',
}
if supports_par or require_par:
self.discovery['pushed_authorization_request_endpoint'] = f'{self.issuer}/par'
if require_par:
self.discovery['require_pushed_authorization_requests'] = True
def _public_key_to_jwk(self) -> dict[str, str]:
numbers = self.public_key.public_numbers()
@@ -48,7 +68,7 @@ class MockOIDCProvider:
now = time()
exp = now - 3600 if expired else now + 3600
payload = {
'sub': 'user123',
'sub': self.sub,
'iss': self.issuer,
'aud': audience,
'iat': now,
@@ -70,6 +90,20 @@ class MockOIDCProvider:
return httpx.Response(200, json=self.discovery)
if path == '/jwks':
return httpx.Response(200, json={'keys': [self._public_key_to_jwk()]})
if path == '/par':
if request.method != 'POST':
return httpx.Response(405)
self.par_requests.append(parse_qs(request.content.decode()))
self.par_auth_headers.append(request.headers.get('Authorization'))
if self.par_error:
return httpx.Response(400, json={
'error': 'invalid_request',
'error_description': 'mock PAR failure',
})
return httpx.Response(201, json={
'request_uri': f'urn:ietf:params:oauth:request_uri:par-{len(self.par_requests)}',
'expires_in': 60,
})
if path == '/token' and request.method == 'POST':
body = request.content.decode()
data = dict(part.split('=') for part in body.split('&')) if body else {}
@@ -87,7 +121,7 @@ class MockOIDCProvider:
'token_type': 'Bearer',
})
if path == '/userinfo':
return httpx.Response(200, json={'sub': 'user123', 'email': 'user@example.com'})
return httpx.Response(200, json={'sub': self.sub, 'email': self.email})
if path == '/logout':
return httpx.Response(200)
return httpx.Response(404)
@@ -194,6 +228,236 @@ class OIDCClientTest(unittest.TestCase):
self.assertEqual(['id-token'], query['id_token_hint'])
self.assertEqual(['/'], query['post_logout_redirect_uri'])
@async_test
async def test_build_authorization_url_with_par(self) -> None:
provider = MockOIDCProvider('par', supports_par=True)
http_client = httpx.AsyncClient(transport=MockTransport(provider))
config = OIDCConfig(
issuer=provider.issuer,
client_id='client',
client_secret='secret',
redirect_uri='http://localhost:8000/auth/callback',
http_client=http_client,
par='always',
)
client = OIDCClient(config)
url, state, nonce, verifier = await client.build_authorization_url(par='always')
query = parse_qs(urlparse(url).query)
self.assertEqual(['client'], query['client_id'])
self.assertEqual(['urn:ietf:params:oauth:request_uri:par-1'], query['request_uri'])
self.assertNotIn('state', query)
self.assertNotIn('nonce', query)
self.assertNotIn('code_challenge', query)
self.assertEqual(1, len(provider.par_requests))
data = provider.par_requests[0]
self.assertEqual(['client'], data['client_id'])
self.assertEqual(['code'], data['response_type'])
self.assertEqual([state], data['state'])
self.assertEqual([nonce], data['nonce'])
self.assertEqual(['S256'], data['code_challenge_method'])
self.assertIn('openid', data['scope'][0])
self.assertIn('code_challenge', data)
self.assertTrue(len(verifier) > 0)
auth = provider.par_auth_headers[0]
self.assertIsNotNone(auth)
assert auth is not None
self.assertTrue(auth.startswith('Basic '))
@async_test
async def test_par_auto_uses_par_only_when_advertised(self) -> None:
provider = MockOIDCProvider('par', supports_par=True)
http_client = httpx.AsyncClient(transport=MockTransport(provider))
config = OIDCConfig(
issuer=provider.issuer,
client_id='client',
redirect_uri='http://localhost:8000/auth/callback',
http_client=http_client,
)
client = OIDCClient(config)
url, *_ = await client.build_authorization_url(par='auto')
query = parse_qs(urlparse(url).query)
self.assertIn('request_uri', query)
self.assertEqual(1, len(provider.par_requests))
# The default mock provider does not advertise a PAR endpoint.
url, *_ = await self.client.build_authorization_url(par='auto')
query = parse_qs(urlparse(url).query)
self.assertNotIn('request_uri', query)
self.assertIn('state', query)
@async_test
async def test_par_always_without_endpoint_raises(self) -> None:
with self.assertRaises(RuntimeError):
await self.client.build_authorization_url(par='always')
@async_test
async def test_required_par_with_never_raises(self) -> None:
provider = MockOIDCProvider('req', require_par=True)
http_client = httpx.AsyncClient(transport=MockTransport(provider))
config = OIDCConfig(
issuer=provider.issuer,
client_id='client',
redirect_uri='http://localhost:8000/auth/callback',
http_client=http_client,
)
client = OIDCClient(config)
with self.assertRaises(RuntimeError) as ctx:
await client.build_authorization_url(par='never')
self.assertIn('set par', str(ctx.exception))
self.assertEqual(0, len(provider.par_requests))
@async_test
async def test_required_par_with_auto_uses_par(self) -> None:
provider = MockOIDCProvider('req', require_par=True)
http_client = httpx.AsyncClient(transport=MockTransport(provider))
config = OIDCConfig(
issuer=provider.issuer,
client_id='client',
redirect_uri='http://localhost:8000/auth/callback',
http_client=http_client,
)
client = OIDCClient(config)
url, *_ = await client.build_authorization_url(par='auto')
self.assertIn('request_uri', parse_qs(urlparse(url).query))
self.assertEqual(1, len(provider.par_requests))
@async_test
async def test_par_error_response_propagates(self) -> None:
provider = MockOIDCProvider('par', supports_par=True, par_error=True)
http_client = httpx.AsyncClient(transport=MockTransport(provider))
config = OIDCConfig(
issuer=provider.issuer,
client_id='client',
redirect_uri='http://localhost:8000/auth/callback',
http_client=http_client,
par='always',
)
client = OIDCClient(config)
with self.assertRaises(httpx.HTTPStatusError):
await client.build_authorization_url(par='always')
self.assertEqual(1, len(provider.par_requests))
@async_test
async def test_par_public_client_uses_no_basic_auth(self) -> None:
provider = MockOIDCProvider('par', supports_par=True)
http_client = httpx.AsyncClient(transport=MockTransport(provider))
config = OIDCConfig(
issuer=provider.issuer,
client_id='client',
redirect_uri='http://localhost:8000/auth/callback',
http_client=http_client,
par='always',
)
client = OIDCClient(config)
await client.build_authorization_url(par='always')
self.assertIsNone(provider.par_auth_headers[0])
self.assertEqual(['client'], provider.par_requests[0]['client_id'])
def test_config_par_defaults_to_never_and_validates(self) -> None:
config = OIDCConfig(
issuer='https://issuer.local',
client_id='client',
redirect_uri='http://localhost/callback',
)
self.assertEqual('never', config.par)
with self.assertRaises(ValueError):
OIDCConfig(
issuer='https://issuer.local',
client_id='client',
redirect_uri='http://localhost/callback',
par='bogus',
)
class OIDCParTest(unittest.TestCase):
@async_test
async def test_login_and_callback_with_par(self) -> None:
provider = MockOIDCProvider('par', email='par@example.com', supports_par=True)
http_client = httpx.AsyncClient(transport=MockTransport(provider))
store = InMemorySessionStore()
session = SessionMixin(store)
config = OIDCConfig(
issuer=provider.issuer,
client_id='client',
client_secret='secret',
redirect_uri='http://localhost:8000/auth/callback',
http_client=http_client,
par='auto',
fetch_userinfo=True,
)
oidc = OIDCMixin(config, session=session)
app = KayaApp(mixins=[session, oidc])
@app.GET('/profile')
@oidc.require_auth
async def profile(ctx: HttpContext) -> None:
user = oidc.get_user(ctx)
await ctx.send_str(200, user.email if user and user.email else 'none')
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
r = await client.get('/auth/login', follow_redirects=False)
self.assertEqual(302, r.status_code)
query = parse_qs(urlparse(r.headers['Location']).query)
self.assertIn('request_uri', query)
self.assertNotIn('state', query)
self.assertEqual(1, len(provider.par_requests))
state = provider.par_requests[0]['state'][0]
provider.nonce = provider.par_requests[0]['nonce'][0]
r = await client.get('/auth/callback', params={'code': 'mock-code', 'state': state}, follow_redirects=False)
self.assertEqual(302, r.status_code)
r = await client.get('/profile')
self.assertEqual('par@example.com', r.text)
@async_test
async def test_container_applies_policy_per_provider(self) -> None:
par_provider = MockOIDCProvider('par', supports_par=True)
plain_provider = MockOIDCProvider('plain')
par_http = httpx.AsyncClient(transport=MockTransport(par_provider))
plain_http = httpx.AsyncClient(transport=MockTransport(plain_provider))
par_config = OIDCConfig(
issuer=par_provider.issuer,
client_id='par-client',
redirect_uri='http://localhost:8000/auth/withpar/callback',
http_client=par_http,
provider='withpar',
par='always',
)
plain_config = OIDCConfig(
issuer=plain_provider.issuer,
client_id='plain-client',
redirect_uri='http://localhost:8000/auth/plain/callback',
http_client=plain_http,
provider='plain',
par='never',
)
session = SessionMixin(InMemorySessionStore())
oidc = OIDCProviders([par_config, plain_config], session=session)
app = KayaApp(mixins=[session, oidc])
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
r = await client.get('/auth/withpar/login', follow_redirects=False)
self.assertEqual(302, r.status_code)
self.assertIn('request_uri', parse_qs(urlparse(r.headers['Location']).query))
r = await client.get('/auth/plain/login', follow_redirects=False)
self.assertEqual(302, r.status_code)
plain_query = parse_qs(urlparse(r.headers['Location']).query)
self.assertNotIn('request_uri', plain_query)
self.assertIn('state', plain_query)
class OIDCMixinTest(unittest.TestCase):
@@ -356,5 +620,362 @@ class OIDCMixinTest(unittest.TestCase):
self.assertIn('Set-Cookie', r.headers)
class OIDCMultiProviderTest(unittest.TestCase):
def _provider_config(self, provider: MockOIDCProvider, name: str) -> tuple[OIDCConfig, httpx.AsyncClient]:
http_client = httpx.AsyncClient(transport=MockTransport(provider))
config = OIDCConfig(
issuer=provider.issuer,
client_id=f'{name}-client',
client_secret='secret',
provider=name,
redirect_uri=f'http://localhost:8000/auth/{name}/callback',
http_client=http_client,
fetch_userinfo=True,
)
return config, http_client
def _build(self, strategy: str = 'multi-identity') -> tuple[KayaApp, OIDCProviders, InMemorySessionStore]:
self.google = MockOIDCProvider('google', sub='g-sub', email='g@example.com')
self.github = MockOIDCProvider('github', sub='h-sub', email='h@example.com')
google_config, self.google_http = self._provider_config(self.google, 'google')
github_config, self.github_http = self._provider_config(self.github, 'github')
store = InMemorySessionStore()
session = SessionMixin(store)
oidc = OIDCProviders([google_config, github_config], session=session, strategy=strategy)
app = KayaApp(mixins=[session, oidc])
@app.GET('/who')
async def who(ctx: HttpContext) -> None:
users = oidc.get_users(ctx)
await ctx.send_str(200, ','.join(sorted(users)) or 'none')
@app.GET('/google-profile')
@oidc.require_auth('google')
async def google_profile(ctx: HttpContext) -> None:
user = oidc.get_user(ctx, 'google')
assert user is not None
await ctx.send_str(200, user.email or '')
@app.GET('/github-profile')
@oidc.require_auth('github')
async def github_profile(ctx: HttpContext) -> None:
user = oidc.get_user(ctx, 'github')
assert user is not None
await ctx.send_str(200, user.email or '')
@app.GET('/refresh-github')
@oidc.require_auth('github')
async def refresh_github(ctx: HttpContext) -> None:
token = await oidc.refresh_access_token(ctx.session, 'github')
await ctx.send_str(200, token or 'no-token')
return app, oidc, store
async def _complete_login(self, client: httpx.AsyncClient, name: str, state: str) -> None:
r = await client.get(
f'/auth/{name}/callback',
params={'code': 'mock-code', 'state': state},
follow_redirects=False,
)
self.assertEqual(302, r.status_code)
async def _login(self, client: httpx.AsyncClient, provider: MockOIDCProvider, name: str) -> None:
r = await client.get(f'/auth/{name}/login', follow_redirects=False)
self.assertEqual(302, r.status_code)
query = parse_qs(urlparse(r.headers['Location']).query)
provider.nonce = query['nonce'][0]
await self._complete_login(client, name, query['state'][0])
@async_test
async def test_independent_logins(self) -> None:
app, oidc, store = self._build()
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
await self._login(client, self.google, 'google')
r = await client.get('/google-profile')
self.assertEqual('g@example.com', r.text)
# The github provider is untouched by the google login.
r = await client.get('/github-profile', follow_redirects=False)
self.assertEqual(302, r.status_code)
self.assertEqual('/auth/github/login', r.headers['Location'])
r = await client.get('/who')
self.assertEqual('google', r.text)
@async_test
async def test_interleaved_login_flows(self) -> None:
app, oidc, store = self._build()
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
# Start both flows before completing either one.
r = await client.get('/auth/google/login', follow_redirects=False)
google_query = parse_qs(urlparse(r.headers['Location']).query)
r = await client.get('/auth/github/login', follow_redirects=False)
github_query = parse_qs(urlparse(r.headers['Location']).query)
self.google.nonce = google_query['nonce'][0]
self.github.nonce = github_query['nonce'][0]
# Completing google must not invalidate the pending github state.
await self._complete_login(client, 'google', google_query['state'][0])
await self._complete_login(client, 'github', github_query['state'][0])
r = await client.get('/who')
self.assertEqual('github,google', r.text)
@async_test
async def test_logout_one_provider_keeps_the_other(self) -> None:
app, oidc, store = self._build()
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
await self._login(client, self.google, 'google')
await self._login(client, self.github, 'github')
session_id = client.cookies['session_id']
self.assertIn(session_id, store._data)
r = await client.get('/auth/github/logout', follow_redirects=False)
self.assertEqual(302, r.status_code)
self.assertIn(self.github.discovery['end_session_endpoint'], r.headers['Location'])
# Google is still logged in and the shared session survives.
r = await client.get('/google-profile')
self.assertEqual('g@example.com', r.text)
r = await client.get('/github-profile', follow_redirects=False)
self.assertEqual(302, r.status_code)
self.assertIn(session_id, store._data)
# Logging out the last provider invalidates the session.
r = await client.get('/auth/google/logout', follow_redirects=False)
self.assertEqual(302, r.status_code)
self.assertNotIn(session_id, store._data)
@async_test
async def test_provider_refresh_is_namespaced(self) -> None:
app, oidc, store = self._build()
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
await self._login(client, self.google, 'google')
await self._login(client, self.github, 'github')
# Refreshing github must not clobber google's session state.
r = await client.get('/refresh-github')
self.assertEqual('new-access-token', r.text)
r = await client.get('/google-profile')
self.assertEqual('g@example.com', r.text)
@async_test
async def test_multiple_mixin_instances_compose(self) -> None:
google = MockOIDCProvider('google', sub='g-sub', email='g@example.com')
github = MockOIDCProvider('github', sub='h-sub', email='h@example.com')
google_config, google_http = self._provider_config(google, 'google')
github_config, github_http = self._provider_config(github, 'github')
store = InMemorySessionStore()
session = SessionMixin(store)
google_mixin = OIDCMixin(
replace(google_config, login_path='/google/login', callback_path='/google/callback', logout_path='/google/logout'),
session=session,
)
github_mixin = OIDCMixin(
replace(github_config, login_path='/github/login', callback_path='/github/callback', logout_path='/github/logout'),
session=session,
)
app = KayaApp(mixins=[session, google_mixin, github_mixin])
@app.GET('/g')
@google_mixin.require_auth
async def g(ctx: HttpContext) -> None:
user = google_mixin.get_user(ctx)
assert user is not None
await ctx.send_str(200, user.email or '')
@app.GET('/h')
@github_mixin.require_auth
async def h(ctx: HttpContext) -> None:
user = github_mixin.get_user(ctx)
assert user is not None
await ctx.send_str(200, user.email or '')
@app.GET('/who')
async def who(ctx: HttpContext) -> None:
users = google_mixin.get_users(ctx)
providers = google_mixin.logged_in_providers(ctx)
summary = ','.join(f'{name}={users[name].email}' for name in sorted(users))
await ctx.send_str(200, f'{",".join(sorted(providers))}|{summary}')
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
r = await client.get('/google/login', follow_redirects=False)
query = parse_qs(urlparse(r.headers['Location']).query)
google.nonce = query['nonce'][0]
r = await client.get('/google/callback', params={'code': 'c', 'state': query['state'][0]}, follow_redirects=False)
self.assertEqual(302, r.status_code)
r = await client.get('/g')
self.assertEqual('g@example.com', r.text)
r = await client.get('/h', follow_redirects=False)
self.assertEqual(302, r.status_code)
# Mixin-level helpers see the namespaced user across providers.
r = await client.get('/who')
self.assertEqual('google|google=g@example.com', r.text)
@async_test
async def test_single_identity_evicts_previous_provider(self) -> None:
app, oidc, store = self._build(strategy='single-identity')
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
await self._login(client, self.google, 'google')
r = await client.get('/google-profile')
self.assertEqual('g@example.com', r.text)
# Logging in with github replaces the google identity.
await self._login(client, self.github, 'github')
r = await client.get('/who')
self.assertEqual('github', r.text)
r = await client.get('/google-profile', follow_redirects=False)
self.assertEqual(302, r.status_code)
self.assertEqual('/auth/google/login', r.headers['Location'])
r = await client.get('/github-profile')
self.assertEqual('h@example.com', r.text)
data = store._data[client.cookies['session_id']]
self.assertNotIn('oidc:google:user', data)
self.assertNotIn('oidc:google:id_token', data)
self.assertIn('oidc:github:user', data)
@async_test
async def test_single_identity_last_completed_login_wins(self) -> None:
app, oidc, store = self._build(strategy='single-identity')
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
# Start both flows before completing either one; pending flow state
# is not evicted, so both callbacks can complete.
r = await client.get('/auth/google/login', follow_redirects=False)
google_query = parse_qs(urlparse(r.headers['Location']).query)
r = await client.get('/auth/github/login', follow_redirects=False)
github_query = parse_qs(urlparse(r.headers['Location']).query)
self.google.nonce = google_query['nonce'][0]
self.github.nonce = github_query['nonce'][0]
await self._complete_login(client, 'google', google_query['state'][0])
await self._complete_login(client, 'github', github_query['state'][0])
r = await client.get('/who')
self.assertEqual('github', r.text)
@async_test
async def test_standalone_mixins_default_to_single_identity(self) -> None:
google = MockOIDCProvider('google', sub='g-sub', email='g@example.com')
github = MockOIDCProvider('github', sub='h-sub', email='h@example.com')
google_config, google_http = self._provider_config(google, 'google')
github_config, github_http = self._provider_config(github, 'github')
session = SessionMixin(InMemorySessionStore())
google_mixin = OIDCMixin(
replace(google_config, login_path='/google/login', callback_path='/google/callback', logout_path='/google/logout'),
session=session,
)
github_mixin = OIDCMixin(
replace(github_config, login_path='/github/login', callback_path='/github/callback', logout_path='/github/logout'),
session=session,
)
app = KayaApp(mixins=[session, google_mixin, github_mixin])
self.assertEqual('single-identity', google_mixin.strategy)
@app.GET('/g')
@google_mixin.require_auth
async def g(ctx: HttpContext) -> None:
await ctx.send_str(200, 'google')
@app.GET('/h')
@github_mixin.require_auth
async def h(ctx: HttpContext) -> None:
await ctx.send_str(200, 'github')
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
r = await client.get('/google/login', follow_redirects=False)
query = parse_qs(urlparse(r.headers['Location']).query)
google.nonce = query['nonce'][0]
r = await client.get('/google/callback', params={'code': 'c', 'state': query['state'][0]}, follow_redirects=False)
self.assertEqual(302, r.status_code)
r = await client.get('/g')
self.assertEqual('google', r.text)
r = await client.get('/github/login', follow_redirects=False)
query = parse_qs(urlparse(r.headers['Location']).query)
github.nonce = query['nonce'][0]
r = await client.get('/github/callback', params={'code': 'c', 'state': query['state'][0]}, follow_redirects=False)
self.assertEqual(302, r.status_code)
r = await client.get('/g', follow_redirects=False)
self.assertEqual(302, r.status_code)
r = await client.get('/h')
self.assertEqual('github', r.text)
def test_default_strategy_is_single_identity(self) -> None:
provider = MockOIDCProvider()
config, http = self._provider_config(provider, 'google')
session = SessionMixin(InMemorySessionStore())
providers = OIDCProviders([config], session=session)
self.assertEqual('single-identity', providers.strategy)
mixin = OIDCMixin(config, session=session)
self.assertEqual('single-identity', mixin.strategy)
def test_invalid_strategy_rejected(self) -> None:
provider = MockOIDCProvider()
config, http = self._provider_config(provider, 'google')
session = SessionMixin(InMemorySessionStore())
with self.assertRaises(ValueError):
OIDCMixin(config, session=session, strategy='bogus')
with self.assertRaises(ValueError):
OIDCProviders([config], session=session, strategy='bogus')
def test_config_rejects_invalid_provider(self) -> None:
with self.assertRaises(ValueError):
OIDCConfig(
issuer='https://issuer.local',
client_id='client',
redirect_uri='http://localhost/callback',
provider='bad:name',
)
def test_providers_reject_duplicates(self) -> None:
config = OIDCConfig(
issuer='https://issuer.local',
client_id='client',
redirect_uri='http://localhost/callback',
)
session = SessionMixin(InMemorySessionStore())
with self.assertRaises(ValueError):
OIDCProviders([config, config], session=session)
def test_providers_require_at_least_one_config(self) -> None:
session = SessionMixin(InMemorySessionStore())
with self.assertRaises(ValueError):
OIDCProviders([], session=session)
def test_providers_reject_bad_base_path(self) -> None:
config = OIDCConfig(
issuer='https://issuer.local',
client_id='client',
redirect_uri='http://localhost/callback',
)
session = SessionMixin(InMemorySessionStore())
with self.assertRaises(ValueError):
OIDCProviders([config], session=session, base_path='auth')
if __name__ == '__main__':
unittest.main()
+89
View File
@@ -0,0 +1,89 @@
# kaya-otel
OpenTelemetry tracing and metrics for the Kaya web framework.
Provides `OTelMixin`, a `KayaMixin` that instruments HTTP requests and
WebSocket connections with OpenTelemetry spans and exports HTTP server
metrics, shipping everything to an OTLP/HTTP collector.
## Usage
```python
from kaya.core import KayaApp, HttpContext
from kaya.otel import OTelMixin
app = KayaApp(mixins=[
OTelMixin(
service_name='my-service',
endpoint='http://localhost:4318',
)
])
@app.GET('/')
async def home(ctx: HttpContext):
await ctx.send_str(200, 'Hello World!')
```
## Parameters
- `service_name`: value of the `service.name` resource attribute.
- `endpoint`: base URL of an OTLP/HTTP collector; the signal paths
`/v1/traces` and `/v1/metrics` are appended. When `None`, the exporters
use their own defaults, including the standard
`OTEL_EXPORTER_OTLP_ENDPOINT` environment variable.
- `headers`: extra HTTP headers sent to the collector (e.g. authentication).
- `metric_export_interval_millis`: metric export interval
(default `60000`).
- `tracer_provider` / `meter_provider`: inject custom providers (e.g. with
in-memory exporters for tests) instead of the OTLP defaults. The mixin
only shuts down providers it created itself.
- `resource_attributes`: extra resource attributes merged with
`service.name` when the mixin creates the providers.
- `excluded_paths` / `excluded_path_regexes`: skip tracing and metrics for
matching paths (useful for health checks and metrics endpoints).
- `capture_request_headers` / `capture_response_headers`: opt-in HTTP header
capture as `http.request.header.<name>` / `http.response.header.<name>`
span attributes. Header names are normalized to lowercase with `-`
replaced by `_`; values are captured as string lists.
- `sanitize_headers`: headers captured as `REDACTED` (default:
`authorization`, `proxy-authorization`, `cookie`, `set-cookie`).
- `server_request_hook` / `server_response_hook` / `websocket_connect_hook` /
`websocket_close_hook`: optional synchronous callbacks invoked with the
span and the Kaya context/websocket at the corresponding lifecycle point.
Hook exceptions are logged and do not fail the request.
- `metrics_include_raw_path`: when `True`, metrics use raw `url.path`
attributes (legacy, potentially high cardinality). The default `False`
keeps metrics low-cardinality by using `http.route` when Kaya can resolve
a route template.
- `websocket_error_close_codes`: close codes that mark a websocket span as
failed. Defaults to protocol/application error codes such as `1002`,
`1003` and `1007`-`1011`.
## Behavior
- Every HTTP request gets a `SERVER` span named `<METHOD> <path>` with the
usual HTTP semantic attributes (`http.request.method`, `url.path`,
`client.address`, `server.address`, ...). When Kaya resolves a route
template, the span name is updated to `<METHOD> <route template>` and
`http.route` is set. The response status code is recorded as
`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, 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:
- `http.server.request.duration` (histogram, seconds), with
`http.request.method`, `http.route` when known and
`http.response.status_code` attributes.
- `http.server.active_requests` (up-down counter), with
`http.request.method` attributes.
`OTelMixin` is a `KayaMixin`, so the app stays a `KayaApp` and both ASGI and
RSGI keep working.
+58
View File
@@ -0,0 +1,58 @@
[build-system]
requires = ["setuptools>=61.0", "setuptools-scm>=8"]
build-backend = "setuptools.build_meta"
[project]
name = "kaya-otel"
dynamic = ["version"]
authors = [
{ name="Walter Oggioni", email="oggioni.walter@gmail.com" },
]
description = "OpenTelemetry tracing and metrics for the Kaya lightweight ASGI web framework"
readme = "README.md"
requires-python = ">=3.10"
license = "MIT"
classifiers = [
'Development Status :: 3 - Alpha',
'Topic :: Utilities',
'Intended Audience :: System Administrators',
'Intended Audience :: Developers',
'Environment :: Console',
'Programming Language :: Python :: 3',
]
dependencies = [
"kaya-core",
"opentelemetry-sdk",
"opentelemetry-exporter-otlp-proto-http",
]
[project.optional-dependencies]
dev = [
"build", "mypy", "ipdb", "twine", "httpx", "httpx-ws", "kaya-rsgi"
]
[project.urls]
"Homepage" = "https://github.com/woggioni/kaya"
"Bug Tracker" = "https://github.com/woggioni/kaya/issues"
[tool.setuptools.packages.find]
where = ["src"]
namespaces = true
[tool.mypy]
python_version = "3.12"
disallow_untyped_defs = true
show_error_codes = true
no_implicit_optional = true
warn_return_any = true
warn_unused_ignores = true
exclude = ["scripts", "docs", "test"]
strict = true
[tool.setuptools_scm]
root = "../.."
version_file = "src/kaya/otel/_version.py"
[tool.setuptools_scm.tag]
prefix = "release/"
@@ -0,0 +1,10 @@
from pkgutil import extend_path
__path__ = extend_path(__path__, __name__)
from ._mixin import OTelMixin
__all__ = [
'OTelMixin',
]
+518
View File
@@ -0,0 +1,518 @@
import re
import time
from asyncio import AbstractEventLoop
from logging import getLogger
from pathlib import Path
from typing import (
Any,
AsyncGenerator,
Callable,
Dict,
Mapping,
Optional,
Sequence,
Tuple,
)
from kaya.core import HttpContext, KayaApp, KayaMixin, WebSocket, WebSocketMessage
from kaya.core._types import StrOrStrings
from opentelemetry import context as otel_context
from opentelemetry import trace
from opentelemetry.context import Context
from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.metrics import Histogram, UpDownCounter
from opentelemetry.propagators.textmap import Getter
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.trace import Span, SpanKind, Status, StatusCode, Tracer
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
log = getLogger(__name__)
RequestHook = Callable[[Span, HttpContext], None]
WebSocketHook = Callable[[Span, WebSocket], None]
_DEFAULT_SANITIZE_HEADERS = (
'authorization',
'proxy-authorization',
'cookie',
'set-cookie',
)
_DEFAULT_WEBSOCKET_ERROR_CLOSE_CODES = frozenset((1002, 1003, 1007, 1008, 1009, 1010, 1011))
class _HeadersGetter(Getter[Mapping[str, Sequence[str]]]):
"""Extract propagation headers from a kaya request/websocket context.
Kaya header mappings have lowercase names with one or more values each,
which is all the W3C trace-context propagator needs.
"""
def get(self, carrier: Mapping[str, Sequence[str]], key: str) -> Optional[list[str]]:
values = carrier.get(key.lower())
if not values:
return None
return list(values)
def keys(self, carrier: Mapping[str, Sequence[str]]) -> list[str]:
return list(carrier.keys())
_PROPAGATOR = TraceContextTextMapPropagator()
_GETTER = _HeadersGetter()
def _server_attributes(ctx: HttpContext) -> Dict[str, Any]:
attributes: Dict[str, Any] = {
'http.request.method': str(ctx.method),
'url.scheme': ctx.scheme,
'url.path': ctx.path,
}
if ctx.query_string:
attributes['url.query'] = ctx.query_string
if ctx.client is not None:
attributes['client.address'] = ctx.client[0]
attributes['client.port'] = ctx.client[1]
if ctx.server is not None:
attributes['server.address'] = ctx.server[0]
if ctx.server[1] is not None:
attributes['server.port'] = ctx.server[1]
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('-', '_')
def _header_values(value: StrOrStrings | Sequence[str]) -> list[str]:
if isinstance(value, str):
return [value]
return list(value)
class _SpanHolder:
"""Mutable per-request slot shared by the tracing hooks and wrappers."""
__slots__ = ('span', 'token', 'status_code', 'start_ns', 'active_attributes', 'close_code')
def __init__(self, span: Span, token: object, active_attributes: Dict[str, Any]) -> None:
self.span = span
self.token = token
self.status_code: Optional[int] = None
self.start_ns = time.perf_counter_ns()
self.active_attributes = active_attributes
self.close_code: Optional[int] = None
class _TracedHttpContext(HttpContext):
"""HttpContext wrapper that records the response status code on the span.
Attributes not explicitly overridden are delegated to the wrapped context
via ``__getattr__``, so it works with any concrete ``HttpContext`` (ASGI
or RSGI) and preserves attributes set by other mixins (e.g. ``session``).
"""
def __init__(self, ctx: HttpContext, holder: _SpanHolder, mixin: 'OTelMixin') -> None:
object.__setattr__(self, '_ctx', ctx)
object.__setattr__(self, 'session', ctx.session)
object.__setattr__(self, '_holder', holder)
object.__setattr__(self, '_mixin', mixin)
def __getattr__(self, name: str) -> Any:
if name == '_ctx':
raise AttributeError(name)
return getattr(self._ctx, name)
def _sent(self, status: int, headers: Optional[Mapping[str, StrOrStrings]]) -> None:
self._holder.status_code = status
self._holder.span.set_attribute('http.response.status_code', status)
self._mixin._capture_response_headers(self._holder.span, headers)
async def stream_body(self,
status: int,
body_generator: AsyncGenerator[bytes, None],
headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
self._sent(status, headers)
await self._ctx.stream_body(status, body_generator, headers)
async def send_bytes(self, status: int, body: bytes, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
self._sent(status, headers)
await self._ctx.send_bytes(status, body, headers)
async def send_str(self, status: int, body: str, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
self._sent(status, headers)
await self._ctx.send_str(status, body, headers)
async def send_file(self, status: int, path: Path, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
self._sent(status, headers)
await self._ctx.send_file(status, path, headers)
async def send_empty(self, status: int, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
self._sent(status, headers)
await self._ctx.send_empty(status, headers)
class _TracedWebSocket(WebSocket):
"""WebSocket wrapper that keeps the connection span reachable.
Everything is delegated to the wrapped socket via ``__getattr__``, so it
works with any concrete ``WebSocket`` (ASGI or RSGI) and preserves
attributes set by other mixins (e.g. ``session``).
"""
def __init__(self, ws: WebSocket, holder: _SpanHolder) -> None:
object.__setattr__(self, '_ws', ws)
object.__setattr__(self, 'session', ws.session)
object.__setattr__(self, '_holder', holder)
def __getattr__(self, name: str) -> Any:
if name == '_ws':
raise AttributeError(name)
return getattr(self._ws, name)
async def accept(self, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
await self._ws.accept(headers)
async def receive(self) -> WebSocketMessage:
return await self._ws.receive()
async def send_text(self, data: str) -> None:
await self._ws.send_text(data)
async def send_bytes(self, data: bytes) -> None:
await self._ws.send_bytes(data)
async def close(self, code: int = 1000) -> None:
self._holder.close_code = code
await self._ws.close(code)
async def __anext__(self) -> WebSocketMessage:
return await self._ws.__anext__()
class OTelMixin(KayaMixin):
"""Kaya mixin adding OpenTelemetry tracing and metrics.
HTTP requests get a ``SERVER`` span named ``<METHOD> <path>`` carrying the
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,
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.
Metrics exported on the meter provider:
- ``http.server.request.duration`` (histogram, seconds), with
``http.request.method``, ``http.route`` when known and
``http.response.status_code`` attributes;
- ``http.server.active_requests`` (up-down counter), with
``http.request.method`` attributes.
Example::
app = KayaApp(mixins=[
OTelMixin(
service_name='my-service',
endpoint='http://localhost:4318',
)
])
``endpoint`` is the base URL of an OTLP/HTTP collector (the signal paths
``/v1/traces`` and ``/v1/metrics`` are appended); when ``None`` the
exporters fall back to their own defaults, including the standard
``OTEL_EXPORTER_OTLP_ENDPOINT`` environment variable. Custom
``tracer_provider``/``meter_provider`` instances (e.g. with in-memory
exporters for tests) can be injected instead of the OTLP defaults; the
mixin only shuts down providers it created itself.
"""
def __init__(self,
service_name: str,
endpoint: Optional[str] = None,
headers: Optional[Mapping[str, str]] = None,
metric_export_interval_millis: float = 60000,
tracer_provider: Optional[TracerProvider] = None,
meter_provider: Optional[MeterProvider] = None,
resource_attributes: Optional[Mapping[str, Any]] = None,
excluded_paths: Sequence[str] = (),
excluded_path_regexes: Sequence[str] = (),
capture_request_headers: Sequence[str] = (),
capture_response_headers: Sequence[str] = (),
sanitize_headers: Sequence[str] = _DEFAULT_SANITIZE_HEADERS,
server_request_hook: Optional[RequestHook] = None,
server_response_hook: Optional[RequestHook] = None,
websocket_connect_hook: Optional[WebSocketHook] = None,
websocket_close_hook: Optional[WebSocketHook] = None,
metrics_include_raw_path: bool = False,
websocket_error_close_codes: Optional[Sequence[int]] = None) -> None:
self._owns_tracer_provider = tracer_provider is None
self._owns_meter_provider = meter_provider is None
if tracer_provider is None:
tracer_provider = TracerProvider(resource=self._create_resource(service_name, resource_attributes))
span_exporter = OTLPSpanExporter(
endpoint=f'{endpoint}/v1/traces' if endpoint else None,
headers=dict(headers) if headers else None,
)
tracer_provider.add_span_processor(BatchSpanProcessor(span_exporter))
if meter_provider is None:
meter_provider = MeterProvider(
resource=self._create_resource(service_name, resource_attributes),
metric_readers=[
PeriodicExportingMetricReader(
OTLPMetricExporter(
endpoint=f'{endpoint}/v1/metrics' if endpoint else None,
headers=dict(headers) if headers else None,
),
export_interval_millis=metric_export_interval_millis,
)
],
)
self._tracer_provider = tracer_provider
self._meter_provider = meter_provider
self._tracer: Tracer = tracer_provider.get_tracer('kaya.otel')
meter = meter_provider.get_meter('kaya.otel')
self._request_duration: Histogram = meter.create_histogram(
'http.server.request.duration',
unit='s',
description='Duration of HTTP server requests',
)
self._active_requests: UpDownCounter = meter.create_up_down_counter(
'http.server.active_requests',
unit='{request}',
description='Number of in-flight HTTP server requests',
)
self._app: Optional[KayaApp] = None
self._excluded_paths = frozenset(excluded_paths)
self._excluded_path_regexes = tuple(re.compile(regex) for regex in excluded_path_regexes)
self._request_headers_to_capture = tuple(capture_request_headers)
self._response_headers_to_capture = tuple(capture_response_headers)
self._sanitize_headers = frozenset(name.lower() for name in sanitize_headers)
self._server_request_hook = server_request_hook
self._server_response_hook = server_response_hook
self._websocket_connect_hook = websocket_connect_hook
self._websocket_close_hook = websocket_close_hook
self._metrics_include_raw_path = metrics_include_raw_path
self._websocket_error_close_codes = (
frozenset(websocket_error_close_codes)
if websocket_error_close_codes is not None
else _DEFAULT_WEBSOCKET_ERROR_CLOSE_CODES
)
# Live spans keyed by the id of the (wrapped) context/websocket the
# after hooks receive; entries are removed when the span ends.
self._spans: Dict[int, _SpanHolder] = {}
@staticmethod
def _create_resource(service_name: str,
resource_attributes: Optional[Mapping[str, Any]]) -> Resource:
attributes: Dict[str, Any] = {'service.name': service_name}
if resource_attributes:
attributes.update(resource_attributes)
return Resource.create(attributes)
def apply(self, app: KayaApp) -> None:
self._app = app
app.add_before_request_hook(self._before_request)
app.add_after_request_hook(self._after_request)
app.add_before_websocket_hook(self._before_websocket)
app.add_after_websocket_hook(self._after_websocket)
def _is_excluded(self, path: str) -> bool:
if path in self._excluded_paths:
return True
return any(regex.search(path) is not None for regex in self._excluded_path_regexes)
@staticmethod
def _call_hook(hook: Optional[Callable[[Span, Any], None]], span: Span, value: Any) -> None:
if hook is None:
return
try:
hook(span, value)
except Exception:
log.exception('OpenTelemetry span hook failed')
def _capture_headers(self,
span: Span,
headers: Mapping[str, StrOrStrings | Sequence[str]],
configured: Sequence[str],
prefix: str) -> None:
if not configured:
return
lower_headers = {name.lower(): value for name, value in headers.items()}
for name in configured:
key = name.lower()
if key not in lower_headers:
continue
values = _header_values(lower_headers[key])
if key in self._sanitize_headers:
values = ['REDACTED'] * len(values)
span.set_attribute(f'{prefix}.{_normalize_header_name(name)}', values)
def _capture_request_headers(self, span: Span, headers: Mapping[str, Sequence[str]]) -> None:
self._capture_headers(span, headers, self._request_headers_to_capture, 'http.request.header')
def _capture_response_headers(self,
span: Span,
headers: Optional[Mapping[str, StrOrStrings]]) -> None:
if headers is None:
return
self._capture_headers(span, headers, self._response_headers_to_capture, 'http.response.header')
def _start_span(self, name: str, attributes: Mapping[str, Any],
headers: Mapping[str, Sequence[str]]) -> _SpanHolder:
parent: Context = _PROPAGATOR.extract(carrier=headers, getter=_GETTER)
span = self._tracer.start_span(
name,
context=parent,
kind=SpanKind.SERVER,
attributes=dict(attributes),
)
self._capture_request_headers(span, headers)
token = otel_context.attach(trace.set_span_in_context(span, parent))
return _SpanHolder(span, token, {})
@staticmethod
def _finish_span(holder: _SpanHolder, error: bool) -> None:
if error:
holder.span.set_status(Status(StatusCode.ERROR))
holder.span.end()
otel_context.detach(holder.token) # type: ignore[arg-type]
def _pop_holder(self, obj: Any) -> Optional[_SpanHolder]:
"""Find and remove the span holder for a context/websocket.
The before hooks key holders by the id of the context they received,
but after hooks run on the outermost wrapper (later mixins may have
wrapped the context again), so the id may not match. In that case the
holder is reachable through the wrapper delegation chain as
``_holder`` and the stale dict entry is reaped by identity.
"""
holder = self._spans.pop(id(obj), None)
if holder is not None:
return holder
candidate = getattr(obj, '_holder', None)
if not isinstance(candidate, _SpanHolder):
return None
for key, value in list(self._spans.items()):
if value is candidate:
del self._spans[key]
return candidate
def _route_template(self, ctx: HttpContext) -> Optional[str]:
if self._app is None:
return None
return self._app.route_template(ctx.path, ctx.method)
async def _before_request(self, ctx: HttpContext) -> Optional[HttpContext]:
if self._is_excluded(ctx.path):
return None
holder = self._start_span(
f'{ctx.method} {ctx.path}',
_server_attributes(ctx),
ctx.headers,
)
holder.active_attributes = {'http.request.method': str(ctx.method)}
if self._metrics_include_raw_path:
holder.active_attributes['url.path'] = ctx.path
self._spans[id(ctx)] = holder
self._active_requests.add(1, holder.active_attributes)
self._call_hook(self._server_request_hook, holder.span, ctx)
return _TracedHttpContext(ctx, holder, self)
async def _after_request(self, ctx: HttpContext) -> None:
holder = self._pop_holder(ctx)
if holder is None:
return
status_code = holder.status_code
exception = getattr(ctx, 'exception', None)
error = status_code is not None and status_code >= 500
if isinstance(exception, BaseException):
holder.span.record_exception(exception)
error = True
route_template = self._route_template(ctx)
if route_template is not None:
holder.span.update_name(f'{ctx.method} {route_template}')
holder.span.set_attribute('http.route', route_template)
self._call_hook(self._server_response_hook, holder.span, ctx)
self._finish_span(holder, error)
self._active_requests.add(-1, holder.active_attributes)
metric_attributes: Dict[str, Any] = {
'http.request.method': str(ctx.method),
}
if self._metrics_include_raw_path:
metric_attributes['url.path'] = ctx.path
elif route_template is not None:
metric_attributes['http.route'] = route_template
if status_code is not None:
metric_attributes['http.response.status_code'] = status_code
duration = (time.perf_counter_ns() - holder.start_ns) / 1e9
self._request_duration.record(duration, metric_attributes)
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}', _websocket_attributes(ws), ws.headers)
self._spans[id(ws)] = holder
self._call_hook(self._websocket_connect_hook, holder.span, ws)
return _TracedWebSocket(ws, holder)
async def _after_websocket(self, ws: WebSocket) -> None:
holder = self._pop_holder(ws)
if holder is None:
return
exception = getattr(ws, 'exception', None)
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)
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)
def shutdown(self, loop: AbstractEventLoop) -> None:
# Only providers created by this mixin are shut down; injected ones
# stay under the caller's control.
if self._owns_tracer_provider:
self._tracer_provider.shutdown()
if self._owns_meter_provider:
self._meter_provider.shutdown()
Whitespace-only changes.
+420
View File
@@ -0,0 +1,420 @@
import json
import unittest
from typing import Any, Optional, Tuple
import httpx
from pwo import async_test
from kaya.core import HttpContext, KayaApp
from kaya.core._asgi import AsgiWebSocket
from kaya.otel import OTelMixin
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import InMemoryMetricReader
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.trace import StatusCode
TRACEPARENT = '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01'
def make_mixin(**kwargs) -> Tuple[OTelMixin, InMemorySpanExporter, InMemoryMetricReader]:
resource = Resource.create({'service.name': 'test-service'})
span_exporter = InMemorySpanExporter()
tracer_provider = TracerProvider(resource=resource)
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
metric_reader = InMemoryMetricReader()
meter_provider = MeterProvider(resource=resource, metric_readers=[metric_reader])
mixin = OTelMixin(
service_name='test-service',
tracer_provider=tracer_provider,
meter_provider=meter_provider,
**kwargs,
)
return mixin, span_exporter, metric_reader
def make_app(mixin: OTelMixin) -> KayaApp:
app = KayaApp(mixins=[mixin])
@app.GET('/hello')
async def hello(ctx: HttpContext) -> None:
await ctx.send_str(200, json.dumps({'ok': True}))
@app.GET('/boom')
async def boom(ctx: HttpContext) -> None:
await ctx.send_str(500, 'boom')
return app
async def request(app: KayaApp, path: str = '/hello', headers: Optional[dict[str, str]] = None) -> httpx.Response:
transport = httpx.ASGITransport(app=app, client=('127.0.0.1', 123))
async with httpx.AsyncClient(transport=transport, base_url="http://127.0.0.1:80") as http_client:
return await http_client.get(path, headers=headers)
class HttpTracingTest(unittest.TestCase):
@async_test
async def test_request_produces_server_span(self):
mixin, span_exporter, _ = make_mixin()
app = make_app(mixin)
r = await request(app)
self.assertEqual(200, r.status_code)
spans = span_exporter.get_finished_spans()
self.assertEqual(1, len(spans))
span = spans[0]
self.assertEqual('GET /hello', span.name)
assert span.attributes is not None
self.assertEqual('GET', span.attributes['http.request.method'])
self.assertEqual('/hello', span.attributes['url.path'])
self.assertEqual(200, span.attributes['http.response.status_code'])
self.assertEqual(StatusCode.UNSET, span.status.status_code)
@async_test
async def test_5xx_marks_span_as_error(self):
mixin, span_exporter, _ = make_mixin()
app = make_app(mixin)
r = await request(app, '/boom')
self.assertEqual(500, r.status_code)
spans = span_exporter.get_finished_spans()
self.assertEqual(1, len(spans))
self.assertEqual(500, spans[0].attributes['http.response.status_code']) # type: ignore[index]
self.assertEqual(StatusCode.ERROR, spans[0].status.status_code)
@async_test
async def test_traceparent_header_propagates(self):
mixin, span_exporter, _ = make_mixin()
app = make_app(mixin)
await request(app, headers={'traceparent': TRACEPARENT})
spans = span_exporter.get_finished_spans()
self.assertEqual(1, len(spans))
span = spans[0]
self.assertEqual(0x4bf92f3577b34da6a3ce929d0e0e4736, span.context.trace_id) # type: ignore[union-attr]
assert span.parent is not None
self.assertEqual(0x00f067aa0ba902b7, span.parent.span_id)
@async_test
async def test_unmatched_route_still_traced(self):
mixin, span_exporter, _ = make_mixin()
app = make_app(mixin)
r = await request(app, '/nowhere')
self.assertEqual(404, r.status_code)
spans = span_exporter.get_finished_spans()
self.assertEqual(1, len(spans))
self.assertEqual(404, spans[0].attributes['http.response.status_code']) # type: ignore[index]
@async_test
async def test_no_leaked_spans_after_requests(self):
mixin, span_exporter, _ = make_mixin()
app = make_app(mixin)
await request(app)
await request(app, '/boom')
self.assertEqual(0, len(mixin._spans))
self.assertEqual(2, len(span_exporter.get_finished_spans()))
@async_test
async def test_exception_marks_span_as_error(self):
mixin, span_exporter, _ = make_mixin()
app = KayaApp(mixins=[mixin])
@app.GET('/raises')
async def raises(ctx: HttpContext) -> None:
raise RuntimeError('boom')
transport = httpx.ASGITransport(app=app, client=('127.0.0.1', 123))
async with httpx.AsyncClient(transport=transport, base_url="http://127.0.0.1:80") as http_client:
with self.assertRaises(RuntimeError):
await http_client.get('/raises')
spans = span_exporter.get_finished_spans()
self.assertEqual(1, len(spans))
span = spans[0]
self.assertEqual(StatusCode.ERROR, span.status.status_code)
self.assertTrue(any(event.name == 'exception' for event in span.events))
self.assertEqual(0, len(mixin._spans))
@async_test
async def test_excluded_path_skips_tracing_and_metrics(self):
mixin, span_exporter, metric_reader = make_mixin(excluded_paths=('/health',))
app = KayaApp(mixins=[mixin])
@app.GET('/health')
async def health(ctx: HttpContext) -> None:
await ctx.send_str(200, 'ok')
@app.GET('/hello')
async def hello(ctx: HttpContext) -> None:
await ctx.send_str(200, 'hello')
await request(app, '/health')
await request(app, '/hello')
spans = span_exporter.get_finished_spans()
self.assertEqual(1, len(spans))
self.assertEqual('GET /hello', spans[0].name)
metrics = MetricsTest._metric_names(metric_reader)
duration_points = list(metrics['http.server.request.duration'].data.data_points)
self.assertEqual(1, len(duration_points))
self.assertEqual(1, duration_points[0].count)
@async_test
async def test_route_template_used_for_span_name_and_metrics(self):
mixin, span_exporter, metric_reader = make_mixin()
app = KayaApp(mixins=[mixin])
@app.GET('/items/${item_id:int}')
async def item(ctx: HttpContext, item_id: int) -> None:
await ctx.send_str(200, str(item_id))
r = await request(app, '/items/123')
self.assertEqual(200, r.status_code)
spans = span_exporter.get_finished_spans()
self.assertEqual(1, len(spans))
span = spans[0]
self.assertEqual('GET /items/${item_id:int}', span.name)
assert span.attributes is not None
self.assertEqual('/items/${item_id:int}', span.attributes['http.route'])
self.assertEqual('/items/123', span.attributes['url.path'])
metrics = MetricsTest._metric_names(metric_reader)
duration_points = list(metrics['http.server.request.duration'].data.data_points)
self.assertEqual(1, len(duration_points))
duration_attributes = dict(duration_points[0].attributes)
self.assertEqual('GET', duration_attributes['http.request.method'])
self.assertEqual('/items/${item_id:int}', duration_attributes['http.route'])
self.assertEqual(200, duration_attributes['http.response.status_code'])
self.assertNotIn('url.path', duration_attributes)
active_points = list(metrics['http.server.active_requests'].data.data_points)
self.assertEqual(1, len(active_points))
self.assertEqual({'http.request.method': 'GET'}, dict(active_points[0].attributes))
@async_test
async def test_header_capture_and_sanitization(self):
mixin, span_exporter, _ = make_mixin(
capture_request_headers=('X-Request-Id', 'Authorization'),
capture_response_headers=('X-Response-Id', 'Set-Cookie'),
)
app = KayaApp(mixins=[mixin])
@app.GET('/headers')
async def headers(ctx: HttpContext) -> None:
await ctx.send_str(200, 'ok', {
'X-Response-Id': 'res-1',
'Set-Cookie': 'sid=secret',
})
r = await request(app, '/headers', headers={
'X-Request-Id': 'req-1',
'Authorization': 'Bearer secret',
})
self.assertEqual(200, r.status_code)
spans = span_exporter.get_finished_spans()
self.assertEqual(1, len(spans))
attributes = spans[0].attributes
assert attributes is not None
self.assertEqual(['req-1'], list(attributes['http.request.header.x_request_id']))
self.assertEqual(['REDACTED'], list(attributes['http.request.header.authorization']))
self.assertEqual(['res-1'], list(attributes['http.response.header.x_response_id']))
self.assertEqual(['REDACTED'], list(attributes['http.response.header.set_cookie']))
@async_test
async def test_http_hooks_are_called(self):
calls = []
def request_hook(span: Any, ctx: HttpContext) -> None:
calls.append(('request', ctx.path))
span.set_attribute('test.request_hook', True)
def response_hook(span: Any, ctx: HttpContext) -> None:
calls.append(('response', ctx.path))
span.set_attribute('test.response_hook', True)
mixin, span_exporter, _ = make_mixin(
server_request_hook=request_hook,
server_response_hook=response_hook,
)
app = make_app(mixin)
await request(app)
self.assertEqual([('request', '/hello'), ('response', '/hello')], calls)
attributes = span_exporter.get_finished_spans()[0].attributes
assert attributes is not None
self.assertTrue(attributes['test.request_hook'])
self.assertTrue(attributes['test.response_hook'])
class MetricsTest(unittest.TestCase):
@staticmethod
def _metric_names(metric_reader: InMemoryMetricReader) -> dict:
data = metric_reader.get_metrics_data()
return {m.name: m for rm in data.resource_metrics for m in rm.scope_metrics[0].metrics
if rm.scope_metrics}
@async_test
async def test_duration_histogram_and_active_requests(self):
mixin, _, metric_reader = make_mixin()
app = make_app(mixin)
await request(app)
metrics = self._metric_names(metric_reader)
self.assertIn('http.server.request.duration', metrics)
self.assertIn('http.server.active_requests', metrics)
duration_points = list(metrics['http.server.request.duration'].data.data_points)
self.assertEqual(1, len(duration_points))
self.assertEqual(1, duration_points[0].count)
self.assertGreaterEqual(duration_points[0].sum, 0)
active_points = list(metrics['http.server.active_requests'].data.data_points)
self.assertEqual(1, len(active_points))
self.assertEqual(0, active_points[0].value)
class WebSocketTracingTest(unittest.TestCase):
@staticmethod
def _make_ws(headers=(), query_string: bytes = b''):
async def send(message):
pass
async def receive():
return {'type': 'websocket.connect'}
scope = {
'type': 'websocket',
'path': '/ws/games/abc',
'query_string': query_string,
'scheme': 'ws',
'client': ('127.0.0.1', 12345),
'server': ('127.0.0.1', 80),
'headers': headers,
}
return AsgiWebSocket(scope, receive, send)
@async_test
async def test_websocket_lifecycle_produces_span(self):
mixin, span_exporter, _ = make_mixin()
ws = self._make_ws()
wrapped = await mixin._before_websocket(ws)
self.assertIsNotNone(wrapped)
self.assertEqual(1, len(mixin._spans))
await mixin._after_websocket(wrapped) # type: ignore[arg-type]
self.assertEqual(0, len(mixin._spans))
spans = span_exporter.get_finished_spans()
self.assertEqual(1, len(spans))
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):
mixin, span_exporter, _ = make_mixin()
ws = self._make_ws(headers=[(b'traceparent', TRACEPARENT.encode())])
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))
self.assertEqual(0x4bf92f3577b34da6a3ce929d0e0e4736, spans[0].context.trace_id) # type: ignore[union-attr]
@async_test
async def test_wrapped_websocket_delegates(self):
mixin, _, _ = make_mixin()
ws = self._make_ws()
wrapped = await mixin._before_websocket(ws)
assert wrapped is not None
self.assertEqual('/ws/games/abc', wrapped.path)
self.assertEqual(('127.0.0.1', 12345), wrapped.client)
await mixin._after_websocket(wrapped)
@async_test
async def test_websocket_error_close_code_marks_span_error(self):
mixin, span_exporter, _ = make_mixin()
ws = self._make_ws()
wrapped = await mixin._before_websocket(ws)
assert wrapped is not None
await wrapped.close(1011)
await mixin._after_websocket(wrapped)
spans = span_exporter.get_finished_spans()
self.assertEqual(1, len(spans))
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):
mixin, span_exporter, _ = make_mixin()
ws = self._make_ws()
wrapped = await mixin._before_websocket(ws)
assert wrapped is not None
wrapped.exception = RuntimeError('ws boom')
await mixin._after_websocket(wrapped)
spans = span_exporter.get_finished_spans()
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):
calls = []
def connect_hook(span: Any, ws: Any) -> None:
calls.append(('connect', ws.path))
span.set_attribute('test.websocket_connect_hook', True)
def close_hook(span: Any, ws: Any) -> None:
calls.append(('close', ws.path))
span.set_attribute('test.websocket_close_hook', True)
mixin, span_exporter, _ = make_mixin(
websocket_connect_hook=connect_hook,
websocket_close_hook=close_hook,
)
ws = self._make_ws()
wrapped = await mixin._before_websocket(ws)
assert wrapped is not None
await wrapped.close(1000)
await mixin._after_websocket(wrapped)
self.assertEqual([('connect', '/ws/games/abc'), ('close', '/ws/games/abc')], calls)
spans = span_exporter.get_finished_spans()
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'])
if __name__ == '__main__':
unittest.main()
+45 -5
View File
@@ -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, resolve_client
from kaya.core import (
AbstractKayaApp,
HttpContext,
HttpMethod,
WebSocket,
WebSocketDisconnect,
WebSocketMessage,
)
from kaya.core._types import StrOrStrings
@@ -51,8 +59,7 @@ class RsgiContext(HttpContext):
fun = cast(Callable[[Mapping[str, Sequence[str]], tuple[str, str]], Mapping[str, Sequence[str]]], acc)
self.headers = reduce(fun, scope.headers.items(), {})
self.client = resolve_client(self.headers,
Maybe.of(scope.client.split(':'))
self.client = (Maybe.of(scope.client.split(':'))
.map(lambda it: (it[0], int(it[1])))
.or_else_throw(RuntimeError))
self.server = (Maybe.of(scope.server.split(':'))
@@ -110,6 +117,7 @@ class RsgiContext(HttpContext):
class RsgiWebSocket(WebSocket):
_protocol: RSGIWebsocketProtocol
_transport: Optional[RSGIWebsocketTransport]
_closed: bool
scheme: str
path: str
query_string: str
@@ -125,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
@@ -135,8 +144,7 @@ class RsgiWebSocket(WebSocket):
fun = cast(Callable[[Mapping[str, Sequence[str]], tuple[str, str]], Mapping[str, Sequence[str]]], acc)
self.headers = reduce(fun, scope.headers.items(), {})
self.client = resolve_client(self.headers,
Maybe.of(scope.client.split(':'))
self.client = (Maybe.of(scope.client.split(':'))
.map(lambda it: (it[0], int(it[1])))
.or_else_throw(RuntimeError))
self.server = (Maybe.of(scope.server.split(':'))
@@ -146,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.
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')
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)
@@ -164,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')
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')
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:
if self._closed:
return
try:
self._protocol.close(code)
except RSGIProtocolClosed:
self._closed = True
async def __anext__(self) -> WebSocketMessage:
message = await self.receive()
+130 -39
View File
@@ -1,59 +1,150 @@
import unittest
from kaya.rsgi import RsgiContext, RsgiWebSocket
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
class RsgiContextTest(unittest.TestCase):
@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)
@staticmethod
def _make_context(headers):
class FakeScope:
scheme = 'http'
method = 'GET'
path = '/'
query_string = ''
client = '127.0.0.1:12345'
server = '127.0.0.1:80'
@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)
def __init__(self, headers):
self.headers = headers
@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)
return RsgiContext(FakeScope(headers), object()) # type: ignore[arg-type]
@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)
def test_forwarded_header(self):
ctx = self._make_context({'forwarded': 'for=203.0.113.5:1234'})
self.assertEqual(('203.0.113.5', 1234), ctx.client)
@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)
def test_x_forwarded_headers(self):
ctx = self._make_context({
'x-forwarded-for': '203.0.113.5, 70.41.3.18',
'x-forwarded-port': '8443',
})
self.assertEqual(('203.0.113.5', 8443), ctx.client)
@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)
def test_x_forwarded_host_fallback(self):
ctx = self._make_context({'x-forwarded-host': '198.51.100.7'})
self.assertEqual(('198.51.100.7', 12345), ctx.client)
@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)
def test_socket_peer_fallback(self):
ctx = self._make_context({})
self.assertEqual(('127.0.0.1', 12345), ctx.client)
@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)
+2
View File
@@ -6,6 +6,8 @@ kaya-session-memcache @ file:./packages/kaya-session-memcache
kaya-oidc @ file:./packages/kaya-oidc
kaya-openapi @ file:./packages/kaya-openapi
kaya-cors @ file:./packages/kaya-cors
kaya-forwarded @ file:./packages/kaya-forwarded
kaya-otel @ file:./packages/kaya-otel
build
fakeredis
mypy
+36
View File
@@ -44,6 +44,8 @@ executing==2.2.1
# via stack-data
fakeredis==2.36.2
# via -r requirements-dev.in
googleapis-common-protos==1.75.3
# via opentelemetry-exporter-otlp-proto-http
granian==2.7.9
# via kaya-rsgi
h11==0.16.0
@@ -90,16 +92,22 @@ file:./packages/kaya-core
# via
# -r requirements-dev.in
# kaya-cors
# kaya-forwarded
# kaya-oidc
# kaya-openapi
# kaya-otel
# kaya-rsgi
# kaya-session
file:./packages/kaya-cors
# via -r requirements-dev.in
file:./packages/kaya-forwarded
# via -r requirements-dev.in
file:./packages/kaya-oidc
# via -r requirements-dev.in
file:./packages/kaya-openapi
# via -r requirements-dev.in
file:./packages/kaya-otel
# via -r requirements-dev.in
file:./packages/kaya-rsgi
# via -r requirements-dev.in
file:./packages/kaya-session
@@ -132,6 +140,25 @@ mypy-extensions==1.1.0
# via mypy
nh3==0.3.6
# via readme-renderer
opentelemetry-api==1.44.0
# via
# opentelemetry-exporter-otlp-proto-http
# opentelemetry-sdk
# opentelemetry-semantic-conventions
opentelemetry-exporter-otlp-proto-common==1.44.0
# via opentelemetry-exporter-otlp-proto-http
opentelemetry-exporter-otlp-proto-http==1.44.0
# via kaya-otel
opentelemetry-proto==1.44.0
# via
# opentelemetry-exporter-otlp-proto-common
# opentelemetry-exporter-otlp-proto-http
opentelemetry-sdk==1.44.0
# via
# kaya-otel
# opentelemetry-exporter-otlp-proto-http
opentelemetry-semantic-conventions==0.65b0
# via opentelemetry-sdk
packaging==26.2
# via
# build
@@ -144,6 +171,10 @@ pexpect==4.9.0
# via ipython
prompt-toolkit==3.0.52
# via ipython
protobuf==7.36.2
# via
# googleapis-common-protos
# opentelemetry-proto
psutil==7.2.2
# via ipython
ptyprocess==0.7.0
@@ -175,6 +206,7 @@ redis==8.0.1
# kaya-session-redis
requests==2.34.2
# via
# opentelemetry-exporter-otlp-proto-http
# requests-toolbelt
# twine
requests-toolbelt==1.0.0
@@ -199,6 +231,10 @@ typing-extensions==4.16.0
# via
# kaya-core
# mypy
# opentelemetry-api
# opentelemetry-exporter-otlp-proto-http
# opentelemetry-sdk
# opentelemetry-semantic-conventions
# pwo
urllib3==2.7.0
# via