Add OIDC back-channel logout with session indexes

packages/kaya-session/README.md

@@ -66,6 +66,37 @@ session and creates a fresh empty one.
 Set `max_age=None` to disable server-side expiry (and the `Max-Age` cookie
 attribute) entirely.
 
+## Session indexes
+
+Authentication layers can attach reverse-index metadata to a session without
+adding it to the user-visible session dictionary:
+
+```python
+ctx.session.add_index('oidc:default:sub', 'user-123')
+```
+
+Index-capable stores persist these memberships on `save()` and maintain them
+across session ID rotation and deletion:
+
+```python
+session_ids = await store.find_sessions('oidc:default:sub', 'user-123')
+deleted = await store.delete_by_index('oidc:default:sub', 'user-123')
+```
+
+Use `session.remove_index(name, value)` to remove a membership while keeping
+the session. `SessionMixin.store` exposes the backing store for features such
+as back-channel logout.
+
+`InMemorySessionStore`, `RedisSessionStore`, and `MemcacheSessionStore` set
+`supports_indexes = True`. Custom stores can opt in by setting that flag,
+implementing `find_sessions()`, maintaining reverse memberships on `save()`,
+and removing them on `delete()`. `delete_by_index()` has a generic default
+implementation based on `find_sessions()` and `delete()`.
+
+Remote stores serialize an envelope containing both the session dictionary and
+its index metadata. Existing plain-dict payloads remain readable and are
+upgraded on the next save.
+
 ## Features
 
 - `Session`: dict-like session object with modification tracking
@@ -74,6 +105,7 @@ attribute) entirely.
 - `SessionMixin`: composable Kaya mixin managing session cookies and persistence
 - Session ID regeneration (`session.regenerate_id()`) and invalidation
   (`session.invalidate()`) for authentication layers
+- Reverse session indexes for looking up and invalidating sessions by metadata
 - WebSocket support: the session is exposed as `ws.session` in websocket
   handlers, loaded at connect time and persisted on close
 

packages/kaya-session/src/kaya/session/__init__.py

@@ -3,12 +3,13 @@ from pkgutil import extend_path
 __path__ = extend_path(__path__, __name__)
 
 from ._mixin import SessionMixin
-from ._session import Session
+from ._session import IndexKey, Session
 from ._store import InMemorySessionStore, SessionStore
 
 
 __all__ = [
     'InMemorySessionStore',
+    'IndexKey',
     'Session',
     'SessionMixin',
     'SessionStore',

packages/kaya-session/src/kaya/session/_mixin.py

@@ -183,6 +183,16 @@ class SessionMixin(KayaMixin):
         self._secure = secure
         self._samesite = samesite
 
+    @property
+    def store(self) -> SessionStore:
+        """Return the backing session store."""
+        return self._store
+
+    @property
+    def max_age(self) -> Optional[int]:
+        """Return the idle timeout used when persisting sessions."""
+        return self._max_age
+
     def apply(self, app: KayaApp) -> None:
         app.add_before_request_hook(self._before_request)
         app.add_after_request_hook(self._after_request)
@@ -273,3 +283,4 @@ class SessionMixin(KayaMixin):
 
         if session.id is not None:
             await self._store.save(session.id, session, self._max_age)
+            session.commit_indexes()

packages/kaya-session/src/kaya/session/_session.py

@@ -1,4 +1,10 @@
-from typing import Any, Iterator, Mapping, MutableMapping, Optional
+from typing import Any, Iterable, Iterator, Mapping, MutableMapping, Optional
+
+
+IndexKey = tuple[str, str]
+
+_PAYLOAD_MARKER = '__kaya_session_v1__'
+_PAYLOAD_VERSION = 1
 
 
 class Session(MutableMapping[str, Any]):
@@ -9,13 +15,21 @@ class Session(MutableMapping[str, Any]):
     it.
     """
 
-    def __init__(self, session_id: Optional[str] = None, data: Optional[Mapping[str, Any]] = None) -> None:
+    def __init__(
+        self,
+        session_id: Optional[str] = None,
+        data: Optional[Mapping[str, Any]] = None,
+        indexes: Iterable[IndexKey] = (),
+    ) -> None:
         self._id: Optional[str] = session_id
         self._old_id: Optional[str] = None
         self._data: dict[str, Any] = dict(data) if data else {}
         self._modified: bool = False
         self._regenerate: bool = False
         self._invalidated: bool = False
+        self._persisted_indexes: set[IndexKey] = set(indexes)
+        self._added_indexes: set[IndexKey] = set()
+        self._removed_indexes: set[IndexKey] = set()
 
     @property
     def id(self) -> Optional[str]:
@@ -60,6 +74,75 @@ class Session(MutableMapping[str, Any]):
     def mark_modified(self) -> None:
         self._modified = True
 
+    def add_index(self, name: str, value: str) -> None:
+        """Index this session under ``(name, value)`` when it is saved.
+
+        Index-capable stores can use this metadata to find and invalidate
+        sessions without knowing their IDs. Adding an already persisted index
+        is harmless and does not create a duplicate membership.
+        """
+        index = (name, value)
+        self._removed_indexes.discard(index)
+        if index not in self._persisted_indexes:
+            self._added_indexes.add(index)
+        self._modified = True
+
+    def remove_index(self, name: str, value: str) -> None:
+        """Remove an index membership when this session is saved."""
+        index = (name, value)
+        if index in self._added_indexes:
+            self._added_indexes.discard(index)
+        elif index in self._persisted_indexes:
+            self._removed_indexes.add(index)
+        self._modified = True
+
+    @property
+    def indexes(self) -> frozenset[IndexKey]:
+        """Return the index memberships that should be persisted."""
+        return frozenset(
+            (self._persisted_indexes - self._removed_indexes) | self._added_indexes
+        )
+
+    @property
+    def added_indexes(self) -> frozenset[IndexKey]:
+        """Return index memberships added since the last successful save."""
+        return frozenset(self._added_indexes)
+
+    @property
+    def removed_indexes(self) -> frozenset[IndexKey]:
+        """Return index memberships removed since the last successful save."""
+        return frozenset(self._removed_indexes)
+
+    def commit_indexes(self) -> None:
+        """Mark the current index memberships as persisted."""
+        self._persisted_indexes = set(self.indexes)
+        self._added_indexes.clear()
+        self._removed_indexes.clear()
+
+    def to_payload(self) -> dict[str, Any]:
+        """Return the serializable session payload used by remote stores."""
+        return {
+            _PAYLOAD_MARKER: _PAYLOAD_VERSION,
+            'data': dict(self),
+            'indexes': sorted(self.indexes),
+        }
+
+    @classmethod
+    def from_payload(cls, session_id: str, payload: Mapping[str, Any]) -> 'Session':
+        """Build a session from an envelope or a legacy plain-data payload."""
+        data = payload.get('data')
+        raw_indexes = payload.get('indexes')
+        if (payload.get(_PAYLOAD_MARKER) == _PAYLOAD_VERSION
+                and isinstance(data, Mapping)
+                and isinstance(raw_indexes, (list, tuple))):
+            indexes: set[IndexKey] = set()
+            for item in raw_indexes:
+                if (isinstance(item, (list, tuple)) and len(item) == 2
+                        and isinstance(item[0], str) and isinstance(item[1], str)):
+                    indexes.add((item[0], item[1]))
+            return cls(session_id, data, indexes=indexes)
+        return cls(session_id, payload)
+
     def __getitem__(self, key: str) -> Any:
         return self._data[key]
 

packages/kaya-session/src/kaya/session/_store.py

@@ -1,13 +1,21 @@
 from abc import ABC, abstractmethod
 from secrets import token_urlsafe
 from time import monotonic
-from typing import Any, Callable, Optional
+from typing import Any, Callable, ClassVar, Optional, Sequence
 
-from ._session import Session
+from ._session import IndexKey, Session
 
 
 class SessionStore(ABC):
-    """Pluggable backend for session persistence."""
+    """Pluggable backend for session persistence.
+
+    Stores with ``supports_indexes = True`` also maintain reverse indexes for
+    :class:`~kaya.session.Session` index metadata. They must add the session ID
+    to every effective index on ``save()``, apply pending removals, and remove
+    all of a session's index memberships on ``delete()``.
+    """
+
+    supports_indexes: ClassVar[bool] = False
 
     @abstractmethod
     async def load(self, session_id: str, max_age: Optional[int] = None) -> Optional[Session]:
@@ -33,6 +41,23 @@ class SessionStore(ABC):
         """Remove the session from the store."""
         pass
 
+    async def find_sessions(self, name: str, value: str) -> Sequence[str]:
+        """Return live session IDs indexed under ``(name, value)``.
+
+        Only available on stores whose ``supports_indexes`` flag is true.
+        Implementations should prune stale index memberships encountered here.
+        """
+        raise NotImplementedError(
+            f'{type(self).__name__} does not support session indexes'
+        )
+
+    async def delete_by_index(self, name: str, value: str) -> int:
+        """Delete every live session indexed under ``(name, value)``."""
+        session_ids = await self.find_sessions(name, value)
+        for session_id in session_ids:
+            await self.delete(session_id)
+        return len(session_ids)
+
     def new_session_id(self) -> str:
         """Return a new opaque session identifier.
 
@@ -51,10 +76,26 @@ class InMemorySessionStore(SessionStore):
     inactivity. Active sessions slide the expiry window on each access.
     """
 
+    supports_indexes = True
+
     def __init__(self, clock: Callable[[], float] = monotonic) -> None:
         self._clock = clock
         self._data: dict[str, dict[str, Any]] = {}
         self._expires: dict[str, float] = {}
+        self._session_indexes: dict[str, set[IndexKey]] = {}
+        self._index_sessions: dict[IndexKey, set[str]] = {}
+
+    def _drop_session(self, session_id: str) -> None:
+        self._data.pop(session_id, None)
+        self._expires.pop(session_id, None)
+        indexes = self._session_indexes.pop(session_id, set())
+        for index in indexes:
+            session_ids = self._index_sessions.get(index)
+            if session_ids is None:
+                continue
+            session_ids.discard(session_id)
+            if not session_ids:
+                self._index_sessions.pop(index, None)
 
     async def load(self, session_id: str, max_age: Optional[int] = None) -> Optional[Session]:
         data = self._data.get(session_id)
@@ -64,14 +105,17 @@ class InMemorySessionStore(SessionStore):
         expires = self._expires.get(session_id)
         now = self._clock()
         if expires is not None and now > expires:
-            self._data.pop(session_id, None)
-            self._expires.pop(session_id, None)
+            self._drop_session(session_id)
             return None
 
         if max_age is not None and expires is not None:
             self._expires[session_id] = now + max_age
 
-        return Session(session_id, data)
+        return Session(
+            session_id,
+            data,
+            indexes=self._session_indexes.get(session_id, ()),
+        )
 
     async def save(self, session_id: str, session: Session, max_age: Optional[int] = None) -> None:
         self._data[session_id] = dict(session)
@@ -80,6 +124,34 @@ class InMemorySessionStore(SessionStore):
         else:
             self._expires.pop(session_id, None)
 
+        for index in session.removed_indexes:
+            session_ids = self._index_sessions.get(index)
+            if session_ids is None:
+                continue
+            session_ids.discard(session_id)
+            if not session_ids:
+                self._index_sessions.pop(index, None)
+
+        indexes = set(session.indexes)
+        self._session_indexes[session_id] = indexes
+        for index in indexes:
+            self._index_sessions.setdefault(index, set()).add(session_id)
+
     async def delete(self, session_id: str) -> None:
-        self._data.pop(session_id, None)
-        self._expires.pop(session_id, None)
+        self._drop_session(session_id)
+
+    async def find_sessions(self, name: str, value: str) -> Sequence[str]:
+        session_ids = self._index_sessions.get((name, value))
+        if session_ids is None:
+            return []
+
+        now = self._clock()
+        live: list[str] = []
+        for session_id in list(session_ids):
+            expires = self._expires.get(session_id)
+            if (session_id not in self._data
+                    or (expires is not None and now > expires)):
+                self._drop_session(session_id)
+            else:
+                live.append(session_id)
+        return live

packages/kaya-session/tests/test_indexes.py

@@ -0,0 +1,169 @@
+import unittest
+
+import httpx
+from pwo import async_test
+
+from kaya.core import HttpContext, KayaApp
+from kaya.session import InMemorySessionStore, Session, SessionMixin
+
+
+class FakeClock:
+    def __init__(self, start: float = 0.0) -> None:
+        self._now = start
+
+    def __call__(self) -> float:
+        return self._now
+
+    def advance(self, seconds: float) -> None:
+        self._now += seconds
+
+
+class SessionIndexUnitTest(unittest.TestCase):
+
+    def test_add_and_remove_index(self) -> None:
+        session = Session()
+        session.add_index('sub', 'alice')
+        self.assertEqual(frozenset({('sub', 'alice')}), session.indexes)
+        self.assertEqual(frozenset({('sub', 'alice')}), session.added_indexes)
+        self.assertTrue(session.modified)
+
+        session.remove_index('sub', 'alice')
+        self.assertEqual(frozenset(), session.indexes)
+        self.assertEqual(frozenset(), session.added_indexes)
+        self.assertEqual(frozenset(), session.removed_indexes)
+
+    def test_remove_persisted_index(self) -> None:
+        session = Session(indexes={('sub', 'alice')})
+        self.assertFalse(session.modified)
+        session.remove_index('sub', 'alice')
+        self.assertEqual(frozenset(), session.indexes)
+        self.assertEqual(frozenset({('sub', 'alice')}), session.removed_indexes)
+
+        session.add_index('sub', 'alice')
+        self.assertEqual(frozenset({('sub', 'alice')}), session.indexes)
+        self.assertEqual(frozenset(), session.removed_indexes)
+
+    def test_payload_round_trip(self) -> None:
+        session = Session('abc', {'user': 'alice'}, indexes={('sub', 'alice')})
+        loaded = Session.from_payload('abc', session.to_payload())
+        self.assertEqual({'user': 'alice'}, dict(loaded))
+        self.assertEqual(frozenset({('sub', 'alice')}), loaded.indexes)
+
+    def test_legacy_payload(self) -> None:
+        loaded = Session.from_payload('abc', {'user': 'alice'})
+        self.assertEqual({'user': 'alice'}, dict(loaded))
+        self.assertEqual(frozenset(), loaded.indexes)
+
+
+class InMemorySessionIndexTest(unittest.TestCase):
+
+    @async_test
+    async def test_save_load_and_find(self) -> None:
+        store = InMemorySessionStore()
+        self.assertTrue(store.supports_indexes)
+        session = Session('abc', {'user': 'alice'})
+        session.add_index('sub', 'alice')
+        await store.save('abc', session)
+        session.commit_indexes()
+
+        loaded = await store.load('abc')
+        assert loaded is not None
+        self.assertEqual(frozenset({('sub', 'alice')}), loaded.indexes)
+        self.assertEqual(['abc'], await store.find_sessions('sub', 'alice'))
+        self.assertEqual([], await store.find_sessions('sub', 'bob'))
+
+    @async_test
+    async def test_delete_cleans_reverse_index(self) -> None:
+        store = InMemorySessionStore()
+        session = Session('abc', indexes={('sub', 'alice')})
+        await store.save('abc', session)
+        await store.delete('abc')
+        self.assertEqual([], await store.find_sessions('sub', 'alice'))
+        self.assertEqual({}, store._session_indexes)
+
+    @async_test
+    async def test_delete_by_index_deletes_all_sessions(self) -> None:
+        store = InMemorySessionStore()
+        await store.save('a', Session('a', indexes={('sub', 'alice')}))
+        await store.save('b', Session('b', indexes={('sub', 'alice')}))
+        await store.save('c', Session('c', indexes={('sub', 'bob')}))
+
+        self.assertEqual(2, await store.delete_by_index('sub', 'alice'))
+        self.assertIsNone(await store.load('a'))
+        self.assertIsNone(await store.load('b'))
+        self.assertIsNotNone(await store.load('c'))
+
+    @async_test
+    async def test_expiry_prunes_index(self) -> None:
+        clock = FakeClock()
+        store = InMemorySessionStore(clock=clock)
+        await store.save(
+            'abc',
+            Session('abc', indexes={('sub', 'alice')}),
+            max_age=60,
+        )
+        clock.advance(61)
+        self.assertEqual([], await store.find_sessions('sub', 'alice'))
+        self.assertEqual({}, store._session_indexes)
+
+    @async_test
+    async def test_removed_index_is_not_readded(self) -> None:
+        store = InMemorySessionStore()
+        session = Session('abc', indexes={('sub', 'alice')})
+        await store.save('abc', session)
+        session.commit_indexes()
+
+        session.remove_index('sub', 'alice')
+        await store.save('abc', session)
+        self.assertEqual([], await store.find_sessions('sub', 'alice'))
+        self.assertEqual(frozenset(), session.indexes)
+
+
+class SessionIndexIntegrationTest(unittest.TestCase):
+
+    def setUp(self) -> None:
+        self.store = InMemorySessionStore()
+        self.mixin = SessionMixin(self.store)
+        self.app = KayaApp(mixins=[self.mixin])
+
+        @self.app.GET('/login')
+        async def login(ctx: HttpContext) -> None:
+            ctx.session['user'] = 'alice'
+            ctx.session.add_index('sub', 'alice')
+            await ctx.send_str(200, 'ok')
+
+        @self.app.GET('/rotate')
+        async def rotate(ctx: HttpContext) -> None:
+            ctx.session.regenerate_id()
+            await ctx.send_str(200, 'ok')
+
+        @self.app.GET('/logout')
+        async def logout(ctx: HttpContext) -> None:
+            ctx.session.invalidate()
+            await ctx.send_str(200, 'ok')
+
+    @async_test
+    async def test_rotation_preserves_index(self) -> None:
+        self.assertIs(self.store, self.mixin.store)
+        transport = httpx.ASGITransport(app=self.app)
+        async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
+            await client.get('/login')
+            old_id = client.cookies['session_id']
+            self.assertEqual([old_id], await self.store.find_sessions('sub', 'alice'))
+
+            await client.get('/rotate')
+            new_id = client.cookies['session_id']
+            self.assertNotEqual(old_id, new_id)
+            self.assertEqual([new_id], await self.store.find_sessions('sub', 'alice'))
+
+    @async_test
+    async def test_invalidation_cleans_index(self) -> None:
+        transport = httpx.ASGITransport(app=self.app)
+        async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client:
+            await client.get('/login')
+            await client.get('/logout')
+            self.assertEqual([], await self.store.find_sessions('sub', 'alice'))
+
+
+if __name__ == '__main__':
+    unittest.main()