Skip to content

Commit 38bb4fa

Browse files
deepthi-rao-scaledeepthi-rao-scale
committed
fix(authz): avoid stale authorization checks
1 parent 9a5b28f commit 38bb4fa

4 files changed

Lines changed: 246 additions & 3 deletions

File tree

‎agentex/src/api/authentication_cache.py‎

Lines changed: 56 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,14 @@ async def clear(self) -> None:
8585
async with self._lock:
8686
self.cache.clear()
8787

88+
async def delete_by_prefix(self, prefix: str) -> int:
89+
"""Delete cache entries whose keys start with the given prefix."""
90+
async with self._lock:
91+
keys = [key for key in self.cache if key.startswith(prefix)]
92+
for key in keys:
93+
del self.cache[key]
94+
return len(keys)
95+
8896
async def remove_expired(self) -> None:
8997
"""Remove all expired entries from cache."""
9098
async with self._lock:
@@ -257,6 +265,19 @@ def _contains_api_key(principal_context: Any) -> bool:
257265

258266
# Authorization Check Cache Methods (Async)
259267

268+
@staticmethod
269+
def _create_authorization_resource_cache_prefix(
270+
resource_type: str,
271+
resource_selector: str,
272+
) -> str:
273+
resource_key = AuthenticationCache._hash_dict(
274+
{
275+
"resource_type": resource_type,
276+
"resource_selector": resource_selector,
277+
}
278+
)
279+
return f"authz:{resource_key}:"
280+
260281
@staticmethod
261282
def _create_authorization_cache_key(
262283
resource_type: str,
@@ -303,13 +324,16 @@ def _create_authorization_cache_key(
303324

304325
# Create the cache key components
305326
cache_data = {
306-
"resource_type": resource_type,
307-
"resource_selector": resource_selector,
308327
"operation": operation,
309328
"principal": principal_key_data,
310329
}
311330

312-
return f"authz:{AuthenticationCache._hash_dict(cache_data)}"
331+
return (
332+
AuthenticationCache._create_authorization_resource_cache_prefix(
333+
resource_type, resource_selector
334+
)
335+
+ AuthenticationCache._hash_dict(cache_data)
336+
)
313337

314338
async def get_authorization_check(
315339
self,
@@ -319,6 +343,12 @@ async def get_authorization_check(
319343
principal_context: Any,
320344
) -> bool | None:
321345
"""Get cached authorization check result."""
346+
if self._contains_api_key(principal_context):
347+
logger.debug(
348+
"Skipping authorization check cache lookup for API key principal"
349+
)
350+
return None
351+
322352
cache_key = self._create_authorization_cache_key(
323353
resource_type, resource_selector, operation, principal_context
324354
)
@@ -339,6 +369,12 @@ async def set_authorization_check(
339369
allowed: bool,
340370
) -> None:
341371
"""Cache authorization check result."""
372+
if self._contains_api_key(principal_context):
373+
logger.debug(
374+
"Skipping authorization check cache write for API key principal"
375+
)
376+
return
377+
342378
cache_key = self._create_authorization_cache_key(
343379
resource_type, resource_selector, operation, principal_context
344380
)
@@ -358,6 +394,23 @@ async def clear_all(self) -> None:
358394
await self.authorization_check_cache.clear()
359395
logger.info("All authentication and authorization caches cleared")
360396

397+
async def clear_authorization_checks_for_resource(
398+
self,
399+
resource_type: str,
400+
resource_selector: str,
401+
) -> None:
402+
"""Clear cached authorization check results for one resource."""
403+
prefix = self._create_authorization_resource_cache_prefix(
404+
resource_type, resource_selector
405+
)
406+
deleted = await self.authorization_check_cache.delete_by_prefix(prefix)
407+
logger.info(
408+
"Authorization check cache cleared for %s:%s entries=%d",
409+
resource_type,
410+
resource_selector,
411+
deleted,
412+
)
413+
361414
async def cleanup_expired(self) -> None:
362415
"""Remove expired entries from all caches."""
363416
await self.agent_identity_cache.remove_expired()

‎agentex/src/domain/services/authorization_service.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,13 @@ def _bypass(self) -> bool:
3939
def is_enabled(self) -> bool:
4040
return self.enabled
4141

42+
async def _clear_authorization_cache(self, resource: AgentexResource) -> None:
43+
auth_cache = await get_auth_cache()
44+
await auth_cache.clear_authorization_checks_for_resource(
45+
resource_type=str(resource.type),
46+
resource_selector=resource.selector,
47+
)
48+
4249
async def grant(
4350
self, resource: AgentexResource, *, commit: bool = True, principal_context=...
4451
) -> None:
@@ -61,6 +68,7 @@ async def grant(
6168
resource,
6269
AuthorizedOperationType.create,
6370
)
71+
await self._clear_authorization_cache(resource)
6472
return result
6573

6674
async def revoke(
@@ -87,6 +95,7 @@ async def revoke(
8795
logger.info(
8896
f"Revoked {AuthorizedOperationType.delete} permission on {resource.type}:{resource.selector}"
8997
)
98+
await self._clear_authorization_cache(resource)
9099
return result
91100

92101
async def check(
@@ -214,6 +223,7 @@ async def register_resource(
214223
f"{parent.type}:{parent.selector}" if parent is not None else None,
215224
)
216225
await self.gateway.register_resource(effective_principal, resource, parent)
226+
await self._clear_authorization_cache(resource)
217227

218228
async def deregister_resource(
219229
self,
@@ -237,6 +247,7 @@ async def deregister_resource(
237247
resource.selector,
238248
)
239249
await self.gateway.deregister_resource(effective_principal, resource)
250+
await self._clear_authorization_cache(resource)
240251

241252

242253
DAuthorizationService = Annotated[AuthorizationService, Depends(AuthorizationService)]

‎agentex/tests/unit/api/test_authentication_cache_metrics.py‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,3 +102,53 @@ async def test_auth_gateway_response_with_api_key_principal_is_not_cached():
102102
await cache.set_auth_gateway_response(headers, principal)
103103

104104
assert await cache.get_auth_gateway_response(headers) is None
105+
106+
107+
@pytest.mark.unit
108+
@pytest.mark.asyncio
109+
async def test_authorization_check_with_api_key_principal_is_not_cached():
110+
cache = AuthenticationCache()
111+
principal = {"user_id": "user-1", "account_id": "acct-1", "api_key": "secret-key"}
112+
113+
await cache.set_authorization_check(
114+
resource_type="agent",
115+
resource_selector="agent-1",
116+
operation="execute",
117+
principal_context=principal,
118+
allowed=True,
119+
)
120+
121+
assert (
122+
await cache.get_authorization_check(
123+
resource_type="agent",
124+
resource_selector="agent-1",
125+
operation="execute",
126+
principal_context=principal,
127+
)
128+
is None
129+
)
130+
131+
132+
@pytest.mark.unit
133+
@pytest.mark.asyncio
134+
async def test_authorization_check_without_api_key_principal_is_cached():
135+
cache = AuthenticationCache()
136+
principal = {"user_id": "user-1", "account_id": "acct-1"}
137+
138+
await cache.set_authorization_check(
139+
resource_type="agent",
140+
resource_selector="agent-1",
141+
operation="read",
142+
principal_context=principal,
143+
allowed=True,
144+
)
145+
146+
assert (
147+
await cache.get_authorization_check(
148+
resource_type="agent",
149+
resource_selector="agent-1",
150+
operation="read",
151+
principal_context=principal,
152+
)
153+
is True
154+
)

‎agentex/tests/unit/services/test_authorization_service_cache.py‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
from types import SimpleNamespace
2+
from unittest.mock import AsyncMock
3+
4+
import pytest
5+
from src.api.authentication_cache import reset_auth_cache
6+
from src.api.schemas.authorization_types import AgentexResource, AuthorizedOperationType
7+
from src.domain.services.authorization_service import AuthorizationService
8+
9+
10+
def _request_with_principal(principal_context):
11+
return SimpleNamespace(
12+
state=SimpleNamespace(
13+
principal_context=principal_context,
14+
agent_identity=None,
15+
)
16+
)
17+
18+
19+
def _service(principal_context, gateway):
20+
return AuthorizationService(
21+
enabled=True,
22+
gateway=gateway,
23+
request=_request_with_principal(principal_context),
24+
)
25+
26+
27+
@pytest.mark.unit
28+
@pytest.mark.asyncio
29+
async def test_api_key_principal_authorization_check_calls_gateway_each_time():
30+
await reset_auth_cache()
31+
try:
32+
gateway = AsyncMock()
33+
gateway.check.return_value = True
34+
service = _service(
35+
{"user_id": "user-1", "account_id": "acct-1", "api_key": "secret-key"},
36+
gateway,
37+
)
38+
resource = AgentexResource.agent("agent-1")
39+
40+
assert await service.check(resource, AuthorizedOperationType.execute) is True
41+
assert await service.check(resource, AuthorizedOperationType.execute) is True
42+
43+
assert gateway.check.await_count == 2
44+
finally:
45+
await reset_auth_cache()
46+
47+
48+
@pytest.mark.unit
49+
@pytest.mark.asyncio
50+
async def test_non_api_key_principal_authorization_check_uses_cache():
51+
await reset_auth_cache()
52+
try:
53+
gateway = AsyncMock()
54+
gateway.check.return_value = True
55+
service = _service({"user_id": "user-1", "account_id": "acct-1"}, gateway)
56+
resource = AgentexResource.agent("agent-1")
57+
58+
assert await service.check(resource, AuthorizedOperationType.read) is True
59+
assert await service.check(resource, AuthorizedOperationType.read) is True
60+
61+
assert gateway.check.await_count == 1
62+
finally:
63+
await reset_auth_cache()
64+
65+
66+
@pytest.mark.unit
67+
@pytest.mark.asyncio
68+
@pytest.mark.parametrize(
69+
"mutation",
70+
["grant", "revoke", "register_resource", "deregister_resource"],
71+
)
72+
async def test_authorization_mutations_clear_cached_authorization_checks(mutation):
73+
await reset_auth_cache()
74+
try:
75+
gateway = AsyncMock()
76+
gateway.check.return_value = True
77+
service = _service({"user_id": "user-1", "account_id": "acct-1"}, gateway)
78+
resource = AgentexResource.agent("agent-1")
79+
80+
assert await service.check(resource, AuthorizedOperationType.read) is True
81+
assert await service.check(resource, AuthorizedOperationType.read) is True
82+
assert gateway.check.await_count == 1
83+
84+
await getattr(service, mutation)(resource)
85+
86+
assert await service.check(resource, AuthorizedOperationType.read) is True
87+
assert gateway.check.await_count == 2
88+
finally:
89+
await reset_auth_cache()
90+
91+
92+
@pytest.mark.unit
93+
@pytest.mark.asyncio
94+
@pytest.mark.parametrize(
95+
"mutation",
96+
["grant", "revoke", "register_resource", "deregister_resource"],
97+
)
98+
async def test_authorization_mutations_only_clear_checks_for_mutated_resource(
99+
mutation,
100+
):
101+
await reset_auth_cache()
102+
try:
103+
gateway = AsyncMock()
104+
gateway.check.return_value = True
105+
service = _service({"user_id": "user-1", "account_id": "acct-1"}, gateway)
106+
changed_resource = AgentexResource.agent("agent-1")
107+
unchanged_resource = AgentexResource.agent("agent-2")
108+
109+
assert (
110+
await service.check(changed_resource, AuthorizedOperationType.read) is True
111+
)
112+
assert (
113+
await service.check(unchanged_resource, AuthorizedOperationType.read)
114+
is True
115+
)
116+
assert gateway.check.await_count == 2
117+
118+
await getattr(service, mutation)(changed_resource)
119+
120+
assert (
121+
await service.check(changed_resource, AuthorizedOperationType.read) is True
122+
)
123+
assert (
124+
await service.check(unchanged_resource, AuthorizedOperationType.read)
125+
is True
126+
)
127+
assert gateway.check.await_count == 3
128+
finally:
129+
await reset_auth_cache()

0 commit comments

Comments
 (0)