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.

packages/kaya-oidc/tests/test_oidc.py

@@ -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,26 @@ 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',
+    ) -> 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.discovery = {
             'issuer': self.issuer,
             'authorization_endpoint': f'{self.issuer}/auth',
@@ -48,7 +58,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,
@@ -87,7 +97,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)
@@ -356,5 +366,245 @@ 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) -> 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)
+        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)
+
+    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()