|
| 1 | +"""Revoke-by-id and the failure paths of the key-rotation lifecycle. |
| 2 | +
|
| 3 | +``tests/unit/test_key_rotation.py`` pins the happy lifecycle. This file covers |
| 4 | +what it does not reach: ``revoke_key_by_id`` -- the whole reason F-02 A stopped |
| 5 | +taking the plaintext key in a URL, and the only revoke path the admin router |
| 6 | +now calls -- plus the three failures the rotator has to distinguish rather than |
| 7 | +crash on: a read-only key store, a store that is broken rather than read-only, |
| 8 | +and a background grace-period cleanup that fails after the caller is gone. |
| 9 | +
|
| 10 | +Direct rotator calls: no TestClient, so the CI key-rotation coverage gate can |
| 11 | +run this file. |
| 12 | +""" |
| 13 | + |
| 14 | +from __future__ import annotations |
| 15 | + |
| 16 | +import errno |
| 17 | +from datetime import date |
| 18 | +from pathlib import Path |
| 19 | + |
| 20 | +import pytest |
| 21 | +import structlog |
| 22 | + |
| 23 | +from agentflow_runtime.serving.api.auth.key_rotation import KeyStoreReadOnlyError |
| 24 | +from agentflow_runtime.serving.api.auth.manager import ( |
| 25 | + ApiKeysConfig, |
| 26 | + AuthManager, |
| 27 | + KeyCreateRequest, |
| 28 | + TenantKey, |
| 29 | +) |
| 30 | + |
| 31 | +SEED_KEY_YAML = ( |
| 32 | + "keys:\n" |
| 33 | + ' - key: "rotation-acme-key"\n' |
| 34 | + ' name: "Rotation Agent"\n' |
| 35 | + ' tenant: "acme"\n' |
| 36 | + " rate_limit_rpm: 100\n" |
| 37 | + " allowed_entity_types: null\n" |
| 38 | + ' created_at: "2026-04-10"\n' |
| 39 | +) |
| 40 | + |
| 41 | + |
| 42 | +@pytest.fixture(autouse=True) |
| 43 | +def _uncached_auth_logger(monkeypatch: pytest.MonkeyPatch) -> None: |
| 44 | + from agentflow_runtime.serving.api import auth as auth_package |
| 45 | + |
| 46 | + monkeypatch.setattr(auth_package, "logger", structlog.get_logger()) |
| 47 | + |
| 48 | + |
| 49 | +@pytest.fixture |
| 50 | +def manager(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> AuthManager: |
| 51 | + monkeypatch.setenv("AGENTFLOW_USAGE_DB_PATH", str(tmp_path / "usage.duckdb")) |
| 52 | + api_keys_path = tmp_path / "config" / "api_keys.yaml" |
| 53 | + api_keys_path.parent.mkdir(parents=True, exist_ok=True) |
| 54 | + api_keys_path.write_text(SEED_KEY_YAML, encoding="utf-8", newline="\n") |
| 55 | + mgr = AuthManager( |
| 56 | + api_keys_path=api_keys_path, |
| 57 | + db_path=tmp_path / "usage.duckdb", |
| 58 | + admin_key="admin-secret", |
| 59 | + ) |
| 60 | + # bcrypt-12 is slow and every create/rotate hashes several times; the hash |
| 61 | + # format is unchanged, only the cost factor. |
| 62 | + mgr.security_policy.bcrypt_rounds = 4 |
| 63 | + mgr.load() |
| 64 | + mgr.ensure_usage_table() |
| 65 | + try: |
| 66 | + yield mgr |
| 67 | + finally: |
| 68 | + mgr.shutdown() |
| 69 | + |
| 70 | + |
| 71 | +def _only_key_id(manager: AuthManager) -> str: |
| 72 | + keys = manager.list_keys_with_usage() |
| 73 | + assert keys, "expected at least one configured key" |
| 74 | + key_id = keys[0]["key_id"] |
| 75 | + assert key_id is not None |
| 76 | + return str(key_id) |
| 77 | + |
| 78 | + |
| 79 | +def test_revoking_by_key_id_removes_the_entry(manager: AuthManager) -> None: |
| 80 | + key_id = _only_key_id(manager) |
| 81 | + |
| 82 | + assert manager.revoke_key_by_id(key_id) is True |
| 83 | + |
| 84 | + assert manager.list_keys_with_usage() == [] |
| 85 | + |
| 86 | + |
| 87 | +def test_revoking_an_unknown_key_id_changes_nothing(manager: AuthManager) -> None: |
| 88 | + before = manager.list_keys_with_usage() |
| 89 | + |
| 90 | + assert manager.revoke_key_by_id("no-such-key-id") is False |
| 91 | + |
| 92 | + # No write, no reload: an admin typo must not disturb the live store. |
| 93 | + assert manager.list_keys_with_usage() == before |
| 94 | + |
| 95 | + |
| 96 | +def test_revoking_by_key_id_stops_a_cached_plaintext_from_authenticating( |
| 97 | + manager: AuthManager, |
| 98 | +) -> None: |
| 99 | + created = manager.create_key( |
| 100 | + KeyCreateRequest(name="Support Agent", tenant="globex", rate_limit_rpm=7) |
| 101 | + ) |
| 102 | + assert created.key is not None |
| 103 | + plaintext = created.key |
| 104 | + assert manager.authenticate(plaintext) is not None |
| 105 | + assert created.key_hash in manager._runtime_plaintext_by_hash |
| 106 | + assert created.key_id is not None |
| 107 | + |
| 108 | + assert manager.revoke_key_by_id(created.key_id) is True |
| 109 | + |
| 110 | + # The plaintext->key cache exists to skip the slow verify. If revoke left an |
| 111 | + # entry behind, the revoked key would keep authenticating from memory for |
| 112 | + # the life of the process. |
| 113 | + assert created.key_hash not in manager._runtime_plaintext_by_hash |
| 114 | + assert manager.authenticate(plaintext) is None |
| 115 | + |
| 116 | + |
| 117 | +def test_rotating_an_entry_with_no_key_material_is_refused(manager: AuthManager) -> None: |
| 118 | + key_id = _only_key_id(manager) |
| 119 | + stored = manager._load_config().keys[0] |
| 120 | + hollow = stored.model_copy(update={"key": None, "key_hash": None}) |
| 121 | + |
| 122 | + # A store on disk cannot produce this: ApiKeysConfig validation rejects an |
| 123 | + # entry with neither key nor key_hash, which is why the config has to be |
| 124 | + # built unvalidated to get here at all. The branch is defensive -- it is the |
| 125 | + # difference between a named refusal and a TypeError inside hash_api_key. |
| 126 | + def _hollow_config() -> ApiKeysConfig: |
| 127 | + return ApiKeysConfig.model_construct(keys=[hollow]) |
| 128 | + |
| 129 | + manager._load_config = _hollow_config # type: ignore[method-assign] |
| 130 | + |
| 131 | + with pytest.raises(ValueError, match="Current key material is unavailable"): |
| 132 | + manager.rotate_key(key_id) |
| 133 | + |
| 134 | + |
| 135 | +def test_a_read_only_store_is_reported_as_read_only( |
| 136 | + manager: AuthManager, monkeypatch: pytest.MonkeyPatch |
| 137 | +) -> None: |
| 138 | + config = manager._load_config() |
| 139 | + |
| 140 | + def _denied(self: Path, *args: object, **kwargs: object) -> int: |
| 141 | + raise PermissionError(errno.EACCES, "Permission denied") |
| 142 | + |
| 143 | + monkeypatch.setattr(Path, "write_text", _denied) |
| 144 | + |
| 145 | + with pytest.raises(KeyStoreReadOnlyError): |
| 146 | + manager._key_rotator.write_config(config) |
| 147 | + # The manager downgrades in place, so the next admin mutation answers 409 |
| 148 | + # instead of probing the mount again. |
| 149 | + assert manager._key_store_writable is False |
| 150 | + |
| 151 | + |
| 152 | +def test_a_broken_store_is_not_reported_as_read_only( |
| 153 | + manager: AuthManager, monkeypatch: pytest.MonkeyPatch |
| 154 | +) -> None: |
| 155 | + config = manager._load_config() |
| 156 | + |
| 157 | + def _io_error(self: Path, *args: object, **kwargs: object) -> int: |
| 158 | + raise OSError(errno.EIO, "I/O error") |
| 159 | + |
| 160 | + monkeypatch.setattr(Path, "write_text", _io_error) |
| 161 | + |
| 162 | + with pytest.raises(OSError) as excinfo: |
| 163 | + manager._key_rotator.write_config(config) |
| 164 | + |
| 165 | + assert not isinstance(excinfo.value, KeyStoreReadOnlyError) |
| 166 | + assert excinfo.value.errno == errno.EIO |
| 167 | + # A failing disk is not a read-only mount: leaving the flag alone keeps the |
| 168 | + # 409 for the case that really means it. |
| 169 | + assert manager._key_store_writable is not False |
| 170 | + |
| 171 | + |
| 172 | +def test_a_failed_grace_period_cleanup_is_logged_and_swallowed(manager: AuthManager) -> None: |
| 173 | + def _boom(_key_id: str) -> bool: |
| 174 | + raise RuntimeError("store went away") |
| 175 | + |
| 176 | + manager._key_rotator.revoke_old_key = _boom # type: ignore[method-assign] |
| 177 | + |
| 178 | + with structlog.testing.capture_logs() as events: |
| 179 | + manager._key_rotator.expire_previous_key("acme-ops-abcd") |
| 180 | + |
| 181 | + # This runs on a timer thread with no caller left to catch it: an exception |
| 182 | + # here would be swallowed by threading with no record at all. |
| 183 | + assert any(event.get("event") == "api_key_rotation_cleanup_failed" for event in events) |
| 184 | + assert any("store went away" in str(event.get("error", "")) for event in events) |
| 185 | + |
| 186 | + |
| 187 | +def test_a_cleanup_for_an_unknown_key_is_silent(manager: AuthManager) -> None: |
| 188 | + def _missing(_key_id: str) -> bool: |
| 189 | + raise KeyError("gone") |
| 190 | + |
| 191 | + manager._key_rotator.revoke_old_key = _missing # type: ignore[method-assign] |
| 192 | + |
| 193 | + with structlog.testing.capture_logs() as events: |
| 194 | + manager._key_rotator.expire_previous_key("acme-ops-abcd") |
| 195 | + |
| 196 | + # A key revoked before its grace timer fired is the expected race, not a |
| 197 | + # fault worth waking anyone for. |
| 198 | + assert not [ |
| 199 | + event for event in events if event.get("event") == "api_key_rotation_cleanup_failed" |
| 200 | + ] |
| 201 | + |
| 202 | + |
| 203 | +def test_a_created_key_can_be_revoked_by_its_id_only(manager: AuthManager) -> None: |
| 204 | + created = manager.create_key( |
| 205 | + KeyCreateRequest(name="Second Agent", tenant="globex", rate_limit_rpm=5) |
| 206 | + ) |
| 207 | + assert created.key_id is not None |
| 208 | + |
| 209 | + # F-02 A: the plaintext is not an identifier the admin surface accepts. It |
| 210 | + # is still valid key material, so revoke must not resolve it as an id. |
| 211 | + assert manager.revoke_key_by_id(str(created.key)) is False |
| 212 | + assert manager.revoke_key_by_id(created.key_id) is True |
| 213 | + |
| 214 | + remaining = {key["key_id"] for key in manager.list_keys_with_usage()} |
| 215 | + assert created.key_id not in remaining |
| 216 | + |
| 217 | + |
| 218 | +def test_the_seeded_entry_survives_an_unrelated_revoke(manager: AuthManager) -> None: |
| 219 | + seeded = _only_key_id(manager) |
| 220 | + created = manager.create_key( |
| 221 | + KeyCreateRequest(name="Third Agent", tenant="globex", rate_limit_rpm=5) |
| 222 | + ) |
| 223 | + assert created.key_id is not None |
| 224 | + |
| 225 | + assert manager.revoke_key_by_id(created.key_id) is True |
| 226 | + |
| 227 | + remaining = {key["key_id"] for key in manager.list_keys_with_usage()} |
| 228 | + assert remaining == {seeded} |
| 229 | + |
| 230 | + |
| 231 | +def test_a_revoked_entry_leaves_no_previous_slot_behind(manager: AuthManager) -> None: |
| 232 | + key_id = _only_key_id(manager) |
| 233 | + rotated, _expires_at = manager.rotate_key(key_id) |
| 234 | + assert rotated.previous_key_hash is not None |
| 235 | + |
| 236 | + assert manager.revoke_key_by_id(key_id) is True |
| 237 | + |
| 238 | + stored: list[TenantKey] = manager._load_config().keys |
| 239 | + assert [item for item in stored if item.key_id == key_id] == [] |
| 240 | + assert manager.list_keys_with_usage() == [] |
| 241 | + |
| 242 | + |
| 243 | +def test_revoking_the_last_key_leaves_a_loadable_store(manager: AuthManager) -> None: |
| 244 | + key_id = _only_key_id(manager) |
| 245 | + |
| 246 | + assert manager.revoke_key_by_id(key_id) is True |
| 247 | + |
| 248 | + reloaded = AuthManager( |
| 249 | + api_keys_path=manager.api_keys_path, |
| 250 | + db_path=manager.db_path, |
| 251 | + admin_key="admin-secret", |
| 252 | + ) |
| 253 | + reloaded.load() |
| 254 | + try: |
| 255 | + # An empty store is a valid store: writing something a later load() |
| 256 | + # cannot parse would take the API down on the next restart. |
| 257 | + assert reloaded.configured_key_count == 0 |
| 258 | + assert reloaded.authenticate("rotation-acme-key") is None |
| 259 | + finally: |
| 260 | + reloaded.shutdown() |
| 261 | + |
| 262 | + |
| 263 | +def test_a_revoked_key_id_is_gone_from_rotation_status(manager: AuthManager) -> None: |
| 264 | + created = manager.create_key( |
| 265 | + KeyCreateRequest(name="Fourth Agent", tenant="globex", rate_limit_rpm=5) |
| 266 | + ) |
| 267 | + assert created.key_id is not None |
| 268 | + assert manager.get_rotation_status(created.key_id)["phase"] in {"idle", "grace_period"} |
| 269 | + |
| 270 | + assert manager.revoke_key_by_id(created.key_id) is True |
| 271 | + |
| 272 | + with pytest.raises(KeyError): |
| 273 | + manager.get_rotation_status(created.key_id) |
| 274 | + |
| 275 | + |
| 276 | +def test_the_seed_entry_is_created_with_a_date(manager: AuthManager) -> None: |
| 277 | + stored = manager._load_config().keys[0] |
| 278 | + |
| 279 | + # ensure_key_ids() rewrites the store on first load; the created_at it |
| 280 | + # writes back has to survive that round trip as a date, not a string. |
| 281 | + assert isinstance(stored.created_at, date) |
0 commit comments