Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e73c4b2c18
|
||
|
|
fa8bb0487d
|
||
|
|
15dc4b6dad
|
||
|
|
3af2e7a975
|
||
|
|
540a7b111e
|
No files matched your search
@@ -40,20 +40,174 @@ async def profile(ctx: HttpContext):
|
||||
`OIDCMixin` depends on `SessionMixin`; passing only `oidc` to `KayaApp(mixins=...)`
|
||||
also works because the app applies mixin dependencies automatically.
|
||||
|
||||
## Multiple providers
|
||||
|
||||
Use `OIDCProviders` to run several OIDC providers in one application. Each
|
||||
config's `provider` becomes a route segment under `base_path`
|
||||
(`/auth/{provider}/login`, `/auth/{provider}/callback` and
|
||||
`/auth/{provider}/logout`), and all providers share a single `SessionMixin`:
|
||||
|
||||
```python
|
||||
from kaya.oidc import OIDCConfig, OIDCProviders
|
||||
|
||||
session = SessionMixin(InMemorySessionStore())
|
||||
oidc = OIDCProviders(
|
||||
[
|
||||
OIDCConfig(
|
||||
issuer='https://accounts.google.com',
|
||||
client_id=os.environ['GOOGLE_CLIENT_ID'],
|
||||
client_secret=os.environ['GOOGLE_CLIENT_SECRET'],
|
||||
redirect_uri='http://localhost:8000/auth/google/callback',
|
||||
provider='google',
|
||||
),
|
||||
OIDCConfig(
|
||||
issuer='https://github.com/login/oauth', # or any OIDC issuer
|
||||
client_id=os.environ['GITHUB_CLIENT_ID'],
|
||||
client_secret=os.environ['GITHUB_CLIENT_SECRET'],
|
||||
redirect_uri='http://localhost:8000/auth/github/callback',
|
||||
provider='github',
|
||||
),
|
||||
],
|
||||
session=session,
|
||||
base_path='/auth', # optional
|
||||
)
|
||||
app = KayaApp(mixins=[session, oidc])
|
||||
|
||||
@app.GET('/google-profile')
|
||||
@oidc.require_auth('google')
|
||||
async def google_profile(ctx: HttpContext):
|
||||
user = oidc.get_user(ctx, 'google')
|
||||
await ctx.send_str(200, f"Hello {user.email or user.sub}")
|
||||
|
||||
@app.GET('/who')
|
||||
async def who(ctx: HttpContext):
|
||||
users = oidc.get_users(ctx) # {provider: OIDCUser}
|
||||
providers = oidc.logged_in_providers(ctx) # ['google', ...]
|
||||
await ctx.send_str(200, ', '.join(users))
|
||||
```
|
||||
|
||||
Each provider's session state is namespaced as `oidc:{provider}:...`. Make sure
|
||||
each config's `redirect_uri` points at that provider's derived callback path.
|
||||
|
||||
### Identity strategy
|
||||
|
||||
`strategy` controls how many identities may coexist in the shared session:
|
||||
|
||||
- `'single-identity'` (default): completing a login evicts every other
|
||||
provider's identity and tokens, so the session holds exactly one user — last
|
||||
login wins, like Spring Security. After a GitHub login, `get_user(ctx,
|
||||
'google')` returns `None` and `require_auth('google')` redirects to the Google
|
||||
login.
|
||||
- `'multi-identity'`: each provider stays logged in independently. Use it for
|
||||
linked accounts, aggregating several provider APIs, or step-up flows. Logging
|
||||
out of one provider keeps the shared session alive while another provider is
|
||||
still logged in.
|
||||
|
||||
```python
|
||||
oidc = OIDCProviders([google_config, github_config], session=session,
|
||||
strategy='multi-identity')
|
||||
```
|
||||
|
||||
Ephemeral flow keys (`state`, `nonce`, `code_verifier`) are never evicted, so
|
||||
logins already in flight can still complete; under `'single-identity'` the last
|
||||
callback to finish wins. Under either strategy, `get_users(ctx)` only reports
|
||||
providers actually present in the session.
|
||||
|
||||
The composable alternative is to create one `OIDCMixin` per provider, pointing
|
||||
each at distinct `login_path`/`callback_path`/`logout_path` values and sharing
|
||||
the same `SessionMixin`. The same `strategy` keyword is accepted and defaults to
|
||||
`'single-identity'`:
|
||||
|
||||
```python
|
||||
google = OIDCMixin(google_config, session=session)
|
||||
github = OIDCMixin(github_config, session=session, strategy='multi-identity')
|
||||
app = KayaApp(mixins=[session, google, github])
|
||||
```
|
||||
|
||||
> **Upgrading:** session keys are now prefixed with `oidc:{provider}:` (the
|
||||
> default provider is `default`). Existing logged-in users are signed out once
|
||||
> and simply re-authenticate.
|
||||
|
||||
## Pushed authorization requests (PAR)
|
||||
|
||||
`OIDCConfig.par` controls RFC 9126 Pushed Authorization Requests:
|
||||
|
||||
- `'never'` (default): the authorization request travels as query parameters,
|
||||
as before.
|
||||
- `'auto'`: push the request whenever the provider advertises a
|
||||
`pushed_authorization_request_endpoint`; otherwise fall back to the plain
|
||||
redirect.
|
||||
- `'always'`: always push; login raises a `RuntimeError` if the provider does
|
||||
not support PAR.
|
||||
|
||||
```python
|
||||
OIDCConfig(..., par='auto')
|
||||
```
|
||||
|
||||
A provider that advertises `require_pushed_authorization_requests: true` must be
|
||||
configured with `par='auto'` or `par='always'`. With `par='never'` the login
|
||||
fails fast with a clear error instead of issuing a redirect the provider would
|
||||
reject.
|
||||
|
||||
The same parameters as the plain authorization request (PKCE challenge,
|
||||
`state`, `nonce`, scopes, ...) are POSTed directly to the provider, and the
|
||||
browser is redirected with only `client_id` and the single-use `request_uri`.
|
||||
Request parameters therefore no longer travel through the user agent. PKCE,
|
||||
`state`, and `nonce` are still used and protect against request URI swapping.
|
||||
|
||||
Client authentication follows the token endpoint: HTTP Basic when
|
||||
`client_secret` is set, or none for public clients. PAR with other client
|
||||
authentication methods (`private_key_jwt`, mTLS) is not supported.
|
||||
|
||||
## Features
|
||||
|
||||
- Generic OIDC discovery
|
||||
- Authorization Code Flow with PKCE (S256)
|
||||
- ID token signature validation with JWKS
|
||||
- Issuer validation against discovery metadata, including support for
|
||||
Microsoft Entra ID `{tenantid}` issuer templates
|
||||
- Refreshed ID token validation on token refresh
|
||||
- `state` and `nonce` protection
|
||||
- Session fixation defense via `regenerate_id()` after login
|
||||
- Optional userinfo endpoint fetch
|
||||
- Refresh token support
|
||||
- RP-initiated logout (when provider advertises `end_session_endpoint`)
|
||||
- Pushed Authorization Requests (RFC 9126) via `OIDCConfig.par`
|
||||
- Multiple providers in one app via `OIDCProviders` (or several `OIDCMixin` instances)
|
||||
- `single-identity` (default) or `multi-identity` session strategy
|
||||
- Composable with any other `KayaMixin` (RSGI, MCP, etc.)
|
||||
|
||||
## Microsoft Entra ID
|
||||
|
||||
Entra ID v2 ID tokens carry an `iss` claim of the form
|
||||
`https://login.microsoftonline.com/<tenant-id>/v2.0`, while the discovery
|
||||
metadata for the multi-tenant endpoints advertises the templated issuer
|
||||
`https://login.microsoftonline.com/{tenantid}/v2.0`. `kaya-oidc` validates the
|
||||
`iss` claim against the discovered metadata and substitutes the token's `tid`
|
||||
claim into `{tenantid}`, so both tenant-specific and multi-tenant issuers work:
|
||||
|
||||
```python
|
||||
OIDCConfig(
|
||||
issuer='https://login.microsoftonline.com/<tenant-id-or-domain>/v2.0',
|
||||
client_id=os.environ['OIDC_CLIENT_ID'],
|
||||
client_secret=os.environ['OIDC_CLIENT_SECRET'],
|
||||
redirect_uri='https://example.com/auth/callback',
|
||||
)
|
||||
```
|
||||
|
||||
Use a tenant-specific issuer (the tenant GUID or a verified domain name) to
|
||||
restrict logins to that tenant. The `common` and `organizations` aliases are
|
||||
multi-tenant: any tenant in that cloud can obtain an ID token that passes
|
||||
validation. If you build a multi-tenant application, check the `tid` claim
|
||||
against the tenants you trust.
|
||||
|
||||
## Security notes
|
||||
|
||||
- The `iss` claim is validated against the `issuer` value of the provider's
|
||||
discovery document, not against `OIDCConfig.issuer` (they can differ, for
|
||||
example for Entra ID multi-tenant endpoints).
|
||||
- Refreshed ID tokens are validated before being stored; the `nonce` is also
|
||||
checked when the refreshed token carries one.
|
||||
- The `none` signing algorithm is rejected by default.
|
||||
- Only algorithms listed in `OIDCConfig.allowed_id_token_algorithms` are accepted.
|
||||
- Always use HTTPS for `redirect_uri` in production.
|
||||
|
||||
@@ -1,11 +1,15 @@
|
||||
from ._client import OIDCClient
|
||||
from ._config import OIDCConfig
|
||||
from ._mixin import OIDCMixin, OIDCUser
|
||||
from ._config import OIDCConfig, PARPolicy
|
||||
from ._mixin import IdentityStrategy, OIDCMixin, OIDCUser
|
||||
from ._providers import OIDCProviders
|
||||
|
||||
|
||||
__all__ = [
|
||||
'IdentityStrategy',
|
||||
'OIDCClient',
|
||||
'OIDCConfig',
|
||||
'OIDCMixin',
|
||||
'OIDCProviders',
|
||||
'OIDCUser',
|
||||
'PARPolicy',
|
||||
]
|
||||
@@ -1,3 +1,4 @@
|
||||
import re
|
||||
from typing import Any, Mapping, Optional, Sequence, cast
|
||||
from urllib.parse import urlencode
|
||||
|
||||
@@ -5,9 +6,11 @@ import httpx
|
||||
import jwt
|
||||
from jwt import PyJWK
|
||||
|
||||
from ._config import OIDCConfig
|
||||
from ._config import PARPolicy, OIDCConfig
|
||||
from ._utils import generate_nonce, generate_pkce, generate_state
|
||||
|
||||
_TENANT_ID_PLACEHOLDER = re.compile(r'\{tenantid\}', re.IGNORECASE)
|
||||
|
||||
|
||||
class OIDCClient:
|
||||
"""Low-level OIDC client implementing discovery, PKCE, and token validation."""
|
||||
@@ -42,8 +45,14 @@ class OIDCClient:
|
||||
async with httpx.AsyncClient() as client:
|
||||
return await client.request(method, url, **kwargs)
|
||||
|
||||
async def build_authorization_url(self) -> tuple[str, str, str, str]:
|
||||
"""Return (authorization_url, state, nonce, code_verifier)."""
|
||||
async def build_authorization_url(self, par: PARPolicy = 'never') -> tuple[str, str, str, str]:
|
||||
"""Return (authorization_url, state, nonce, code_verifier).
|
||||
|
||||
When ``par`` selects Pushed Authorization Requests (RFC 9126), the
|
||||
authorization request parameters are POSTed to the provider's PAR
|
||||
endpoint first and the returned URL carries only ``client_id`` and
|
||||
``request_uri``.
|
||||
"""
|
||||
metadata = await self._ensure_metadata()
|
||||
authorization_endpoint = metadata.get('authorization_endpoint')
|
||||
if not isinstance(authorization_endpoint, str):
|
||||
@@ -54,7 +63,6 @@ class OIDCClient:
|
||||
code_verifier, code_challenge = generate_pkce()
|
||||
|
||||
params = {
|
||||
'client_id': self._config.client_id,
|
||||
'response_type': 'code',
|
||||
'scope': ' '.join(self._config.scopes),
|
||||
'redirect_uri': self._config.redirect_uri,
|
||||
@@ -63,9 +71,67 @@ class OIDCClient:
|
||||
'code_challenge': code_challenge,
|
||||
'code_challenge_method': 'S256',
|
||||
}
|
||||
url = authorization_endpoint + '?' + urlencode(params)
|
||||
|
||||
if self._use_par(metadata, par):
|
||||
request_uri = await self._push_authorization_request(metadata, params)
|
||||
redirect_params = {'client_id': self._config.client_id, 'request_uri': request_uri}
|
||||
else:
|
||||
redirect_params = {'client_id': self._config.client_id, **params}
|
||||
|
||||
url = authorization_endpoint + '?' + urlencode(redirect_params)
|
||||
return url, state, nonce, code_verifier
|
||||
|
||||
def _use_par(self, metadata: Mapping[str, Any], policy: PARPolicy) -> bool:
|
||||
endpoint = metadata.get('pushed_authorization_request_endpoint')
|
||||
required = metadata.get('require_pushed_authorization_requests') is True
|
||||
|
||||
if required:
|
||||
if not isinstance(endpoint, str):
|
||||
raise RuntimeError(
|
||||
'OIDC provider requires pushed authorization requests but does not '
|
||||
'advertise a pushed_authorization_request_endpoint'
|
||||
)
|
||||
if policy == 'never':
|
||||
raise RuntimeError(
|
||||
"OIDC provider requires pushed authorization requests; set par to "
|
||||
"'auto' or 'always'"
|
||||
)
|
||||
return True
|
||||
|
||||
if policy == 'always':
|
||||
if not isinstance(endpoint, str):
|
||||
raise RuntimeError(
|
||||
'OIDC provider does not advertise a pushed_authorization_request_endpoint'
|
||||
)
|
||||
return True
|
||||
|
||||
return policy == 'auto' and isinstance(endpoint, str)
|
||||
|
||||
async def _push_authorization_request(
|
||||
self, metadata: Mapping[str, Any], params: Mapping[str, str]
|
||||
) -> str:
|
||||
endpoint = metadata.get('pushed_authorization_request_endpoint')
|
||||
if not isinstance(endpoint, str):
|
||||
raise RuntimeError(
|
||||
'OIDC provider does not advertise a pushed_authorization_request_endpoint'
|
||||
)
|
||||
|
||||
# RFC 9126 section 2.1: client_id is required in the body even when the
|
||||
# client authenticates with the Authorization header.
|
||||
data = {'client_id': self._config.client_id, **params}
|
||||
resp = await self._request('POST', endpoint, data=data, auth=self._client_auth())
|
||||
resp.raise_for_status()
|
||||
body = resp.json()
|
||||
request_uri = body.get('request_uri') if isinstance(body, Mapping) else None
|
||||
if not isinstance(request_uri, str):
|
||||
raise RuntimeError('PAR endpoint response is missing request_uri')
|
||||
return request_uri
|
||||
|
||||
def _client_auth(self) -> Optional[tuple[str, str]]:
|
||||
if self._config.client_secret is None:
|
||||
return None
|
||||
return (self._config.client_id, self._config.client_secret)
|
||||
|
||||
async def fetch_token(self, code: str, code_verifier: str) -> Mapping[str, Any]:
|
||||
metadata = await self._ensure_metadata()
|
||||
token_endpoint = metadata.get('token_endpoint')
|
||||
@@ -79,11 +145,8 @@ class OIDCClient:
|
||||
'client_id': self._config.client_id,
|
||||
'code_verifier': code_verifier,
|
||||
}
|
||||
auth: Optional[tuple[str, str]] = None
|
||||
if self._config.client_secret is not None:
|
||||
auth = (self._config.client_id, self._config.client_secret)
|
||||
|
||||
resp = await self._request('POST', token_endpoint, data=data, auth=auth)
|
||||
resp = await self._request('POST', token_endpoint, data=data, auth=self._client_auth())
|
||||
resp.raise_for_status()
|
||||
return cast(Mapping[str, Any], resp.json())
|
||||
|
||||
@@ -98,15 +161,35 @@ class OIDCClient:
|
||||
'refresh_token': refresh_token,
|
||||
'client_id': self._config.client_id,
|
||||
}
|
||||
auth: Optional[tuple[str, str]] = None
|
||||
if self._config.client_secret is not None:
|
||||
auth = (self._config.client_id, self._config.client_secret)
|
||||
|
||||
resp = await self._request('POST', token_endpoint, data=data, auth=auth)
|
||||
resp = await self._request('POST', token_endpoint, data=data, auth=self._client_auth())
|
||||
resp.raise_for_status()
|
||||
return cast(Mapping[str, Any], resp.json())
|
||||
|
||||
async def validate_id_token(self, id_token: str, nonce: str) -> Mapping[str, Any]:
|
||||
payload = await self._decode_id_token(id_token)
|
||||
if payload.get('nonce') != nonce:
|
||||
raise ValueError('ID token nonce mismatch')
|
||||
return payload
|
||||
|
||||
async def validate_refreshed_id_token(
|
||||
self,
|
||||
id_token: str,
|
||||
expected_nonce: Optional[str] = None,
|
||||
) -> Mapping[str, Any]:
|
||||
"""Validate an ID token returned by the refresh token grant.
|
||||
|
||||
Unlike the authorization code flow, the provider may omit the ``nonce``
|
||||
claim; when it is present and the original nonce is known, the two must
|
||||
match (OIDC Core 12.2).
|
||||
"""
|
||||
payload = await self._decode_id_token(id_token)
|
||||
token_nonce = payload.get('nonce')
|
||||
if token_nonce is not None and expected_nonce is not None and token_nonce != expected_nonce:
|
||||
raise ValueError('ID token nonce mismatch')
|
||||
return payload
|
||||
|
||||
async def _decode_id_token(self, id_token: str) -> Mapping[str, Any]:
|
||||
header = jwt.get_unverified_header(id_token)
|
||||
alg = header.get('alg')
|
||||
if not isinstance(alg, str):
|
||||
@@ -115,29 +198,59 @@ class OIDCClient:
|
||||
if alg == 'none':
|
||||
if not self._config.allow_unsigned_id_tokens:
|
||||
raise ValueError('Unsigned ID tokens are not allowed')
|
||||
metadata_issuer = await self._get_metadata_issuer()
|
||||
payload = jwt.decode(
|
||||
id_token,
|
||||
options={'verify_signature': False},
|
||||
audience=self._config.client_id,
|
||||
options={
|
||||
'verify_signature': False,
|
||||
'verify_aud': True,
|
||||
'verify_exp': True,
|
||||
'verify_iat': True,
|
||||
'verify_nbf': True,
|
||||
},
|
||||
)
|
||||
else:
|
||||
if alg not in self._config.allowed_id_token_algorithms:
|
||||
raise ValueError(f'ID token algorithm {alg} is not allowed')
|
||||
metadata_issuer = await self._get_metadata_issuer()
|
||||
jwks = await self._ensure_jwks()
|
||||
signing_key = self._get_signing_key(jwks, header.get('kid'), alg)
|
||||
payload = jwt.decode(
|
||||
id_token,
|
||||
signing_key.key,
|
||||
algorithms=[alg],
|
||||
issuer=self._config.issuer,
|
||||
audience=self._config.client_id,
|
||||
)
|
||||
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError('ID token payload is not a JSON object')
|
||||
if payload.get('nonce') != nonce:
|
||||
raise ValueError('ID token nonce mismatch')
|
||||
self._validate_issuer(payload, metadata_issuer)
|
||||
return payload
|
||||
|
||||
async def _get_metadata_issuer(self) -> str:
|
||||
metadata = await self._ensure_metadata()
|
||||
issuer = metadata.get('issuer')
|
||||
if not isinstance(issuer, str):
|
||||
raise RuntimeError('OIDC provider does not advertise an issuer')
|
||||
return issuer
|
||||
|
||||
@staticmethod
|
||||
def _validate_issuer(payload: Mapping[str, Any], metadata_issuer: str) -> None:
|
||||
iss = payload.get('iss')
|
||||
if not isinstance(iss, str):
|
||||
raise jwt.MissingRequiredClaimError('iss')
|
||||
|
||||
expected = metadata_issuer
|
||||
if _TENANT_ID_PLACEHOLDER.search(expected):
|
||||
tenant_id = payload.get('tid')
|
||||
if not isinstance(tenant_id, str):
|
||||
raise jwt.InvalidIssuerError('Invalid issuer')
|
||||
expected = _TENANT_ID_PLACEHOLDER.sub(tenant_id, expected, count=1)
|
||||
|
||||
if iss != expected:
|
||||
raise jwt.InvalidIssuerError('Invalid issuer')
|
||||
|
||||
def _get_signing_key(self, jwks: Mapping[str, Any], kid: Optional[str], alg: str) -> PyJWK:
|
||||
keys = jwks.get('keys')
|
||||
if not isinstance(keys, Sequence):
|
||||
|
||||
@@ -1,13 +1,31 @@
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Optional, Sequence
|
||||
from typing import Any, Literal, Optional, Sequence
|
||||
|
||||
import httpx
|
||||
|
||||
_PROVIDER_PATTERN = re.compile(r'^[A-Za-z0-9_-]+$')
|
||||
|
||||
type PARPolicy = Literal['never', 'auto', 'always']
|
||||
|
||||
_PAR_POLICIES: tuple[PARPolicy, ...] = ('never', 'auto', 'always')
|
||||
|
||||
|
||||
@dataclass
|
||||
class OIDCConfig:
|
||||
"""Configuration for an OpenID Connect provider.
|
||||
|
||||
``provider`` is a stable identifier for this provider. It namespaces the
|
||||
session state (``oidc:{provider}:...``) so several providers can be used in
|
||||
the same application without interfering with each other.
|
||||
|
||||
``par`` selects when to use Pushed Authorization Requests (RFC 9126):
|
||||
``'never'`` uses a plain query-string authorization request, ``'auto'``
|
||||
pushes the request whenever the provider advertises a
|
||||
``pushed_authorization_request_endpoint``, and ``'always'`` requires PAR.
|
||||
A provider that sets ``require_pushed_authorization_requests`` is always
|
||||
treated as requiring PAR.
|
||||
|
||||
``client_secret`` is optional; omit it for public clients.
|
||||
``fetch_userinfo`` controls whether the userinfo endpoint is queried after
|
||||
token exchange.
|
||||
@@ -35,8 +53,21 @@ class OIDCConfig:
|
||||
'EdDSA',
|
||||
))
|
||||
http_client: Optional[httpx.AsyncClient] = None
|
||||
provider: str = 'default'
|
||||
par: PARPolicy = 'never'
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not self.provider or _PROVIDER_PATTERN.match(self.provider) is None:
|
||||
raise ValueError(
|
||||
'provider must be a non-empty string of letters, digits, '
|
||||
"underscores or hyphens (got "
|
||||
f'{self.provider!r})'
|
||||
)
|
||||
if self.par not in _PAR_POLICIES:
|
||||
raise ValueError(
|
||||
f'par must be one of {", ".join(map(repr, _PAR_POLICIES))} '
|
||||
f'(got {self.par!r})'
|
||||
)
|
||||
scopes_list = list(self.scopes)
|
||||
if 'openid' not in scopes_list:
|
||||
self.scopes = ('openid', *scopes_list)
|
||||
@@ -1,5 +1,6 @@
|
||||
from typing import Any, Awaitable, Callable, Mapping, Optional, Sequence
|
||||
from typing import Any, Awaitable, Callable, Literal, Mapping, Optional, Sequence, cast
|
||||
|
||||
import jwt
|
||||
from kaya.core import HttpContext, KayaApp, KayaMixin
|
||||
from kaya.session import Session, SessionMixin
|
||||
from urllib.parse import parse_qs
|
||||
@@ -10,6 +11,13 @@ from ._config import OIDCConfig
|
||||
|
||||
type HttpHandler = Callable[..., Awaitable[None]]
|
||||
|
||||
type IdentityStrategy = Literal['single-identity', 'multi-identity']
|
||||
|
||||
_SESSION_NAMESPACE = 'oidc'
|
||||
_USER_SUFFIX = 'user'
|
||||
_IDENTITY_FIELDS = frozenset({'user', 'id_token', 'access_token', 'refresh_token'})
|
||||
_IDENTITY_STRATEGIES: tuple[IdentityStrategy, ...] = ('single-identity', 'multi-identity')
|
||||
|
||||
|
||||
class OIDCUser(Mapping[str, Any]):
|
||||
"""Read-only view of the OIDC userinfo / ID token claims."""
|
||||
@@ -28,7 +36,7 @@ class OIDCUser(Mapping[str, Any]):
|
||||
|
||||
@property
|
||||
def sub(self) -> str:
|
||||
return self._data['sub']
|
||||
return cast(str, self._data['sub'])
|
||||
|
||||
@property
|
||||
def email(self) -> Optional[str]:
|
||||
@@ -56,6 +64,17 @@ class OIDCMixin(KayaMixin):
|
||||
- ``callback_path`` (default ``/auth/callback``): handles the provider callback.
|
||||
- ``logout_path`` (default ``/auth/logout``): logs the user out.
|
||||
|
||||
Session state is namespaced by ``OIDCConfig.provider`` as
|
||||
``oidc:{provider}:...`` so several providers can share one session. Use one
|
||||
``OIDCMixin`` per provider (with distinct route paths), or
|
||||
:class:`~kaya.oidc.OIDCProviders` to wire several at once.
|
||||
|
||||
``strategy`` controls how many identities may coexist in the session. The
|
||||
default ``'single-identity'`` evicts every other provider's identity after a
|
||||
successful login, so the session has exactly one user (last login wins).
|
||||
``'multi-identity'`` keeps each provider logged in independently, which is
|
||||
useful for linked accounts, provider API aggregation, and step-up flows.
|
||||
|
||||
Example::
|
||||
|
||||
session = SessionMixin(InMemorySessionStore())
|
||||
@@ -69,10 +88,33 @@ class OIDCMixin(KayaMixin):
|
||||
...
|
||||
"""
|
||||
|
||||
def __init__(self, config: OIDCConfig, session: SessionMixin) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
config: OIDCConfig,
|
||||
session: SessionMixin,
|
||||
strategy: IdentityStrategy = 'single-identity',
|
||||
) -> None:
|
||||
if strategy not in _IDENTITY_STRATEGIES:
|
||||
raise ValueError(
|
||||
'strategy must be one of '
|
||||
f'{", ".join(map(repr, _IDENTITY_STRATEGIES))} (got {strategy!r})'
|
||||
)
|
||||
self._config = config
|
||||
self._session = session
|
||||
self._client = OIDCClient(config)
|
||||
self._strategy = strategy
|
||||
self._prefix = f'{_SESSION_NAMESPACE}:{config.provider}:'
|
||||
|
||||
@property
|
||||
def provider(self) -> str:
|
||||
return self._config.provider
|
||||
|
||||
@property
|
||||
def strategy(self) -> IdentityStrategy:
|
||||
return self._strategy
|
||||
|
||||
def _key(self, name: str) -> str:
|
||||
return f'{self._prefix}{name}'
|
||||
|
||||
@property
|
||||
def dependencies(self) -> Sequence[KayaMixin]:
|
||||
@@ -81,11 +123,13 @@ class OIDCMixin(KayaMixin):
|
||||
def apply(self, app: KayaApp) -> None:
|
||||
@app.GET(self._config.login_path)
|
||||
async def login(ctx: HttpContext) -> None:
|
||||
auth_url, state, nonce, code_verifier = await self._client.build_authorization_url()
|
||||
auth_url, state, nonce, code_verifier = await self._client.build_authorization_url(
|
||||
par=self._config.par
|
||||
)
|
||||
session = self._session_of(ctx)
|
||||
session['oidc_state'] = state
|
||||
session['oidc_nonce'] = nonce
|
||||
session['oidc_code_verifier'] = code_verifier
|
||||
session[self._key('state')] = state
|
||||
session[self._key('nonce')] = nonce
|
||||
session[self._key('code_verifier')] = code_verifier
|
||||
await ctx.send_empty(302, {'Location': auth_url})
|
||||
|
||||
@app.GET(self._config.callback_path)
|
||||
@@ -108,13 +152,13 @@ class OIDCMixin(KayaMixin):
|
||||
return
|
||||
|
||||
session = self._session_of(ctx)
|
||||
expected_state = session.get('oidc_state')
|
||||
expected_state = session.get(self._key('state'))
|
||||
if state != expected_state:
|
||||
await ctx.send_str(400, 'Invalid state')
|
||||
return
|
||||
|
||||
nonce = session.get('oidc_nonce', '')
|
||||
code_verifier = session.get('oidc_code_verifier', '')
|
||||
nonce = session.get(self._key('nonce'), '')
|
||||
code_verifier = session.get(self._key('code_verifier'), '')
|
||||
|
||||
try:
|
||||
token_response = await self._client.fetch_token(code, code_verifier)
|
||||
@@ -131,31 +175,59 @@ class OIDCMixin(KayaMixin):
|
||||
userinfo = await self._client.fetch_userinfo(access_token)
|
||||
user.update(userinfo)
|
||||
|
||||
session.pop('oidc_state', None)
|
||||
session.pop('oidc_nonce', None)
|
||||
session.pop('oidc_code_verifier', None)
|
||||
session.pop(self._key('state'), None)
|
||||
session.pop(self._key('code_verifier'), None)
|
||||
# Keep this provider's nonce so refreshed ID tokens can be
|
||||
# checked for nonce continuity (OIDC Core 12.2).
|
||||
|
||||
session['oidc_user'] = user
|
||||
session['oidc_id_token'] = id_token
|
||||
session[self._key('user')] = user
|
||||
session[self._key('id_token')] = id_token
|
||||
if isinstance(access_token, str):
|
||||
session['oidc_access_token'] = access_token
|
||||
session[self._key('access_token')] = access_token
|
||||
if 'refresh_token' in token_response:
|
||||
session['oidc_refresh_token'] = token_response['refresh_token']
|
||||
session[self._key('refresh_token')] = token_response['refresh_token']
|
||||
|
||||
if self._strategy == 'single-identity':
|
||||
self._strip_other_identities(session)
|
||||
|
||||
session.regenerate_id()
|
||||
await ctx.send_empty(302, {'Location': self._config.post_login_redirect})
|
||||
except ValueError as exc:
|
||||
except (ValueError, jwt.PyJWTError) as exc:
|
||||
await ctx.send_str(400, f'Authentication failed: {exc}')
|
||||
|
||||
@app.GET(self._config.logout_path)
|
||||
async def logout(ctx: HttpContext) -> None:
|
||||
session = self._session_of(ctx)
|
||||
id_token = session.get('oidc_id_token')
|
||||
session.invalidate()
|
||||
id_token = session.get(self._key('id_token'))
|
||||
self._logout_session(session)
|
||||
logout_url = await self._client.build_logout_url(id_token if isinstance(id_token, str) else None)
|
||||
location = logout_url if logout_url is not None else self._config.post_logout_redirect
|
||||
await ctx.send_empty(302, {'Location': location})
|
||||
|
||||
def _logout_session(self, session: Session) -> None:
|
||||
"""Drop this provider's session state.
|
||||
|
||||
The shared session is only invalidated when no other provider remains
|
||||
logged in, so logging out of one provider does not sign the user out of
|
||||
the others.
|
||||
"""
|
||||
for key in [key for key in session if key.startswith(self._prefix)]:
|
||||
del session[key]
|
||||
if not any(_is_user_key(key) for key in session):
|
||||
session.invalidate()
|
||||
|
||||
def _strip_other_identities(self, session: Session) -> None:
|
||||
"""Remove every other provider's identity and tokens from the session.
|
||||
|
||||
Only identity keys are dropped; ephemeral flow keys (``state``,
|
||||
``nonce``, ``code_verifier``) are left alone so logins already in flight
|
||||
can still complete. Under this strategy the last callback to finish
|
||||
wins.
|
||||
"""
|
||||
for key in [key for key in session
|
||||
if _is_identity_key(key) and not key.startswith(self._prefix)]:
|
||||
del session[key]
|
||||
|
||||
@staticmethod
|
||||
def _session_of(ctx: HttpContext) -> Session:
|
||||
session = ctx.session
|
||||
@@ -169,14 +241,31 @@ class OIDCMixin(KayaMixin):
|
||||
return None
|
||||
|
||||
def is_authenticated(self, ctx: HttpContext) -> bool:
|
||||
return 'oidc_user' in self._session_of(ctx)
|
||||
return self._key('user') in self._session_of(ctx)
|
||||
|
||||
def get_user(self, ctx: HttpContext) -> Optional[OIDCUser]:
|
||||
user = self._session_of(ctx).get('oidc_user')
|
||||
user = self._session_of(ctx).get(self._key('user'))
|
||||
if user is None or not isinstance(user, Mapping):
|
||||
return None
|
||||
return OIDCUser(user)
|
||||
|
||||
def logged_in_providers(self, ctx: HttpContext) -> list[str]:
|
||||
"""Return the names of every provider with a user in this session."""
|
||||
return [_provider_of_user_key(key)
|
||||
for key in self._session_of(ctx) if _is_user_key(key)]
|
||||
|
||||
def get_users(self, ctx: HttpContext) -> dict[str, OIDCUser]:
|
||||
"""Return the users of every provider logged into this session."""
|
||||
session = self._session_of(ctx)
|
||||
users: dict[str, OIDCUser] = {}
|
||||
for key in session:
|
||||
if not _is_user_key(key):
|
||||
continue
|
||||
data = session.get(key)
|
||||
if isinstance(data, Mapping):
|
||||
users[_provider_of_user_key(key)] = OIDCUser(data)
|
||||
return users
|
||||
|
||||
def require_auth(self, handler: HttpHandler) -> HttpHandler:
|
||||
async def wrapper(ctx: HttpContext, *args: Any, **kwargs: Any) -> None:
|
||||
if not self.is_authenticated(ctx):
|
||||
@@ -186,18 +275,44 @@ class OIDCMixin(KayaMixin):
|
||||
return wrapper
|
||||
|
||||
async def refresh_access_token(self, session: Session) -> Optional[str]:
|
||||
refresh_token = session.get('oidc_refresh_token')
|
||||
refresh_token = session.get(self._key('refresh_token'))
|
||||
if not isinstance(refresh_token, str):
|
||||
return None
|
||||
try:
|
||||
token_response = await self._client.refresh_token(refresh_token)
|
||||
access_token = token_response.get('access_token')
|
||||
id_token = token_response.get('id_token')
|
||||
if isinstance(id_token, str):
|
||||
stored_nonce = session.get(self._key('nonce'))
|
||||
claims = await self._client.validate_refreshed_id_token(
|
||||
id_token,
|
||||
stored_nonce if isinstance(stored_nonce, str) else None,
|
||||
)
|
||||
user = session.get(self._key('user'))
|
||||
original_sub = user.get('sub') if isinstance(user, Mapping) else None
|
||||
if isinstance(original_sub, str) and claims.get('sub') != original_sub:
|
||||
raise ValueError('ID token subject mismatch')
|
||||
|
||||
if isinstance(access_token, str):
|
||||
session['oidc_access_token'] = access_token
|
||||
if 'id_token' in token_response:
|
||||
session['oidc_id_token'] = token_response['id_token']
|
||||
session[self._key('access_token')] = access_token
|
||||
if isinstance(id_token, str):
|
||||
session[self._key('id_token')] = id_token
|
||||
if 'refresh_token' in token_response:
|
||||
session['oidc_refresh_token'] = token_response['refresh_token']
|
||||
session[self._key('refresh_token')] = token_response['refresh_token']
|
||||
return access_token if isinstance(access_token, str) else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _is_user_key(key: str) -> bool:
|
||||
return key.startswith(f'{_SESSION_NAMESPACE}:') and key.endswith(f':{_USER_SUFFIX}')
|
||||
|
||||
|
||||
def _is_identity_key(key: str) -> bool:
|
||||
return (key.startswith(f'{_SESSION_NAMESPACE}:')
|
||||
and key.rsplit(':', 1)[-1] in _IDENTITY_FIELDS)
|
||||
|
||||
|
||||
def _provider_of_user_key(key: str) -> str:
|
||||
return key[len(_SESSION_NAMESPACE) + 1:-(len(_USER_SUFFIX) + 1)]
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
from dataclasses import replace
|
||||
from types import MappingProxyType
|
||||
from typing import Callable, Mapping, Optional, Sequence
|
||||
|
||||
from kaya.core import HttpContext, KayaApp, KayaMixin
|
||||
from kaya.session import Session, SessionMixin
|
||||
|
||||
from ._config import OIDCConfig
|
||||
from ._mixin import (
|
||||
_IDENTITY_STRATEGIES,
|
||||
HttpHandler,
|
||||
IdentityStrategy,
|
||||
OIDCMixin,
|
||||
OIDCUser,
|
||||
)
|
||||
|
||||
|
||||
class OIDCProviders(KayaMixin):
|
||||
"""Wire several OpenID Connect providers into one application.
|
||||
|
||||
Each config's ``provider`` becomes a route segment under ``base_path``:
|
||||
``{base_path}/{provider}/login``, ``{base_path}/{provider}/callback`` and
|
||||
``{base_path}/{provider}/logout``. All providers share the supplied
|
||||
:class:`~kaya.session.SessionMixin` and namespace their session state by
|
||||
provider.
|
||||
|
||||
``strategy`` controls how many identities may coexist in the session. The
|
||||
default ``'single-identity'`` evicts every other provider's identity after a
|
||||
successful login (last login wins). ``'multi-identity'`` keeps each provider
|
||||
logged in independently, so logging out of one provider does not sign the
|
||||
user out of the others.
|
||||
|
||||
Example::
|
||||
|
||||
session = SessionMixin(InMemorySessionStore())
|
||||
oidc = OIDCProviders([google_config, github_config], session=session)
|
||||
app = KayaApp(mixins=[session, oidc])
|
||||
|
||||
@app.GET('/profile')
|
||||
@oidc.require_auth('google')
|
||||
async def profile(ctx: HttpContext):
|
||||
user = oidc.get_user(ctx, 'google')
|
||||
...
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
configs: Sequence[OIDCConfig],
|
||||
session: SessionMixin,
|
||||
base_path: str = '/auth',
|
||||
strategy: IdentityStrategy = 'single-identity',
|
||||
) -> None:
|
||||
if not configs:
|
||||
raise ValueError('At least one OIDCConfig is required')
|
||||
if not base_path.startswith('/'):
|
||||
raise ValueError('base_path must start with "/"')
|
||||
if strategy not in _IDENTITY_STRATEGIES:
|
||||
raise ValueError(
|
||||
'strategy must be one of '
|
||||
f'{", ".join(map(repr, _IDENTITY_STRATEGIES))} (got {strategy!r})'
|
||||
)
|
||||
|
||||
self._session = session
|
||||
self._strategy = strategy
|
||||
self._base_path = base_path.rstrip('/')
|
||||
self._providers: dict[str, OIDCMixin] = {}
|
||||
for config in configs:
|
||||
name = config.provider
|
||||
if name in self._providers:
|
||||
raise ValueError(f'Duplicate OIDC provider {name!r}')
|
||||
self._providers[name] = OIDCMixin(
|
||||
replace(
|
||||
config,
|
||||
login_path=f'{self._base_path}/{name}/login',
|
||||
callback_path=f'{self._base_path}/{name}/callback',
|
||||
logout_path=f'{self._base_path}/{name}/logout',
|
||||
),
|
||||
session=session,
|
||||
strategy=strategy,
|
||||
)
|
||||
|
||||
@property
|
||||
def base_path(self) -> str:
|
||||
return self._base_path
|
||||
|
||||
@property
|
||||
def strategy(self) -> IdentityStrategy:
|
||||
return self._strategy
|
||||
|
||||
@property
|
||||
def providers(self) -> Mapping[str, OIDCMixin]:
|
||||
return MappingProxyType(self._providers)
|
||||
|
||||
def provider(self, name: str) -> OIDCMixin:
|
||||
try:
|
||||
return self._providers[name]
|
||||
except KeyError:
|
||||
raise KeyError(f'Unknown OIDC provider {name!r}') from None
|
||||
|
||||
def __getitem__(self, name: str) -> OIDCMixin:
|
||||
return self.provider(name)
|
||||
|
||||
@property
|
||||
def dependencies(self) -> Sequence[KayaMixin]:
|
||||
return [self._session]
|
||||
|
||||
def apply(self, app: KayaApp) -> None:
|
||||
for mixin in self._providers.values():
|
||||
mixin.apply(app)
|
||||
|
||||
def is_authenticated(self, ctx: HttpContext, provider: str) -> bool:
|
||||
return self.provider(provider).is_authenticated(ctx)
|
||||
|
||||
def get_user(self, ctx: HttpContext, provider: str) -> Optional[OIDCUser]:
|
||||
return self.provider(provider).get_user(ctx)
|
||||
|
||||
def get_users(self, ctx: HttpContext) -> dict[str, OIDCUser]:
|
||||
"""Return the user of every provider with a session on ``ctx``."""
|
||||
users: dict[str, OIDCUser] = {}
|
||||
for name, mixin in self._providers.items():
|
||||
user = mixin.get_user(ctx)
|
||||
if user is not None:
|
||||
users[name] = user
|
||||
return users
|
||||
|
||||
def logged_in_providers(self, ctx: HttpContext) -> list[str]:
|
||||
"""Return the names of the providers logged into ``ctx``."""
|
||||
return [name for name, mixin in self._providers.items() if mixin.is_authenticated(ctx)]
|
||||
|
||||
def require_auth(self, provider: str) -> Callable[[HttpHandler], HttpHandler]:
|
||||
return self.provider(provider).require_auth
|
||||
|
||||
async def refresh_access_token(self, session: Session, provider: str) -> Optional[str]:
|
||||
return await self.provider(provider).refresh_access_token(session)
|
||||
File diff suppressed because it is too large.
Load diff
@@ -69,9 +69,13 @@ async def home(ctx: HttpContext):
|
||||
`http.response.status_code` when the handler sends the response; 5xx
|
||||
statuses mark the span as failed. Exceptions escaping the handler are
|
||||
recorded on the span and also mark it as failed.
|
||||
- Every WebSocket connection gets one span for its whole lifetime. The close
|
||||
code is recorded as `kaya.websocket.close_code`; exceptions and configured
|
||||
error close codes mark the span as failed.
|
||||
- Every WebSocket connection gets one span for its whole lifetime, carrying
|
||||
the applicable standard attributes: `url.scheme`, `url.path`, `url.query`,
|
||||
`network.protocol.name` (`websocket`), `network.transport`,
|
||||
`client.address`/`client.port` and `server.address`/`server.port`. The
|
||||
close code is recorded as `kaya.websocket.close_code`; exceptions and
|
||||
configured error close codes mark the span as failed and set `error.type`
|
||||
(the exception class name, or the close code).
|
||||
- W3C `traceparent`/`tracestate` headers on incoming requests are honored,
|
||||
so traces propagate from upstream services.
|
||||
- Metrics:
|
||||
|
||||
@@ -84,6 +84,30 @@ def _server_attributes(ctx: HttpContext) -> Dict[str, Any]:
|
||||
return attributes
|
||||
|
||||
|
||||
def _websocket_attributes(ws: WebSocket) -> Dict[str, Any]:
|
||||
"""Standard attributes applicable to a websocket connection span.
|
||||
|
||||
OpenTelemetry defines no websocket-specific conventions; these are the
|
||||
general URL, network and peer attributes that apply to the connection.
|
||||
"""
|
||||
attributes: Dict[str, Any] = {
|
||||
'url.scheme': ws.scheme,
|
||||
'url.path': ws.path,
|
||||
'network.protocol.name': 'websocket',
|
||||
'network.transport': 'tcp',
|
||||
}
|
||||
if ws.query_string:
|
||||
attributes['url.query'] = ws.query_string
|
||||
if ws.client is not None:
|
||||
attributes['client.address'] = ws.client[0]
|
||||
attributes['client.port'] = ws.client[1]
|
||||
if ws.server is not None:
|
||||
attributes['server.address'] = ws.server[0]
|
||||
if ws.server[1] is not None:
|
||||
attributes['server.port'] = ws.server[1]
|
||||
return attributes
|
||||
|
||||
|
||||
def _normalize_header_name(name: str) -> str:
|
||||
return name.lower().replace('-', '_')
|
||||
|
||||
@@ -201,7 +225,11 @@ class OTelMixin(KayaMixin):
|
||||
usual HTTP semantic attributes; the response status code is recorded when
|
||||
the handler sends the response and 5xx statuses mark the span as failed.
|
||||
Exceptions escaping the handler are recorded and mark the span as failed.
|
||||
WebSocket connections get one span for the whole connection lifetime.
|
||||
WebSocket connections get one span for the whole connection lifetime,
|
||||
carrying the applicable standard websocket URL, network and peer
|
||||
attributes (``url.scheme``, ``url.path``, ``network.protocol.name``,
|
||||
``client.address``/``server.address``, ...); exceptions and configured
|
||||
error close codes set ``error.type`` and mark the span as failed.
|
||||
W3C ``traceparent``/``tracestate`` headers on incoming requests are
|
||||
honored, so traces propagate from upstream services.
|
||||
|
||||
@@ -455,14 +483,7 @@ class OTelMixin(KayaMixin):
|
||||
async def _before_websocket(self, ws: WebSocket) -> Optional[WebSocket]:
|
||||
if self._is_excluded(ws.path):
|
||||
return None
|
||||
holder = self._start_span(
|
||||
f'WS {ws.path}',
|
||||
{
|
||||
'url.scheme': ws.scheme,
|
||||
'url.path': ws.path,
|
||||
},
|
||||
ws.headers,
|
||||
)
|
||||
holder = self._start_span(f'WS {ws.path}', _websocket_attributes(ws), ws.headers)
|
||||
self._spans[id(ws)] = holder
|
||||
self._call_hook(self._websocket_connect_hook, holder.span, ws)
|
||||
return _TracedWebSocket(ws, holder)
|
||||
@@ -475,10 +496,16 @@ class OTelMixin(KayaMixin):
|
||||
error = isinstance(exception, BaseException)
|
||||
if isinstance(exception, BaseException):
|
||||
holder.span.record_exception(exception)
|
||||
holder.span.set_attribute('error.type', type(exception).__qualname__)
|
||||
close_code = holder.close_code
|
||||
if close_code is not None:
|
||||
holder.span.set_attribute('kaya.websocket.close_code', close_code)
|
||||
error = error or close_code in self._websocket_error_close_codes
|
||||
if close_code in self._websocket_error_close_codes:
|
||||
error = True
|
||||
if not isinstance(exception, BaseException):
|
||||
# Mirror the HTTP conventions, where server-side error
|
||||
# statuses are recorded as ``error.type`` values.
|
||||
holder.span.set_attribute('error.type', str(close_code))
|
||||
self._call_hook(self._websocket_close_hook, holder.span, ws)
|
||||
self._finish_span(holder, error)
|
||||
|
||||
|
||||
@@ -279,7 +279,7 @@ class MetricsTest(unittest.TestCase):
|
||||
class WebSocketTracingTest(unittest.TestCase):
|
||||
|
||||
@staticmethod
|
||||
def _make_ws(headers=()):
|
||||
def _make_ws(headers=(), query_string: bytes = b''):
|
||||
async def send(message):
|
||||
pass
|
||||
|
||||
@@ -289,7 +289,7 @@ class WebSocketTracingTest(unittest.TestCase):
|
||||
scope = {
|
||||
'type': 'websocket',
|
||||
'path': '/ws/games/abc',
|
||||
'query_string': b'',
|
||||
'query_string': query_string,
|
||||
'scheme': 'ws',
|
||||
'client': ('127.0.0.1', 12345),
|
||||
'server': ('127.0.0.1', 80),
|
||||
@@ -308,7 +308,29 @@ class WebSocketTracingTest(unittest.TestCase):
|
||||
self.assertEqual(0, len(mixin._spans))
|
||||
spans = span_exporter.get_finished_spans()
|
||||
self.assertEqual(1, len(spans))
|
||||
self.assertEqual('WS /ws/games/abc', spans[0].name)
|
||||
span = spans[0]
|
||||
self.assertEqual('WS /ws/games/abc', span.name)
|
||||
assert span.attributes is not None
|
||||
self.assertEqual('ws', span.attributes['url.scheme'])
|
||||
self.assertEqual('/ws/games/abc', span.attributes['url.path'])
|
||||
self.assertNotIn('url.query', span.attributes)
|
||||
self.assertEqual('websocket', span.attributes['network.protocol.name'])
|
||||
self.assertEqual('tcp', span.attributes['network.transport'])
|
||||
self.assertEqual('127.0.0.1', span.attributes['client.address'])
|
||||
self.assertEqual(12345, span.attributes['client.port'])
|
||||
self.assertEqual('127.0.0.1', span.attributes['server.address'])
|
||||
self.assertEqual(80, span.attributes['server.port'])
|
||||
|
||||
@async_test
|
||||
async def test_websocket_query_string_recorded(self):
|
||||
mixin, span_exporter, _ = make_mixin()
|
||||
ws = self._make_ws(query_string=b'room=1&token=abc')
|
||||
wrapped = await mixin._before_websocket(ws)
|
||||
await mixin._after_websocket(wrapped) # type: ignore[arg-type]
|
||||
spans = span_exporter.get_finished_spans()
|
||||
self.assertEqual(1, len(spans))
|
||||
assert spans[0].attributes is not None
|
||||
self.assertEqual('room=1&token=abc', spans[0].attributes['url.query'])
|
||||
|
||||
@async_test
|
||||
async def test_websocket_traceparent_propagates(self):
|
||||
@@ -344,6 +366,7 @@ class WebSocketTracingTest(unittest.TestCase):
|
||||
self.assertEqual(StatusCode.ERROR, spans[0].status.status_code)
|
||||
assert spans[0].attributes is not None
|
||||
self.assertEqual(1011, spans[0].attributes['kaya.websocket.close_code'])
|
||||
self.assertEqual('1011', spans[0].attributes['error.type'])
|
||||
|
||||
@async_test
|
||||
async def test_websocket_exception_marks_span_error(self):
|
||||
@@ -358,6 +381,8 @@ class WebSocketTracingTest(unittest.TestCase):
|
||||
self.assertEqual(1, len(spans))
|
||||
self.assertEqual(StatusCode.ERROR, spans[0].status.status_code)
|
||||
self.assertTrue(any(event.name == 'exception' for event in spans[0].events))
|
||||
assert spans[0].attributes is not None
|
||||
self.assertEqual('RuntimeError', spans[0].attributes['error.type'])
|
||||
|
||||
@async_test
|
||||
async def test_websocket_hooks_are_called(self):
|
||||
@@ -386,6 +411,7 @@ class WebSocketTracingTest(unittest.TestCase):
|
||||
self.assertEqual(1, len(spans))
|
||||
self.assertEqual(StatusCode.UNSET, spans[0].status.status_code)
|
||||
assert spans[0].attributes is not None
|
||||
self.assertNotIn('error.type', spans[0].attributes)
|
||||
self.assertTrue(spans[0].attributes['test.websocket_connect_hook'])
|
||||
self.assertTrue(spans[0].attributes['test.websocket_close_hook'])
|
||||
|
||||
|
||||
Reference in new issue
Block a user