From d4059af45e33d9decc5dee89e1a06f3effd6a262 Mon Sep 17 00:00:00 2001 From: Tim O'Farrell Date: Thu, 2 Jul 2026 13:08:34 -0600 Subject: [PATCH] feat(saas-auth): accept api_key cookie as a fallback credential (#15101) Co-authored-by: openhands --- enterprise/server/auth/saas_user_auth.py | 9 +- enterprise/server/middleware.py | 2 + enterprise/tests/unit/test_auth_middleware.py | 38 +++++ enterprise/tests/unit/test_saas_user_auth.py | 137 +++++++++++++++++- 4 files changed, 182 insertions(+), 4 deletions(-) diff --git a/enterprise/server/auth/saas_user_auth.py b/enterprise/server/auth/saas_user_auth.py index 17fa3eadf0..bc01d10fa9 100644 --- a/enterprise/server/auth/saas_user_auth.py +++ b/enterprise/server/auth/saas_user_auth.py @@ -741,7 +741,14 @@ def get_api_key_from_header(request: Request): return session_api_key # Fallback to X-Access-Token header as an additional option - return request.headers.get('X-Access-Token') + x_access_token = request.headers.get('X-Access-Token') + if x_access_token: + return x_access_token + + # Fallback to the `api_key` cookie, which mirrors the X-Access-Token header + # Security note: This cookie MUST be marked `Secure; HttpOnly; SameSite=Strict` + # (or `Lax`) to mitigate CSRF and XSS risks. + return request.cookies.get('api_key') async def saas_user_auth_from_bearer(request: Request) -> SaasUserAuth | None: diff --git a/enterprise/server/middleware.py b/enterprise/server/middleware.py index d8ce74dcc9..0200658452 100644 --- a/enterprise/server/middleware.py +++ b/enterprise/server/middleware.py @@ -117,12 +117,14 @@ class SetAuthCookieMiddleware: auth_header = request.headers.get('Authorization') mcp_auth_header = request.headers.get('X-Session-API-Key') api_auth_header = request.headers.get('X-Access-Token') + api_key_cookie = request.cookies.get('api_key') accepted_tos: bool | None = False if ( keycloak_auth_cookie is None and (auth_header is None or not auth_header.startswith('Bearer ')) and mcp_auth_header is None and api_auth_header is None + and api_key_cookie is None ): raise NoCredentialsError diff --git a/enterprise/tests/unit/test_auth_middleware.py b/enterprise/tests/unit/test_auth_middleware.py index c25dcfa7b1..289c2640e1 100644 --- a/enterprise/tests/unit/test_auth_middleware.py +++ b/enterprise/tests/unit/test_auth_middleware.py @@ -122,6 +122,44 @@ async def test_middleware_with_cookie_and_refresh( ) +@pytest.mark.asyncio +async def test_middleware_with_api_key_cookie_only( + middleware, mock_request, mock_response +): + """A request authenticated solely via the api_key cookie must not be treated as unauthenticated. + + The middleware's _check_tos would otherwise raise NoCredentialsError + when the only credential on the request is the new ``api_key`` cookie. + """ + mock_request.cookies = {'api_key': 'test_api_key'} + mock_request.headers = {} + mock_request.url = MagicMock() + mock_request.url.hostname = 'localhost' + mock_request.url.path = '/api/some/endpoint' + mock_call_next = AsyncMock(return_value=mock_response) + + result = await middleware(mock_request, mock_call_next) + + assert result == mock_response + mock_call_next.assert_called_once_with(mock_request) + + +@pytest.mark.asyncio +async def test_middleware_no_auth_at_all(middleware, mock_request): + """Without any credential source the middleware should still surface NoCredentialsError.""" + mock_request.cookies = {} + mock_request.headers = {} + mock_request.url = MagicMock() + mock_request.url.hostname = 'localhost' + mock_request.url.path = '/api/some/endpoint' + mock_call_next = AsyncMock(return_value=MagicMock(spec=Response)) + + result = await middleware(mock_request, mock_call_next) + + assert isinstance(result, JSONResponse) + assert result.status_code == status.HTTP_401_UNAUTHORIZED + + def decode_body(body: bytes | memoryview): if isinstance(body, memoryview): return body.tobytes().decode() diff --git a/enterprise/tests/unit/test_saas_user_auth.py b/enterprise/tests/unit/test_saas_user_auth.py index 523b79629b..b23fc15d6e 100644 --- a/enterprise/tests/unit/test_saas_user_auth.py +++ b/enterprise/tests/unit/test_saas_user_auth.py @@ -662,9 +662,10 @@ async def test_saas_user_auth_from_bearer_success(): @pytest.mark.asyncio async def test_saas_user_auth_from_bearer_no_auth_header(): - """Test that saas_user_auth_from_bearer returns None if no auth header.""" + """Test that saas_user_auth_from_bearer returns None if no auth header or cookie.""" mock_request = MagicMock() mock_request.headers = {} + mock_request.cookies = {} result = await saas_user_auth_from_bearer(mock_request) @@ -842,10 +843,11 @@ def test_get_api_key_from_header_with_both_headers(): def test_get_api_key_from_header_with_no_headers(): - """Test that get_api_key_from_header returns None when no relevant headers are present.""" - # Create a mock request with no relevant headers + """Test that get_api_key_from_header returns None when no relevant headers or cookies are present.""" + # Create a mock request with no relevant headers or cookies mock_request = MagicMock(spec=Request) mock_request.headers = {'Other-Header': 'some_value'} + mock_request.cookies = {} # Call the function api_key = get_api_key_from_header(mock_request) @@ -859,6 +861,7 @@ def test_get_api_key_from_header_with_invalid_authorization_format(): # Create a mock request with incorrectly formatted Authorization header mock_request = MagicMock(spec=Request) mock_request.headers = {'Authorization': 'InvalidFormat api_key'} + mock_request.cookies = {} # Call the function api_key = get_api_key_from_header(mock_request) @@ -979,6 +982,134 @@ def test_get_api_key_from_header_bearer_with_empty_token(): assert api_key == '' +def test_get_api_key_from_header_with_api_key_cookie(): + """Test that get_api_key_from_header extracts API key from the api_key cookie.""" + # Create a mock request with the api_key cookie set + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + mock_request.cookies = {'api_key': 'cookie_api_key'} + + # Call the function + api_key = get_api_key_from_header(mock_request) + + # Assert that the API key from the cookie was correctly extracted + assert api_key == 'cookie_api_key' + + +def test_get_api_key_from_header_priority_authorization_over_api_key_cookie(): + """Test that the Authorization header takes priority over the api_key cookie.""" + mock_request = MagicMock(spec=Request) + mock_request.headers = {'Authorization': 'Bearer auth_api_key'} + mock_request.cookies = {'api_key': 'cookie_api_key'} + + api_key = get_api_key_from_header(mock_request) + + # The Authorization header value should win. + assert api_key == 'auth_api_key' + + +def test_get_api_key_from_header_priority_x_session_over_api_key_cookie(): + """Test that the X-Session-API-Key header takes priority over the api_key cookie.""" + mock_request = MagicMock(spec=Request) + mock_request.headers = {'X-Session-API-Key': 'session_api_key'} + mock_request.cookies = {'api_key': 'cookie_api_key'} + + api_key = get_api_key_from_header(mock_request) + + # The X-Session-API-Key header value should win. + assert api_key == 'session_api_key' + + +def test_get_api_key_from_header_priority_x_access_token_over_api_key_cookie(): + """Test that the X-Access-Token header takes priority over the api_key cookie.""" + mock_request = MagicMock(spec=Request) + mock_request.headers = {'X-Access-Token': 'access_token_key'} + mock_request.cookies = {'api_key': 'cookie_api_key'} + + api_key = get_api_key_from_header(mock_request) + + # The X-Access-Token header value should win over the cookie. + assert api_key == 'access_token_key' + + +def test_get_api_key_from_header_with_empty_api_key_cookie(): + """An empty api_key cookie value should be treated as absent and fall through.""" + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + mock_request.cookies = {'api_key': ''} + + api_key = get_api_key_from_header(mock_request) + + # Empty cookie value yields an empty string (mirrors the empty-header + # behaviour). Callers like `saas_user_auth_from_bearer` treat falsy + # values as missing credentials. + assert api_key == '' + + +def test_get_api_key_from_header_with_unrelated_cookies(): + """Unrelated cookies must not be picked up as an API key.""" + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + mock_request.cookies = {'session': 'abc123', 'preferences': 'dark'} + + api_key = get_api_key_from_header(mock_request) + + assert api_key is None + + +@pytest.mark.asyncio +async def test_saas_user_auth_from_bearer_via_api_key_cookie(): + """A valid api_key cookie should authenticate the same as an X-Access-Token header.""" + mock_request = MagicMock() + mock_request.headers = {} + mock_request.cookies = {'api_key': 'cookie_api_key'} + + mock_org_id = uuid.uuid4() + mock_validation_result = ApiKeyValidationResult( + user_id='test_user_id', + org_id=mock_org_id, + key_id=7, + key_name='Cookie Key', + ) + + with patch('server.auth.saas_user_auth.ApiKeyStore') as mock_api_key_store_cls: + mock_api_key_store = MagicMock() + mock_api_key_store.validate_api_key = AsyncMock( + return_value=mock_validation_result + ) + mock_api_key_store_cls.get_instance.return_value = mock_api_key_store + + result = await saas_user_auth_from_bearer(mock_request) + + assert isinstance(result, SaasUserAuth) + assert result.user_id == 'test_user_id' + assert result.api_key_org_id == mock_org_id + assert result.api_key_id == 7 + assert result.api_key_name == 'Cookie Key' + assert result.auth_type == AuthType.BEARER + mock_api_key_store.validate_api_key.assert_called_once_with('cookie_api_key') + + +@pytest.mark.asyncio +async def test_saas_user_auth_from_bearer_via_api_key_cookie_invalid(): + """An api_key cookie with an invalid key should produce no auth, same as a bad header.""" + mock_request = MagicMock() + mock_request.headers = {} + mock_request.cookies = {'api_key': 'invalid_cookie_key'} + + with patch('server.auth.saas_user_auth.ApiKeyStore') as mock_api_key_store_cls: + mock_api_key_store = MagicMock() + mock_api_key_store.validate_api_key = AsyncMock(return_value=None) + mock_api_key_store_cls.get_instance.return_value = mock_api_key_store + + result = await saas_user_auth_from_bearer(mock_request) + + assert result is None + mock_api_key_store.validate_api_key.assert_called_once_with( + 'invalid_cookie_key' + ) + + @pytest.mark.asyncio async def test_saas_user_auth_from_signed_token_blocked_domain(mock_config): """Test that saas_user_auth_from_signed_token raises AuthError when email domain is blocked."""