feat(server): add attribute based access control for resources

This commit is contained in:
Ashwin Bharambe 2025-03-19 08:10:56 -07:00
parent 7c0448456e
commit b937a49436
8 changed files with 862 additions and 35 deletions

View file

@ -13,6 +13,15 @@ from fastapi.testclient import TestClient
from llama_stack.distribution.server.auth import AuthenticationMiddleware
class MockResponse:
def __init__(self, status_code, json_data):
self.status_code = status_code
self._json_data = json_data
def json(self):
return self._json_data
@pytest.fixture
def mock_auth_endpoint():
return "http://mock-auth-service/validate"
@ -45,6 +54,26 @@ def client(app):
return TestClient(app)
@pytest.fixture
def mock_scope():
return {
"type": "http",
"path": "/models/list",
"headers": [
(b"content-type", b"application/json"),
(b"authorization", b"Bearer test-api-key"),
(b"user-agent", b"test-user-agent"),
],
"query_string": b"limit=100&offset=0",
}
@pytest.fixture
def mock_middleware(mock_auth_endpoint):
mock_app = AsyncMock()
return AuthenticationMiddleware(mock_app, mock_auth_endpoint), mock_app
async def mock_post_success(*args, **kwargs):
mock_response = AsyncMock()
mock_response.status_code = 200
@ -122,3 +151,59 @@ def test_auth_request_payload(client, valid_api_key, mock_auth_endpoint):
assert "authorization" in payload["request"]["headers"]
assert "param1" in payload["request"]["params"]
assert "param2" in payload["request"]["params"]
@pytest.mark.asyncio
async def test_auth_middleware_with_access_attributes(mock_middleware, mock_scope):
middleware, mock_app = mock_middleware
mock_receive = AsyncMock()
mock_send = AsyncMock()
with patch("httpx.AsyncClient") as mock_client:
mock_client_instance = AsyncMock()
mock_client.return_value.__aenter__.return_value = mock_client_instance
mock_client_instance.post.return_value = MockResponse(
200,
{
"access_attributes": {
"roles": ["admin", "user"],
"teams": ["ml-team"],
"projects": ["project-x", "project-y"],
}
},
)
await middleware(mock_scope, mock_receive, mock_send)
assert "user_attributes" in mock_scope
assert mock_scope["user_attributes"]["roles"] == ["admin", "user"]
assert mock_scope["user_attributes"]["teams"] == ["ml-team"]
assert mock_scope["user_attributes"]["projects"] == ["project-x", "project-y"]
mock_app.assert_called_once_with(mock_scope, mock_receive, mock_send)
@pytest.mark.asyncio
async def test_auth_middleware_no_attributes(mock_middleware, mock_scope):
"""Test middleware behavior with no access attributes"""
middleware, mock_app = mock_middleware
mock_receive = AsyncMock()
mock_send = AsyncMock()
with patch("httpx.AsyncClient") as mock_client:
mock_client_instance = AsyncMock()
mock_client.return_value.__aenter__.return_value = mock_client_instance
mock_client_instance.post.return_value = MockResponse(
200,
{
"message": "Authentication successful"
# No access_attributes
},
)
await middleware(mock_scope, mock_receive, mock_send)
assert "user_attributes" in mock_scope
assert mock_scope["user_attributes"] == {}