diff --git a/.github/pr-assets/6166-full-project-id-cache-keys.png b/.github/pr-assets/6166-full-project-id-cache-keys.png new file mode 100644 index 0000000000..d20f47d202 Binary files /dev/null and b/.github/pr-assets/6166-full-project-id-cache-keys.png differ diff --git a/api/oss/src/utils/caching.py b/api/oss/src/utils/caching.py index b2c6f97b62..c8e5f3d994 100644 --- a/api/oss/src/utils/caching.py +++ b/api/oss/src/utils/caching.py @@ -40,24 +40,46 @@ # HELPERS ---------------------------------------------------------------------- +AGENTA_CACHE_SCOPE_WIDTH = 12 # minimum width of a scope segment + + +def _scope( + value: Optional[str], + legacy_truncated: Optional[bool] = False, +) -> str: + """One scope segment (`p:` or `u:`) of a cache key. + + Scope segments used to be cut down to their last 12 characters. Two projects whose + ids share that suffix then shared every cache entry in every namespace — including + `check_permissions` and `check_action_access`, which decide authorization. Ids are + server-generated UUID4s so a caller cannot steer a collision, but one project reading + another's cached permission result is not a thing to leave to chance. The id is + carried whole now. + + The padding stays: an absent or short id still produces the same fixed-width segment + it always did, so key shape is unchanged everywhere the id was not being cut. + + `legacy_truncated` reproduces the pre-change shape and exists only so lock holders + can span a rolling deploy — see `oss.src.utils.locking`. + """ + value = value or "" + + if legacy_truncated and len(value) > AGENTA_CACHE_SCOPE_WIDTH: + value = value[-AGENTA_CACHE_SCOPE_WIDTH:] + + return value + "-" * (AGENTA_CACHE_SCOPE_WIDTH - len(value)) + + def _pack( namespace: Optional[str] = None, key: Optional[Union[str, dict]] = None, project_id: Optional[str] = None, user_id: Optional[str] = None, pattern: Optional[bool] = False, + legacy_truncated_scope: Optional[bool] = False, ) -> str: - if project_id: - project_id = project_id[-12:] if len(project_id) > 12 else project_id - else: - project_id = "" - project_id = project_id + "-" * (12 - len(project_id)) - - if user_id: - user_id = user_id[-12:] if len(user_id) > 12 else user_id - else: - user_id = "" - user_id = user_id + "-" * (12 - len(user_id)) + project_id = _scope(project_id, legacy_truncated_scope) + user_id = _scope(user_id, legacy_truncated_scope) namespace = namespace or ("" if not pattern else "*") diff --git a/api/oss/src/utils/locking.py b/api/oss/src/utils/locking.py index cf0f5c1460..fddb67fdf7 100644 --- a/api/oss/src/utils/locking.py +++ b/api/oss/src/utils/locking.py @@ -38,6 +38,73 @@ """ +# TRANSITIONAL: SPANNING THE FULL-PROJECT-ID DEPLOY ----------------------------- +# +# Scope segments in cache and lock keys used to carry only the last 12 characters of +# an id (see `caching._scope`). During a rolling deploy, pods still on the previous +# release take the truncated key, so a lock held only under the full-id key would not +# exclude them and mutual exclusion would be lost for the length of the deploy. +# Every lock operation therefore covers both keys for one release. +# +# REMOVE once no pod predating the full-id change is running: drop the second element +# of `_lock_keys` and the `legacy_key` branches below. That is the whole surface. + + +def _lock_keys( + namespace: str, + key: Optional[Union[str, dict]] = None, + project_id: Optional[str] = None, + user_id: Optional[str] = None, +) -> tuple[str, Optional[str]]: + """This release's lock key, and the previous release's key when it differs.""" + lock_key = _pack( + namespace=f"lock:{namespace}", + key=key, + project_id=project_id, + user_id=user_id, + ) + legacy_key = _pack( + namespace=f"lock:{namespace}", + key=key, + project_id=project_id, + user_id=user_id, + legacy_truncated_scope=True, + ) + + # Ids no longer than the segment width were never truncated, so the two shapes + # coincide and there is no second key to cover. + return lock_key, (legacy_key if legacy_key != lock_key else None) + + +async def _renew_if_owner(lock_key: str, owner: Optional[str], ttl: int) -> bool: + if owner: + return bool( + await _lock_engine.eval( + _LOCK_RENEW_IF_OWNER_SCRIPT, + 1, + lock_key, + owner, + str(ttl), + ) + ) + + return bool(await _lock_engine.expire(lock_key, ttl)) + + +async def _release_if_owner(lock_key: str, owner: Optional[str]) -> bool: + if owner: + return bool( + await _lock_engine.eval( + _LOCK_RELEASE_IF_OWNER_SCRIPT, + 1, + lock_key, + owner, + ) + ) + + return bool(await _lock_engine.delete(lock_key)) + + # LOCK-STORE PRIMITIVES -------------------------------------------------------- # # Thin pass-throughs to the lock Redis client for callers that manage their own @@ -117,14 +184,26 @@ async def acquire_lock( ) """ try: - lock_key = _pack( - namespace=f"lock:{namespace}", + lock_key, legacy_key = _lock_keys( + namespace=namespace, key=key, project_id=project_id, user_id=user_id, ) lock_owner = uuid4().hex + # The legacy key is claimed first: a pod on the previous release sets only that + # one, so taking it is what makes the two generations exclude each other. Claiming + # it second would let both generations hold their own key and enter together. + if legacy_key is not None: + if not await _lock_engine.set(legacy_key, lock_owner, nx=True, ex=ttl): + if LOCK_DEBUG: + log.debug( + "[lock] BLOCKED", + key=legacy_key, + ) + return None + # Atomic SET NX: Returns True if lock acquired, False if already held acquired = await _lock_engine.set(lock_key, lock_owner, nx=True, ex=ttl) @@ -137,6 +216,11 @@ async def acquire_lock( ) return lock_owner else: + # This caller is not entering the critical section, so it must not leave the + # legacy key held until its TTL — that would block everyone for `ttl`. + if legacy_key is not None: + await _release_if_owner(legacy_key, lock_owner) + if LOCK_DEBUG: log.debug( "[lock] BLOCKED", @@ -180,23 +264,19 @@ async def renew_lock( True if lock was renewed, False if lock has already expired or on error """ try: - lock_key = _pack( - namespace=f"lock:{namespace}", + lock_key, legacy_key = _lock_keys( + namespace=namespace, key=key, project_id=project_id, user_id=user_id, ) - if owner: - renewed = await _lock_engine.eval( - _LOCK_RENEW_IF_OWNER_SCRIPT, - 1, - lock_key, - owner, - str(ttl), - ) - else: - renewed = await _lock_engine.expire(lock_key, ttl) + renewed = await _renew_if_owner(lock_key, owner, ttl) + + # Held for as long as the lock itself, or a pod on the previous release would + # take it the moment it lapsed while this holder was still inside the section. + if legacy_key is not None: + await _renew_if_owner(legacy_key, owner, ttl) if renewed: if LOCK_DEBUG: @@ -250,22 +330,19 @@ async def release_lock( await release_lock(namespace="account-creation", key=email) """ try: - lock_key = _pack( - namespace=f"lock:{namespace}", + lock_key, legacy_key = _lock_keys( + namespace=namespace, key=key, project_id=project_id, user_id=user_id, ) - if owner: - deleted = await _lock_engine.eval( - _LOCK_RELEASE_IF_OWNER_SCRIPT, - 1, - lock_key, - owner, - ) - else: - deleted = await _lock_engine.delete(lock_key) + deleted = await _release_if_owner(lock_key, owner) + + # Released even when the primary was already gone: the two were taken together, + # so leaving this one behind would block the section for the rest of its TTL. + if legacy_key is not None: + await _release_if_owner(legacy_key, owner) if deleted: if LOCK_DEBUG: diff --git a/api/oss/tests/pytest/unit/utils/test_cache_key_tenancy.py b/api/oss/tests/pytest/unit/utils/test_cache_key_tenancy.py new file mode 100644 index 0000000000..37e7f8a96f --- /dev/null +++ b/api/oss/tests/pytest/unit/utils/test_cache_key_tenancy.py @@ -0,0 +1,285 @@ +"""Cache and lock keys must not merge two tenants that share an id suffix. + +Scope segments used to carry only the last 12 characters of an id, so two projects +whose ids ended the same way shared every cache entry in every namespace — including +`check_permissions` and `check_action_access`, which decide authorization (#6166). + +The lock tests use a real in-memory fakeredis instance so they run without an external +Redis process; they skip when `fakeredis` is not installed. +""" + +from unittest.mock import patch +from uuid import uuid4 + +import pytest +import pytest_asyncio + +from oss.src.utils.caching import _pack, _scope, AGENTA_CACHE_SCOPE_WIDTH +from oss.src.utils import locking + + +# Two distinct UUID4s contrived to end in the same 12 characters: the collision the old +# truncation produced. Server-generated ids make this astronomically unlikely rather than +# impossible, and the consequence is a cross-tenant read. +COLLIDING_SUFFIX = "b2c3d4e5f6a7" +PROJECT_A = f"11111111-1111-4111-8111-1111{COLLIDING_SUFFIX}" +PROJECT_B = f"22222222-2222-4222-8222-2222{COLLIDING_SUFFIX}" + +USER_A = f"33333333-3333-4333-8333-3333{COLLIDING_SUFFIX}" +USER_B = f"44444444-4444-4444-8444-4444{COLLIDING_SUFFIX}" + + +def _key(namespace="check_action_access", project_id=None, user_id=None, **kwargs): + return _pack( + namespace=namespace, + key="read", + project_id=project_id, + user_id=user_id, + **kwargs, + ) + + +# KEY SHAPE -------------------------------------------------------------------- + + +def test_projects_sharing_an_id_suffix_do_not_share_a_cache_key(): + assert _key(project_id=PROJECT_A) != _key(project_id=PROJECT_B) + + +def test_users_sharing_an_id_suffix_do_not_share_a_cache_key(): + assert _key(project_id=PROJECT_A, user_id=USER_A) != _key( + project_id=PROJECT_A, user_id=USER_B + ) + + +def test_the_whole_id_is_in_the_key(): + assert PROJECT_A in _key(project_id=PROJECT_A) + assert USER_A in _key(project_id=PROJECT_A, user_id=USER_A) + + +def test_an_invalidation_pattern_is_scoped_to_one_project(): + """`invalidate_cache` without a key scans a pattern; it must not span both projects.""" + pattern = _pack(project_id=PROJECT_A, pattern=True) + + assert PROJECT_A in pattern + assert PROJECT_B not in pattern + + +def test_a_short_or_absent_id_keeps_its_historical_segment(): + """Only ids that were being cut change shape; everything else is untouched.""" + assert _scope(None) == "-" * AGENTA_CACHE_SCOPE_WIDTH + assert _scope("") == "-" * AGENTA_CACHE_SCOPE_WIDTH + assert _scope("abc") == "abc" + "-" * (AGENTA_CACHE_SCOPE_WIDTH - 3) + + exactly_wide = "a" * AGENTA_CACHE_SCOPE_WIDTH + assert _scope(exactly_wide) == exactly_wide + # At or below the width, truncating was already a no-op, so both shapes agree. + assert _scope(exactly_wide, legacy_truncated=True) == _scope(exactly_wide) + + +def test_the_legacy_shape_still_collides(): + """The opt-in legacy shape reproduces the bug — that is what makes it transitional.""" + assert _key(project_id=PROJECT_A, legacy_truncated_scope=True) == _key( + project_id=PROJECT_B, legacy_truncated_scope=True + ) + + +# LOCKS ACROSS A ROLLING DEPLOY ------------------------------------------------ + + +@pytest_asyncio.fixture +async def fake_redis(): + """Point the lock engine at an in-memory redis, as `test_evaluation_runtime_locks`. + + fakeredis executes Lua only with the optional `lupa` backend, which is not a + dependency here, so the two ownership scripts are supplied as an equivalent shim. + Everything else — `acquire_lock`, `renew_lock`, `release_lock` — runs as written. + """ + fakeredis = pytest.importorskip("fakeredis") + aioredis = pytest.importorskip("fakeredis.aioredis") + engine = pytest.importorskip("oss.src.dbs.redis.shared.engine") + + server = fakeredis.FakeServer() + client = aioredis.FakeRedis(server=server, decode_responses=False) + + async def _eval(script, numkeys, *args): + """`GET`, compare against the owner token, then `EXPIRE` or `DEL`.""" + lock_key, owner, *rest = args + + if await client.get(lock_key) != owner.encode(): + return 0 + + if script == locking._LOCK_RENEW_IF_OWNER_SCRIPT: + return int(bool(await client.expire(lock_key, int(rest[0])))) + + return int(bool(await client.delete(lock_key))) + + lock_engine = engine.get_lock_engine() + + with ( + patch.object(lock_engine, "_client", return_value=client), + # `LockEngine.__getattr__` forwards to the client; a real attribute shadows it. + patch.object(lock_engine, "eval", _eval, create=True), + ): + yield client + + await client.aclose() + + +async def test_two_ordinary_projects_do_not_share_a_lock(fake_redis): + project_a, project_b = str(uuid4()), str(uuid4()) + + assert ( + await locking.acquire_lock(namespace="eval", key="run", project_id=project_a) + is not None + ) + assert ( + await locking.acquire_lock(namespace="eval", key="run", project_id=project_b) + is not None + ) + + +async def test_colliding_projects_still_share_a_lock_while_the_cover_lasts(fake_redis): + """A deliberate trade-off, confined to the lock namespace and to one release. + + Covering the previous release's key means two projects whose ids share a suffix keep + serializing against each other on locks. Dropping the cover instead would let pods on + either side of a rolling deploy into the same critical section, and losing mutual + exclusion is worse than two unrelated tenants queueing. Removing the transitional + cover (see `locking`) is what closes this last case. + + Cache keys are separated immediately either way — that is where the permission caches + live, and it is the part with a security consequence. + """ + assert _key(project_id=PROJECT_A) != _key(project_id=PROJECT_B) + + held = await locking.acquire_lock(namespace="eval", key="run", project_id=PROJECT_A) + assert held is not None + + assert ( + await locking.acquire_lock(namespace="eval", key="run", project_id=PROJECT_B) + is None + ) + + +async def test_a_lock_still_excludes_the_same_project(fake_redis): + held = await locking.acquire_lock(namespace="eval", key="run", project_id=PROJECT_A) + assert held is not None + + assert ( + await locking.acquire_lock(namespace="eval", key="run", project_id=PROJECT_A) + is None + ) + + await locking.release_lock( + namespace="eval", key="run", project_id=PROJECT_A, owner=held + ) + assert ( + await locking.acquire_lock(namespace="eval", key="run", project_id=PROJECT_A) + is not None + ) + + +async def test_a_holder_on_the_previous_release_still_blocks_this_one(fake_redis): + """The rolling-deploy case: an old pod holds only the truncated key. + + Without the transitional cover this is the window where two pods both believe they + hold the same logical lock. + """ + legacy_key = _pack( + namespace="lock:eval", + key="run", + project_id=PROJECT_A, + legacy_truncated_scope=True, + ) + # Exactly what a pod running the previous release writes. + await fake_redis.set(legacy_key, b"old-pod-owner", nx=True, ex=30) + + assert ( + await locking.acquire_lock(namespace="eval", key="run", project_id=PROJECT_A) + is None + ) + + +async def test_a_blocked_acquire_does_not_strand_the_legacy_key(fake_redis): + """Losing the race on the primary key must not leave the legacy one held. + + The legacy key is claimed first, so a caller that then loses the primary has to give + it back — otherwise a failed acquire would block the section for the full TTL. + """ + lock_key, legacy_key = locking._lock_keys( + namespace="eval", key="run", project_id=PROJECT_A + ) + assert legacy_key is not None + + # Someone already holds the primary; the caller below will claim the legacy key, + # fail on the primary, and must then release what it took. + await fake_redis.set(lock_key, b"another-owner", nx=True, ex=30) + + assert ( + await locking.acquire_lock(namespace="eval", key="run", project_id=PROJECT_A) + is None + ) + assert await fake_redis.get(legacy_key) is None + + +async def test_release_clears_both_generations(fake_redis): + lock_key, legacy_key = locking._lock_keys( + namespace="eval", key="run", project_id=PROJECT_A + ) + assert legacy_key is not None + + owner = await locking.acquire_lock( + namespace="eval", key="run", project_id=PROJECT_A + ) + assert owner is not None + assert await fake_redis.get(legacy_key) is not None + + await locking.release_lock( + namespace="eval", key="run", project_id=PROJECT_A, owner=owner + ) + + assert await fake_redis.get(lock_key) is None + assert await fake_redis.get(legacy_key) is None + + +async def test_a_short_scope_takes_only_one_key(fake_redis): + """When the two shapes coincide there is no second key — and no double SET NX. + + Claiming the same key twice with NX would fail the second call and make every + acquire in this shape look blocked. + """ + short_project = "abc" + lock_key, legacy_key = locking._lock_keys( + namespace="eval", key="run", project_id=short_project + ) + assert legacy_key is None + + assert ( + await locking.acquire_lock( + namespace="eval", key="run", project_id=short_project + ) + is not None + ) + assert await fake_redis.get(lock_key) is not None + + +async def test_renew_keeps_both_generations_alive(fake_redis): + lock_key, legacy_key = locking._lock_keys( + namespace="eval", key="run", project_id=PROJECT_A + ) + assert legacy_key is not None + + owner = await locking.acquire_lock( + namespace="eval", key="run", project_id=PROJECT_A, ttl=5 + ) + assert owner is not None + + assert await locking.renew_lock( + namespace="eval", key="run", project_id=PROJECT_A, ttl=90, owner=owner + ) + + # A legacy key left on the original TTL would lapse mid-section and let a pod on the + # previous release in. + assert await fake_redis.ttl(lock_key) > 5 + assert await fake_redis.ttl(legacy_key) > 5