Add OIDC back-channel logout with session indexes

packages/kaya-session-memcache/README.md

@@ -32,11 +32,25 @@ and slides on each access when the session mixin passes a `max_age`. TTLs
 larger than 30 days are automatically converted to absolute Unix timestamps,
 as required by the memcached protocol.
 
+## Session indexes
+
+`MemcacheSessionStore` supports `Session.add_index()`,
+`SessionStore.find_sessions()`, and `delete_by_index()`. Reverse memberships
+are stored as sets and updated with memcached CAS operations. Bucket keys use
+a SHA-256 digest of the index name and value, and stale members are pruned
+during lookups.
+
+Memcached's item-size limit (normally 1 MB) applies to each index bucket. Use
+the Redis store when one index value can have a very large number of active
+sessions.
+
 ## Serialization
 
 Session data is serialized with `pickle` by default, so arbitrary Python
-objects can be stored. A different serializer can be plugged in via the
-`dumps`/`loads` arguments:
+objects can be stored. The serialized value is an envelope containing the
+session dictionary and index metadata; legacy plain-dict values are still
+accepted. A different serializer can be plugged in via the `dumps`/`loads`
+arguments:
 
 ```python
 import json

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

@@ -1,12 +1,15 @@
+import hashlib
 import pickle
 from time import time
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Collection, Optional, Sequence
 
 import aiomcache
 
-from kaya.session import Session, SessionStore
+from kaya.session import IndexKey, Session, SessionStore
 
 _THIRTY_DAYS = 30 * 24 * 60 * 60
+_INDEX_GRACE_PERIOD = 60
+_CAS_RETRIES = 5
 
 
 def _exptime(max_age: int, clock: Callable[[], float]) -> int:
@@ -36,8 +39,14 @@ class MemcacheSessionStore(SessionStore):
     Python objects can be stored. Custom serializers can be plugged in via
     the ``dumps``/``loads`` arguments. Only connect this store to a trusted
     memcached server, as pickle deserialization of untrusted data is unsafe.
+
+    Reverse index buckets are pickled sets updated with memcached CAS
+    operations. Memcached's item-size limit applies to each bucket; Redis is a
+    better fit when one index value can have a very large number of sessions.
     """
 
+    supports_indexes = True
+
     def __init__(
         self,
         client: aiomcache.Client,
@@ -55,6 +64,69 @@ class MemcacheSessionStore(SessionStore):
     def _key(self, session_id: str) -> bytes:
         return f"{self._prefix}{session_id}".encode()
 
+    def _index_key(self, name: str, value: str) -> bytes:
+        digest = hashlib.sha256(f"{name}\0{value}".encode()).hexdigest()
+        return f"{self._prefix}idx:{digest}".encode()
+
+    def _index_exptime(self, max_age: Optional[int]) -> int:
+        if max_age is None:
+            return 0
+        return _exptime(max_age + _INDEX_GRACE_PERIOD, self._clock)
+
+    @staticmethod
+    def _decode_bucket(payload: bytes) -> tuple[set[str], Optional[int]]:
+        members, max_age = pickle.loads(payload)
+        return set(members), max_age
+
+    async def _update_bucket(
+        self,
+        index_key: bytes,
+        *,
+        add: Collection[str] = (),
+        remove: Collection[str] = (),
+        max_age: Optional[int] = None,
+    ) -> None:
+        for _ in range(_CAS_RETRIES):
+            payload, cas_token = await self._client.gets(index_key)
+            if cas_token is None:
+                if not add:
+                    return
+                members = set(add)
+                created = await self._client.add(
+                    index_key,
+                    pickle.dumps((members, max_age)),
+                    exptime=self._index_exptime(max_age),
+                )
+                if created:
+                    return
+                continue
+
+            assert payload is not None
+            members, stored_max_age = self._decode_bucket(payload)
+            updated = (members | set(add)) - set(remove)
+            if updated == members:
+                if not add:
+                    return
+                if await self._client.touch(
+                    index_key,
+                    self._index_exptime(max_age),
+                ):
+                    return
+                # The bucket may have expired between gets() and touch().
+                # Retry so it can be recreated instead of losing the index.
+                continue
+
+            effective_max_age = max_age if add else stored_max_age
+            if await self._client.cas(
+                index_key,
+                pickle.dumps((updated, effective_max_age)),
+                cas_token,
+                exptime=self._index_exptime(effective_max_age),
+            ):
+                return
+
+        raise RuntimeError('Could not update session index after repeated CAS conflicts')
+
     async def load(self, session_id: str, max_age: Optional[int] = None) -> Optional[Session]:
         key = self._key(session_id)
         payload = await self._client.get(key)
@@ -65,12 +137,63 @@ class MemcacheSessionStore(SessionStore):
             await self._client.touch(key, _exptime(max_age, self._clock))
 
         data = self._loads(payload)
-        return Session(session_id, data)
+        return Session.from_payload(session_id, data)
 
     async def save(self, session_id: str, session: Session, max_age: Optional[int] = None) -> None:
-        payload = self._dumps(dict(session))
+        # Index first: if the session write subsequently fails, a stale index
+        # entry is harmless and will be pruned, while a live unindexed session
+        # would be invisible to index-based invalidation.
+        for name, value in session.indexes:
+            await self._update_bucket(
+                self._index_key(name, value),
+                add={session_id},
+                max_age=max_age,
+            )
+        for name, value in session.removed_indexes:
+            await self._update_bucket(
+                self._index_key(name, value),
+                remove={session_id},
+            )
+
+        payload = self._dumps(session.to_payload())
         exptime = _exptime(max_age, self._clock) if max_age is not None else 0
         await self._client.set(self._key(session_id), payload, exptime=exptime)
 
     async def delete(self, session_id: str) -> None:
-        await self._client.delete(self._key(session_id))
+        key = self._key(session_id)
+        payload = await self._client.get(key)
+        indexes: set[IndexKey] = set()
+        if payload is not None:
+            try:
+                indexes = set(Session.from_payload(
+                    session_id,
+                    self._loads(payload),
+                ).indexes)
+            except Exception:
+                # A corrupt payload must not prevent session deletion.
+                pass
+
+        await self._client.delete(key)
+        for name, value in indexes:
+            await self._update_bucket(
+                self._index_key(name, value),
+                remove={session_id},
+            )
+
+    async def find_sessions(self, name: str, value: str) -> Sequence[str]:
+        index_key = self._index_key(name, value)
+        payload = await self._client.get(index_key)
+        if payload is None:
+            return []
+
+        members, _ = self._decode_bucket(payload)
+        live: list[str] = []
+        stale: list[str] = []
+        for session_id in members:
+            if await self._client.get(self._key(session_id)) is None:
+                stale.append(session_id)
+            else:
+                live.append(session_id)
+        if stale:
+            await self._update_bucket(index_key, remove=stale)
+        return live

packages/kaya-session-memcache/tests/test_indexes.py

@@ -0,0 +1,122 @@
+import pickle
+import unittest
+
+from pwo import async_test
+from test_store import FakeClock, StubMemcacheClient
+
+from kaya.session import Session
+from kaya.session.memcache import MemcacheSessionStore
+
+
+class MemcacheSessionIndexTest(unittest.TestCase):
+
+    def setUp(self) -> None:
+        self.clock = FakeClock()
+        self.client = StubMemcacheClient(self.clock)
+        self.store = MemcacheSessionStore(self.client, clock=self.clock)  # type: ignore[arg-type]
+
+    @async_test
+    async def test_save_load_and_find(self) -> None:
+        self.assertTrue(self.store.supports_indexes)
+        await self.store.save('abc', Session('abc', {'user': 'alice'}, indexes={('sub', 'alice')}))
+        loaded = await self.store.load('abc')
+        assert loaded is not None
+        self.assertEqual({'user': 'alice'}, dict(loaded))
+        self.assertEqual(frozenset({('sub', 'alice')}), loaded.indexes)
+        self.assertEqual(['abc'], await self.store.find_sessions('sub', 'alice'))
+
+    @async_test
+    async def test_rotation_preserves_index(self) -> None:
+        await self.store.save('old', Session('old', indexes={('sub', 'alice')}))
+        loaded = await self.store.load('old')
+        assert loaded is not None
+        loaded.set_id('new')
+        await self.store.delete('old')
+        await self.store.save('new', loaded)
+        self.assertEqual(['new'], await self.store.find_sessions('sub', 'alice'))
+
+    @async_test
+    async def test_delete_cleans_reverse_index(self) -> None:
+        await self.store.save('abc', Session('abc', indexes={('sub', 'alice')}))
+        await self.store.delete('abc')
+        self.assertEqual([], await self.store.find_sessions('sub', 'alice'))
+
+    @async_test
+    async def test_delete_by_index_deletes_all_sessions(self) -> None:
+        await self.store.save('a', Session('a', indexes={('sub', 'alice')}))
+        await self.store.save('b', Session('b', indexes={('sub', 'alice')}))
+        await self.store.save('c', Session('c', indexes={('sub', 'bob')}))
+        self.assertEqual(2, await self.store.delete_by_index('sub', 'alice'))
+        self.assertIsNone(await self.store.load('a'))
+        self.assertIsNone(await self.store.load('b'))
+        self.assertIsNotNone(await self.store.load('c'))
+
+    @async_test
+    async def test_stale_member_is_pruned(self) -> None:
+        await self.store.save('abc', Session('abc', indexes={('sub', 'alice')}))
+        await self.client.delete(b'kaya:session:abc')
+        self.assertEqual([], await self.store.find_sessions('sub', 'alice'))
+        index_key = self.store._index_key('sub', 'alice')
+        payload = await self.client.get(index_key)
+        assert payload is not None
+        members, _ = pickle.loads(payload)
+        self.assertEqual(set(), members)
+
+    @async_test
+    async def test_index_ttl_tracks_session_max_age(self) -> None:
+        await self.store.save(
+            'abc',
+            Session('abc', indexes={('sub', 'alice')}),
+            max_age=60,
+        )
+        index_key = self.store._index_key('sub', 'alice')
+        self.assertIn(index_key, self.client._expires)
+        self.clock.advance(61)
+        await self.store.save(
+            'abc',
+            Session('abc', indexes={('sub', 'alice')}),
+            max_age=60,
+        )
+        self.assertIsNotNone(await self.client.get(index_key))
+
+    @async_test
+    async def test_legacy_payload_is_upgraded_on_save(self) -> None:
+        await self.client.set(b'kaya:session:abc', pickle.dumps({'foo': 'bar'}))
+        loaded = await self.store.load('abc')
+        assert loaded is not None
+        self.assertEqual({'foo': 'bar'}, dict(loaded))
+        self.assertEqual(frozenset(), loaded.indexes)
+
+        loaded.add_index('sub', 'alice')
+        await self.store.save('abc', loaded)
+        raw = await self.client.get(b'kaya:session:abc')
+        assert raw is not None
+        payload = pickle.loads(raw)
+        self.assertEqual(1, payload['__kaya_session_v1__'])
+        self.assertEqual({'foo': 'bar'}, payload['data'])
+        self.assertEqual(['abc'], await self.store.find_sessions('sub', 'alice'))
+
+    @async_test
+    async def test_cas_conflict_is_retried(self) -> None:
+        await self.store.save('existing', Session('existing', indexes={('sub', 'alice')}))
+        original_cas = self.client.cas
+        attempts = 0
+
+        async def fail_once(key: bytes, value: bytes, cas_token: int, exptime: int = 0) -> bool:
+            nonlocal attempts
+            attempts += 1
+            if attempts == 1:
+                return False
+            return await original_cas(key, value, cas_token, exptime)
+
+        self.client.cas = fail_once  # type: ignore[method-assign]
+        await self.store.save('abc', Session('abc', indexes={('sub', 'alice')}))
+        self.assertEqual(
+            {'existing', 'abc'},
+            set(await self.store.find_sessions('sub', 'alice')),
+        )
+        self.assertGreaterEqual(attempts, 2)
+
+
+if __name__ == '__main__':
+    unittest.main()

packages/kaya-session-memcache/tests/test_store.py

@@ -28,6 +28,8 @@ class StubMemcacheClient:
         self._clock = clock
         self._data: dict[bytes, bytes] = {}
         self._expires: dict[bytes, float] = {}
+        self._cas: dict[bytes, int] = {}
+        self._next_cas = 1
         self.set_exptimes: list[int] = []
         self.touch_exptimes: list[int] = []
 
@@ -42,19 +44,46 @@ class StubMemcacheClient:
             return None
         return self._data.get(key)
 
-    async def set(self, key: bytes, value: bytes, exptime: int = 0) -> bool:
-        self.set_exptimes.append(exptime)
+    def _store(self, key: bytes, value: bytes, exptime: int) -> None:
         self._data[key] = value
+        self._cas[key] = self._next_cas
+        self._next_cas += 1
         if exptime > 0:
             self._expires[key] = self._clock() + exptime
         else:
             self._expires.pop(key, None)
+
+    async def set(self, key: bytes, value: bytes, exptime: int = 0) -> bool:
+        self.set_exptimes.append(exptime)
+        self._store(key, value, exptime)
+        return True
+
+    async def add(self, key: bytes, value: bytes, exptime: int = 0) -> bool:
+        if self._expired(key):
+            self._data.pop(key, None)
+            self._expires.pop(key, None)
+        if key in self._data:
+            return False
+        self._store(key, value, exptime)
+        return True
+
+    async def gets(self, key: bytes) -> tuple[Optional[bytes], Optional[int]]:
+        value = await self.get(key)
+        if value is None:
+            return None, None
+        return value, self._cas[key]
+
+    async def cas(self, key: bytes, value: bytes, cas_token: int, exptime: int = 0) -> bool:
+        if self._expired(key) or self._cas.get(key) != cas_token:
+            return False
+        self._store(key, value, exptime)
         return True
 
     async def delete(self, key: bytes) -> bool:
         existed = key in self._data
         self._data.pop(key, None)
         self._expires.pop(key, None)
+        self._cas.pop(key, None)
         return existed
 
     async def touch(self, key: bytes, exptime: int) -> bool:
@@ -169,7 +198,11 @@ class MemcacheSessionStoreTest(unittest.TestCase):
             clock=self.clock,
         )
         await store.save('abc', Session('abc', {'foo': 'bar'}))
-        self.assertEqual(b'{"foo": "bar"}', self.client._data[b'kaya:session:abc'])
+        self.assertEqual({
+            '__kaya_session_v1__': 1,
+            'data': {'foo': 'bar'},
+            'indexes': [],
+        }, json.loads(self.client._data[b'kaya:session:abc']))
         loaded = await store.load('abc')
         self.assertIsNotNone(loaded)
         assert loaded is not None