fix(backend): validate API key org_id during authorization to prevent cross-org access (org project) (#13468)

2026-03-22 13:47:19 +08:00 · 2026-03-19 16:09:37 +07:00
parent 8039807c3f
commit e02dbb8974
7 changed files with 354 additions and 32 deletions
--- a/enterprise/tests/unit/test_api_key_store.py
+++ b/enterprise/tests/unit/test_api_key_store.py
@@ -5,7 +5,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
 import pytest
 from sqlalchemy import select
 from storage.api_key import ApiKey
-from storage.api_key_store import ApiKeyStore
+from storage.api_key_store import ApiKeyStore, ApiKeyValidationResult


@pytest.fixture
@@ -110,8 +110,8 @@ async def test_create_api_key(

@pytest.mark.asyncio
 async def test_validate_api_key_valid(api_key_store, async_session_maker):
-    """Test validating a valid API key."""
-    # Setup - create an API key in the database
+    """Test validating a valid API key returns user_id and org_id."""
+    # Arrange
    user_id = str(uuid.uuid4())
    org_id = uuid.uuid4()
    api_key_value = 'test-api-key'
@@ -128,11 +128,12 @@ async def test_validate_api_key_valid(api_key_store, async_session_maker):
        await session.commit()
        key_id = key_record.id

-    # Execute - patch a_session_maker to use test's async session maker
+    # Act
    with patch('storage.api_key_store.a_session_maker', async_session_maker):
        result = await api_key_store.validate_api_key(api_key_value)

-    # Verify - result is now ApiKeyValidationResult
+    # Assert
+    assert isinstance(result, ApiKeyValidationResult)
    assert result is not None
    assert result.user_id == user_id
    assert result.org_id == org_id
@@ -202,7 +203,7 @@ async def test_validate_api_key_valid_timezone_naive(
    api_key_store, async_session_maker
 ):
    """Test validating a valid API key with timezone-naive datetime from database."""
-    # Setup - create a valid API key with timezone-naive datetime (future date)
+    # Arrange
    user_id = str(uuid.uuid4())
    org_id = uuid.uuid4()
    api_key_value = 'test-valid-naive-key'
@@ -219,13 +220,44 @@ async def test_validate_api_key_valid_timezone_naive(
        session.add(key_record)
        await session.commit()

-    # Execute - patch a_session_maker to use test's async session maker
+    # Act
    with patch('storage.api_key_store.a_session_maker', async_session_maker):
        result = await api_key_store.validate_api_key(api_key_value)

-    # Verify - result is now ApiKeyValidationResult
+    # Assert
+    assert isinstance(result, ApiKeyValidationResult)
+    assert result.user_id == user_id
+    assert result.org_id == org_id
+
+
+@pytest.mark.asyncio
+async def test_validate_api_key_legacy_without_org_id(
+    api_key_store, async_session_maker
+):
+    """Test validating a legacy API key without org_id returns None for org_id."""
+    # Arrange
+    user_id = str(uuid.uuid4())
+    api_key_value = 'test-legacy-key-no-org'
+
+    async with async_session_maker() as session:
+        key_record = ApiKey(
+            key=api_key_value,
+            user_id=user_id,
+            org_id=None,  # Legacy key without org binding
+            name='Legacy Key',
+        )
+        session.add(key_record)
+        await session.commit()
+
+    # Act
+    with patch('storage.api_key_store.a_session_maker', async_session_maker):
+        result = await api_key_store.validate_api_key(api_key_value)
+
+    # Assert
+    assert isinstance(result, ApiKeyValidationResult)
    assert result is not None
    assert result.user_id == user_id
+    assert result.org_id is None


@pytest.mark.asyncio
--- a/enterprise/tests/unit/test_authorization.py
+++ b/enterprise/tests/unit/test_authorization.py
@@ -13,6 +13,7 @@ from server.auth.authorization import (
    ROLE_PERMISSIONS,
    Permission,
    RoleName,
+    get_api_key_org_id_from_request,
    get_role_permissions,
    get_user_org_role,
    has_permission,
@@ -444,6 +445,15 @@ class TestGetUserOrgRole:
 # =============================================================================


+def _create_mock_request(api_key_org_id=None):
+    """Helper to create a mock request with optional api_key_org_id."""
+    mock_request = MagicMock()
+    mock_user_auth = MagicMock()
+    mock_user_auth.get_api_key_org_id.return_value = api_key_org_id
+    mock_request.state.user_auth = mock_user_auth
+    return mock_request
+
+
 class TestRequirePermission:
    """Tests for require_permission dependency factory."""

@@ -456,6 +466,7 @@ class TestRequirePermission:
        """
        user_id = str(uuid4())
        org_id = uuid4()
+        mock_request = _create_mock_request()

        mock_role = MagicMock()
        mock_role.name = 'admin'
@@ -465,7 +476,9 @@ class TestRequirePermission:
            AsyncMock(return_value=mock_role),
        ):
            permission_checker = require_permission(Permission.VIEW_LLM_SETTINGS)
-            result = await permission_checker(org_id=org_id, user_id=user_id)
+            result = await permission_checker(
+                request=mock_request, org_id=org_id, user_id=user_id
+            )
            assert result == user_id

    @pytest.mark.asyncio
@@ -476,10 +489,11 @@ class TestRequirePermission:
        THEN: 401 Unauthorized is raised
        """
        org_id = uuid4()
+        mock_request = _create_mock_request()

        permission_checker = require_permission(Permission.VIEW_LLM_SETTINGS)
        with pytest.raises(HTTPException) as exc_info:
-            await permission_checker(org_id=org_id, user_id=None)
+            await permission_checker(request=mock_request, org_id=org_id, user_id=None)

        assert exc_info.value.status_code == 401
        assert 'not authenticated' in exc_info.value.detail.lower()
@@ -493,6 +507,7 @@ class TestRequirePermission:
        """
        user_id = str(uuid4())
        org_id = uuid4()
+        mock_request = _create_mock_request()

        with patch(
            'server.auth.authorization.get_user_org_role',
@@ -500,7 +515,9 @@ class TestRequirePermission:
        ):
            permission_checker = require_permission(Permission.VIEW_LLM_SETTINGS)
            with pytest.raises(HTTPException) as exc_info:
-                await permission_checker(org_id=org_id, user_id=user_id)
+                await permission_checker(
+                    request=mock_request, org_id=org_id, user_id=user_id
+                )

            assert exc_info.value.status_code == 403
            assert 'not a member' in exc_info.value.detail.lower()
@@ -514,6 +531,7 @@ class TestRequirePermission:
        """
        user_id = str(uuid4())
        org_id = uuid4()
+        mock_request = _create_mock_request()

        mock_role = MagicMock()
        mock_role.name = 'member'
@@ -524,7 +542,9 @@ class TestRequirePermission:
        ):
            permission_checker = require_permission(Permission.DELETE_ORGANIZATION)
            with pytest.raises(HTTPException) as exc_info:
-                await permission_checker(org_id=org_id, user_id=user_id)
+                await permission_checker(
+                    request=mock_request, org_id=org_id, user_id=user_id
+                )

            assert exc_info.value.status_code == 403
            assert 'delete_organization' in exc_info.value.detail.lower()
@@ -538,6 +558,7 @@ class TestRequirePermission:
        """
        user_id = str(uuid4())
        org_id = uuid4()
+        mock_request = _create_mock_request()

        mock_role = MagicMock()
        mock_role.name = 'owner'
@@ -547,7 +568,9 @@ class TestRequirePermission:
            AsyncMock(return_value=mock_role),
        ):
            permission_checker = require_permission(Permission.DELETE_ORGANIZATION)
-            result = await permission_checker(org_id=org_id, user_id=user_id)
+            result = await permission_checker(
+                request=mock_request, org_id=org_id, user_id=user_id
+            )
            assert result == user_id

    @pytest.mark.asyncio
@@ -559,6 +582,7 @@ class TestRequirePermission:
        """
        user_id = str(uuid4())
        org_id = uuid4()
+        mock_request = _create_mock_request()

        mock_role = MagicMock()
        mock_role.name = 'admin'
@@ -569,7 +593,9 @@ class TestRequirePermission:
        ):
            permission_checker = require_permission(Permission.DELETE_ORGANIZATION)
            with pytest.raises(HTTPException) as exc_info:
-                await permission_checker(org_id=org_id, user_id=user_id)
+                await permission_checker(
+                    request=mock_request, org_id=org_id, user_id=user_id
+                )

            assert exc_info.value.status_code == 403

@@ -582,6 +608,7 @@ class TestRequirePermission:
        """
        user_id = str(uuid4())
        org_id = uuid4()
+        mock_request = _create_mock_request()

        mock_role = MagicMock()
        mock_role.name = 'member'
@@ -595,7 +622,9 @@ class TestRequirePermission:
        ):
            permission_checker = require_permission(Permission.DELETE_ORGANIZATION)
            with pytest.raises(HTTPException):
-                await permission_checker(org_id=org_id, user_id=user_id)
+                await permission_checker(
+                    request=mock_request, org_id=org_id, user_id=user_id
+                )

            mock_logger.warning.assert_called()
            call_args = mock_logger.warning.call_args
@@ -611,6 +640,7 @@ class TestRequirePermission:
        THEN: User ID is returned
        """
        user_id = str(uuid4())
+        mock_request = _create_mock_request()

        mock_role = MagicMock()
        mock_role.name = 'admin'
@@ -620,7 +650,9 @@ class TestRequirePermission:
            AsyncMock(return_value=mock_role),
        ) as mock_get_role:
            permission_checker = require_permission(Permission.VIEW_LLM_SETTINGS)
-            result = await permission_checker(org_id=None, user_id=user_id)
+            result = await permission_checker(
+                request=mock_request, org_id=None, user_id=user_id
+            )
            assert result == user_id
            mock_get_role.assert_called_once_with(user_id, None)

@@ -632,6 +664,7 @@ class TestRequirePermission:
        THEN: HTTPException with 403 status is raised
        """
        user_id = str(uuid4())
+        mock_request = _create_mock_request()

        with patch(
            'server.auth.authorization.get_user_org_role',
@@ -639,7 +672,9 @@ class TestRequirePermission:
        ):
            permission_checker = require_permission(Permission.VIEW_LLM_SETTINGS)
            with pytest.raises(HTTPException) as exc_info:
-                await permission_checker(org_id=None, user_id=user_id)
+                await permission_checker(
+                    request=mock_request, org_id=None, user_id=user_id
+                )

            assert exc_info.value.status_code == 403
            assert 'not a member' in exc_info.value.detail
@@ -662,6 +697,7 @@ class TestPermissionScenarios:
        """
        user_id = str(uuid4())
        org_id = uuid4()
+        mock_request = _create_mock_request()

        mock_role = MagicMock()
        mock_role.name = 'member'
@@ -671,7 +707,9 @@ class TestPermissionScenarios:
            AsyncMock(return_value=mock_role),
        ):
            permission_checker = require_permission(Permission.MANAGE_SECRETS)
-            result = await permission_checker(org_id=org_id, user_id=user_id)
+            result = await permission_checker(
+                request=mock_request, org_id=org_id, user_id=user_id
+            )
            assert result == user_id

    @pytest.mark.asyncio
@@ -683,6 +721,7 @@ class TestPermissionScenarios:
        """
        user_id = str(uuid4())
        org_id = uuid4()
+        mock_request = _create_mock_request()

        mock_role = MagicMock()
        mock_role.name = 'member'
@@ -695,7 +734,9 @@ class TestPermissionScenarios:
                Permission.INVITE_USER_TO_ORGANIZATION
            )
            with pytest.raises(HTTPException) as exc_info:
-                await permission_checker(org_id=org_id, user_id=user_id)
+                await permission_checker(
+                    request=mock_request, org_id=org_id, user_id=user_id
+                )

            assert exc_info.value.status_code == 403

@@ -708,6 +749,7 @@ class TestPermissionScenarios:
        """
        user_id = str(uuid4())
        org_id = uuid4()
+        mock_request = _create_mock_request()

        mock_role = MagicMock()
        mock_role.name = 'admin'
@@ -719,7 +761,9 @@ class TestPermissionScenarios:
            permission_checker = require_permission(
                Permission.INVITE_USER_TO_ORGANIZATION
            )
-            result = await permission_checker(org_id=org_id, user_id=user_id)
+            result = await permission_checker(
+                request=mock_request, org_id=org_id, user_id=user_id
+            )
            assert result == user_id

    @pytest.mark.asyncio
@@ -731,6 +775,7 @@ class TestPermissionScenarios:
        """
        user_id = str(uuid4())
        org_id = uuid4()
+        mock_request = _create_mock_request()

        mock_role = MagicMock()
        mock_role.name = 'admin'
@@ -741,7 +786,9 @@ class TestPermissionScenarios:
        ):
            permission_checker = require_permission(Permission.CHANGE_USER_ROLE_OWNER)
            with pytest.raises(HTTPException) as exc_info:
-                await permission_checker(org_id=org_id, user_id=user_id)
+                await permission_checker(
+                    request=mock_request, org_id=org_id, user_id=user_id
+                )

            assert exc_info.value.status_code == 403

@@ -754,6 +801,7 @@ class TestPermissionScenarios:
        """
        user_id = str(uuid4())
        org_id = uuid4()
+        mock_request = _create_mock_request()

        mock_role = MagicMock()
        mock_role.name = 'owner'
@@ -763,5 +811,200 @@ class TestPermissionScenarios:
            AsyncMock(return_value=mock_role),
        ):
            permission_checker = require_permission(Permission.CHANGE_USER_ROLE_OWNER)
-            result = await permission_checker(org_id=org_id, user_id=user_id)
+            result = await permission_checker(
+                request=mock_request, org_id=org_id, user_id=user_id
+            )
            assert result == user_id
+
+
+# =============================================================================
+# Tests for API key organization validation
+# =============================================================================
+
+
+class TestApiKeyOrgValidation:
+    """Tests for API key organization binding validation in require_permission."""
+
+    @pytest.mark.asyncio
+    async def test_allows_access_when_api_key_org_matches_target_org(self):
+        """
+        GIVEN: API key with org_id that matches the target org_id in the request
+        WHEN: Permission checker is called
+        THEN: User ID is returned (access allowed)
+        """
+        # Arrange
+        user_id = str(uuid4())
+        org_id = uuid4()
+        mock_request = _create_mock_request(api_key_org_id=org_id)
+
+        mock_role = MagicMock()
+        mock_role.name = 'admin'
+
+        # Act & Assert
+        with patch(
+            'server.auth.authorization.get_user_org_role',
+            AsyncMock(return_value=mock_role),
+        ):
+            permission_checker = require_permission(Permission.VIEW_LLM_SETTINGS)
+            result = await permission_checker(
+                request=mock_request, org_id=org_id, user_id=user_id
+            )
+            assert result == user_id
+
+    @pytest.mark.asyncio
+    async def test_denies_access_when_api_key_org_mismatches_target_org(self):
+        """
+        GIVEN: API key created for Org A, but user tries to access Org B
+        WHEN: Permission checker is called
+        THEN: 403 Forbidden is raised with org mismatch message
+        """
+        # Arrange
+        user_id = str(uuid4())
+        api_key_org_id = uuid4()  # Org A - where API key was created
+        target_org_id = uuid4()  # Org B - where user is trying to access
+        mock_request = _create_mock_request(api_key_org_id=api_key_org_id)
+
+        # Act & Assert
+        permission_checker = require_permission(Permission.VIEW_LLM_SETTINGS)
+        with pytest.raises(HTTPException) as exc_info:
+            await permission_checker(
+                request=mock_request, org_id=target_org_id, user_id=user_id
+            )
+
+        assert exc_info.value.status_code == 403
+        assert (
+            'API key is not authorized for this organization' in exc_info.value.detail
+        )
+
+    @pytest.mark.asyncio
+    async def test_allows_access_for_legacy_api_key_without_org_binding(self):
+        """
+        GIVEN: Legacy API key without org_id binding (org_id is None)
+        WHEN: Permission checker is called
+        THEN: Falls through to normal permission check (backward compatible)
+        """
+        # Arrange
+        user_id = str(uuid4())
+        org_id = uuid4()
+        mock_request = _create_mock_request(api_key_org_id=None)
+
+        mock_role = MagicMock()
+        mock_role.name = 'admin'
+
+        # Act & Assert
+        with patch(
+            'server.auth.authorization.get_user_org_role',
+            AsyncMock(return_value=mock_role),
+        ):
+            permission_checker = require_permission(Permission.VIEW_LLM_SETTINGS)
+            result = await permission_checker(
+                request=mock_request, org_id=org_id, user_id=user_id
+            )
+            assert result == user_id
+
+    @pytest.mark.asyncio
+    async def test_allows_access_for_cookie_auth_without_api_key_org_id(self):
+        """
+        GIVEN: Cookie-based authentication (no api_key_org_id in user_auth)
+        WHEN: Permission checker is called
+        THEN: Falls through to normal permission check
+        """
+        # Arrange
+        user_id = str(uuid4())
+        org_id = uuid4()
+        mock_request = _create_mock_request(api_key_org_id=None)
+
+        mock_role = MagicMock()
+        mock_role.name = 'admin'
+
+        # Act & Assert
+        with patch(
+            'server.auth.authorization.get_user_org_role',
+            AsyncMock(return_value=mock_role),
+        ):
+            permission_checker = require_permission(Permission.VIEW_LLM_SETTINGS)
+            result = await permission_checker(
+                request=mock_request, org_id=org_id, user_id=user_id
+            )
+            assert result == user_id
+
+    @pytest.mark.asyncio
+    async def test_logs_warning_on_api_key_org_mismatch(self):
+        """
+        GIVEN: API key org_id doesn't match target org_id
+        WHEN: Permission checker is called
+        THEN: Warning is logged with org mismatch details
+        """
+        # Arrange
+        user_id = str(uuid4())
+        api_key_org_id = uuid4()
+        target_org_id = uuid4()
+        mock_request = _create_mock_request(api_key_org_id=api_key_org_id)
+
+        # Act & Assert
+        with patch('server.auth.authorization.logger') as mock_logger:
+            permission_checker = require_permission(Permission.VIEW_LLM_SETTINGS)
+            with pytest.raises(HTTPException):
+                await permission_checker(
+                    request=mock_request, org_id=target_org_id, user_id=user_id
+                )
+
+            mock_logger.warning.assert_called()
+            call_args = mock_logger.warning.call_args
+            assert call_args[1]['extra']['user_id'] == user_id
+            assert call_args[1]['extra']['api_key_org_id'] == str(api_key_org_id)
+            assert call_args[1]['extra']['target_org_id'] == str(target_org_id)
+
+
+class TestGetApiKeyOrgIdFromRequest:
+    """Tests for get_api_key_org_id_from_request helper function."""
+
+    @pytest.mark.asyncio
+    async def test_returns_org_id_when_user_auth_has_api_key_org_id(self):
+        """
+        GIVEN: Request with user_auth that has api_key_org_id
+        WHEN: get_api_key_org_id_from_request is called
+        THEN: Returns the api_key_org_id
+        """
+        # Arrange
+        org_id = uuid4()
+        mock_request = _create_mock_request(api_key_org_id=org_id)
+
+        # Act
+        result = await get_api_key_org_id_from_request(mock_request)
+
+        # Assert
+        assert result == org_id
+
+    @pytest.mark.asyncio
+    async def test_returns_none_when_user_auth_has_no_api_key_org_id(self):
+        """
+        GIVEN: Request with user_auth that has no api_key_org_id (cookie auth)
+        WHEN: get_api_key_org_id_from_request is called
+        THEN: Returns None
+        """
+        # Arrange
+        mock_request = _create_mock_request(api_key_org_id=None)
+
+        # Act
+        result = await get_api_key_org_id_from_request(mock_request)
+
+        # Assert
+        assert result is None
+
+    @pytest.mark.asyncio
+    async def test_returns_none_when_no_user_auth_in_request(self):
+        """
+        GIVEN: Request without user_auth in state
+        WHEN: get_api_key_org_id_from_request is called
+        THEN: Returns None
+        """
+        # Arrange
+        mock_request = MagicMock()
+        mock_request.state.user_auth = None
+
+        # Act
+        result = await get_api_key_org_id_from_request(mock_request)
+
+        # Assert
+        assert result is None
--- a/enterprise/tests/unit/test_saas_user_auth.py
+++ b/enterprise/tests/unit/test_saas_user_auth.py
@@ -459,7 +459,8 @@ async def test_get_instance_no_auth(mock_request):

@pytest.mark.asyncio
 async def test_saas_user_auth_from_bearer_success():
-    """Test successful authentication from bearer token."""
+    """Test successful authentication from bearer token sets user_id and api_key_org_id."""
+    # Arrange
    mock_request = MagicMock()
    mock_request.headers = {'Authorization': 'Bearer test_api_key'}