From a42cbf65e46204fbd259eee9bb3f4093c0a61125 Mon Sep 17 00:00:00 2001 From: yaojin3616 Date: Thu, 6 Aug 2026 14:24:30 +0800 Subject: [PATCH] fix(enterprise): enforce multi-tenant isolation for EnterpriseInfo and Agent sync --- .../v1_0_0_f061_enterprise_info_tenant_id.py | 50 +++++++ backend/app/api/enterprise.py | 21 ++- backend/app/models/audit.py | 9 +- backend/app/services/enterprise_sync.py | 38 +++-- .../test_enterprise_info_tenant_isolation.py | 141 ++++++++++++++++++ 5 files changed, 239 insertions(+), 20 deletions(-) create mode 100644 backend/alembic/versions/v1_0_0_f061_enterprise_info_tenant_id.py create mode 100644 backend/tests/test_enterprise_info_tenant_isolation.py diff --git a/backend/alembic/versions/v1_0_0_f061_enterprise_info_tenant_id.py b/backend/alembic/versions/v1_0_0_f061_enterprise_info_tenant_id.py new file mode 100644 index 000000000..e2f2b48fe --- /dev/null +++ b/backend/alembic/versions/v1_0_0_f061_enterprise_info_tenant_id.py @@ -0,0 +1,50 @@ +"""Add tenant_id and composite unique constraint to enterprise_info. + +Background: + EnterpriseInfo currently lacks a tenant_id column, causing multi-tenant data bleed where an update + from one tenant administrator overwrote global EnterpriseInfo entries and pushed synced files to all running agents across tenants. + +Scope: + Add tenant_id UUID column (indexed) to enterprise_info. + Drop legacy single info_type unique constraint. + Add composite unique constraint uq_enterprise_info_tenant_type on (tenant_id, info_type). + +Idempotence: + Safe for retry. Pure DDL migration without blocking data locks. + +Revision ID: f061_enterprise_info_tenant_id +Revises: f060_tenant_id_backfill +Create Date: 2026-08-06 14:00:00.000000 + +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision: str = "f061_enterprise_info_tenant_id" +down_revision: Union[str, None] = "f060_tenant_id_backfill" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # 1. Add tenant_id column with default uuid generator or nullable first if populated + op.add_column("enterprise_info", sa.Column("tenant_id", postgresql.UUID(as_uuid=True), nullable=True)) + op.create_index(op.f("ix_enterprise_info_tenant_id"), "enterprise_info", ["tenant_id"], unique=False) + + # 2. Drop legacy single info_type unique constraint + op.drop_constraint("enterprise_info_info_type_key", "enterprise_info", type_="unique") + + # 3. Create new composite unique constraint (tenant_id, info_type) + op.create_unique_constraint("uq_enterprise_info_tenant_type", "enterprise_info", ["tenant_id", "info_type"]) + + +def downgrade() -> None: + op.drop_constraint("uq_enterprise_info_tenant_type", "enterprise_info", type_="unique") + op.create_unique_constraint("enterprise_info_info_type_key", "enterprise_info", ["info_type"]) + op.drop_index(op.f("ix_enterprise_info_tenant_id"), table_name="enterprise_info") + op.drop_column("enterprise_info", "tenant_id") diff --git a/backend/app/api/enterprise.py b/backend/app/api/enterprise.py index 2c4892db5..a3cfb911b 100644 --- a/backend/app/api/enterprise.py +++ b/backend/app/api/enterprise.py @@ -585,8 +585,14 @@ async def list_enterprise_info( current_user: User = Depends(get_current_user), db: Any = None, ): - """List all enterprise information entries.""" - result = await db.execute(select(EnterpriseInfo).order_by(EnterpriseInfo.info_type)) + """List enterprise information entries for current tenant.""" + if not current_user.tenant_id: + return [] + result = await db.execute( + select(EnterpriseInfo) + .where(EnterpriseInfo.tenant_id == current_user.tenant_id) + .order_by(EnterpriseInfo.info_type) + ) return [EnterpriseInfoOut.model_validate(e) for e in result.scalars().all()] @@ -597,12 +603,15 @@ async def update_enterprise_info( current_user: User = Depends(get_current_admin), db: Any = None, ): - """Create or update enterprise information. Triggers sync to agents.""" + """Create or update enterprise information for current tenant. Triggers sync to tenant agents.""" + if not current_user.tenant_id: + raise HTTPException(status_code=403, detail="User must belong to a tenant") + info = await enterprise_sync_service.update_enterprise_info( - db, info_type, data.content, data.visible_roles, current_user.id + db, current_user.tenant_id, info_type, data.content, data.visible_roles, current_user.id ) - # Sync to all running agents - await enterprise_sync_service.sync_to_all_agents(db) + # Sync only to running agents in the current tenant + await enterprise_sync_service.sync_to_all_agents(db, tenant_id=current_user.tenant_id) return EnterpriseInfoOut.model_validate(info) diff --git a/backend/app/models/audit.py b/backend/app/models/audit.py index 71181a408..1ebaf296e 100644 --- a/backend/app/models/audit.py +++ b/backend/app/models/audit.py @@ -3,7 +3,7 @@ import uuid from datetime import datetime -from sqlalchemy import DateTime, Enum, ForeignKey, Integer, String, Text, func, text +from sqlalchemy import DateTime, Enum, ForeignKey, Integer, String, Text, UniqueConstraint, func, text from sqlalchemy.dialects.postgresql import JSON, JSONB, UUID from sqlalchemy.orm import Mapped, mapped_column @@ -83,9 +83,14 @@ class EnterpriseInfo(Base): """Centralized enterprise information with versioning for sync.""" __tablename__ = "enterprise_info" + __tenant_scoped__ = True + __table_args__ = ( + UniqueConstraint("tenant_id", "info_type", name="uq_enterprise_info_tenant_type"), + ) id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) - info_type: Mapped[str] = mapped_column(String(50), nullable=False, unique=True) # org_structure, company_profile, etc. + tenant_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), nullable=False, index=True) + info_type: Mapped[str] = mapped_column(String(50), nullable=False) # org_structure, company_profile, etc. content: Mapped[dict] = mapped_column(JSON, nullable=False) version: Mapped[int] = mapped_column(Integer, default=1, nullable=False) visible_roles: Mapped[list] = mapped_column(JSON, default=[]) # Which agent roles can see this diff --git a/backend/app/services/enterprise_sync.py b/backend/app/services/enterprise_sync.py index 87f8cc9d6..a094181c6 100644 --- a/backend/app/services/enterprise_sync.py +++ b/backend/app/services/enterprise_sync.py @@ -22,15 +22,18 @@ class EnterpriseSyncService: - """Synchronize enterprise information to all online Agent containers.""" + """Synchronize enterprise information to online Agent containers within tenant scope.""" async def update_enterprise_info( - self, db: AsyncSession, info_type: str, content: dict, + self, db: AsyncSession, tenant_id: uuid.UUID, info_type: str, content: dict, visible_roles: list[str], updated_by: uuid.UUID ) -> EnterpriseInfo: - """Update enterprise info in database and notify all agents.""" + """Update enterprise info in database for a specific tenant and notify tenant agents.""" result = await query_dao.execute(db, - select(EnterpriseInfo).where(EnterpriseInfo.info_type == info_type) + select(EnterpriseInfo).where( + EnterpriseInfo.tenant_id == tenant_id, + EnterpriseInfo.info_type == info_type, + ) ) info = result.scalar_one_or_none() @@ -41,6 +44,7 @@ async def update_enterprise_info( info.updated_by = updated_by else: info = EnterpriseInfo( + tenant_id=tenant_id, info_type=info_type, content=content, visible_roles=visible_roles, @@ -50,22 +54,31 @@ async def update_enterprise_info( await query_dao.flush(db) - # Publish update event + # Publish update event with tenant_id scope await publish_event(ENTERPRISE_INFO_CHANNEL, { + "tenant_id": str(tenant_id), "info_type": info_type, "version": info.version, "visible_roles": visible_roles, }) - logger.info(f"Published enterprise_info update: {info_type} v{info.version}") + logger.info(f"Published enterprise_info update for tenant {tenant_id}: {info_type} v{info.version}") return info async def sync_to_agent(self, db: AsyncSession, agent_id: uuid.UUID, agent_role: str = "") -> None: """Pull enterprise info from DB and write to agent's enterprise_info/ directory. - Filters by visible_roles — if empty, all roles can see it. + Strictly filters EnterpriseInfo entries by the agent's tenant_id and role. """ - result = await query_dao.execute(db, select(EnterpriseInfo)) + agent_result = await query_dao.execute(db, select(Agent).where(Agent.id == agent_id)) + agent = agent_result.scalar_one_or_none() + if not agent or not agent.tenant_id: + logger.warning(f"Skipping enterprise_info sync for invalid agent {agent_id}") + return + + result = await query_dao.execute( + db, select(EnterpriseInfo).where(EnterpriseInfo.tenant_id == agent.tenant_id) + ) all_info = result.scalars().all() for info in all_info: @@ -84,13 +97,14 @@ async def sync_to_agent(self, db: AsyncSession, agent_id: uuid.UUID, agent_role: content_type="application/json", ) - logger.info(f"Synced enterprise info to agent {agent_id}") + logger.info(f"Synced tenant {agent.tenant_id} enterprise info to agent {agent_id}") - async def sync_to_all_agents(self, db: AsyncSession) -> int: - """Sync enterprise info to all running agents. Returns count.""" + async def sync_to_all_agents(self, db: AsyncSession, tenant_id: uuid.UUID) -> int: + """Sync enterprise info to running agents strictly belonging to the given tenant. Returns count.""" result = await query_dao.execute( db, select(Agent).where( + Agent.tenant_id == tenant_id, Agent.status == "running", Agent.deleted_at.is_(None), ) @@ -100,7 +114,7 @@ async def sync_to_all_agents(self, db: AsyncSession) -> int: for agent in agents: await self.sync_to_agent(db, agent.id, agent.role_description) - logger.info(f"Synced enterprise info to {len(agents)} agents") + logger.info(f"Synced enterprise info to {len(agents)} agents in tenant {tenant_id}") return len(agents) diff --git a/backend/tests/test_enterprise_info_tenant_isolation.py b/backend/tests/test_enterprise_info_tenant_isolation.py new file mode 100644 index 000000000..896581ef5 --- /dev/null +++ b/backend/tests/test_enterprise_info_tenant_isolation.py @@ -0,0 +1,141 @@ +"""Unit tests verifying multi-tenant isolation for EnterpriseInfo updates and agent file sync.""" + +import json +import uuid +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest +from fastapi import HTTPException + +from app.api import enterprise as enterprise_api +from app.models.agent import Agent +from app.models.audit import EnterpriseInfo +from app.models.user import User +from app.schemas.schemas import EnterpriseInfoUpdate +from app.services.enterprise_sync import enterprise_sync_service + + +class _MockResult: + def __init__(self, items: list) -> None: + self._items = items + + def scalar_one_or_none(self): + return self._items[0] if self._items else None + + def scalars(self): + return self + + def all(self): + return self._items + + +class _MockSession: + def __init__(self) -> None: + self.added = [] + self.flushed = False + + async def execute(self, statement): + return _MockResult([]) + + def add(self, item): + self.added.append(item) + + async def flush(self): + self.flushed = True + + +@pytest.mark.asyncio +async def test_update_enterprise_info_binds_to_current_user_tenant(): + """EnterpriseInfo creation must bind tenant_id to the active user's tenant.""" + tenant_a = uuid.uuid4() + user_a = User(id=uuid.uuid4(), tenant_id=tenant_a, role="org_admin") + db = _MockSession() + + stored_info = None + + async def mock_store(agent_id, path, content, content_type): + pass + + with patch("app.services.enterprise_sync.publish_event", AsyncMock()), \ + patch("app.services.enterprise_sync.store_agent_bytes", mock_store): + info = await enterprise_sync_service.update_enterprise_info( + db=db, + tenant_id=tenant_a, + info_type="company_profile", + content={"name": "Company A"}, + visible_roles=[], + updated_by=user_a.id, + ) + + assert info.tenant_id == tenant_a + assert info.info_type == "company_profile" + assert info.content == {"name": "Company A"} + + +@pytest.mark.asyncio +async def test_sync_to_all_agents_restricts_to_target_tenant(): + """Agent sync must only target running agents belonging to the specified tenant.""" + from datetime import datetime, timezone + + tenant_a = uuid.uuid4() + tenant_b = uuid.uuid4() + + agent_a = Agent(id=uuid.uuid4(), tenant_id=tenant_a, status="running", role_description="dev") + + now = datetime.now(timezone.utc) + info_a = EnterpriseInfo( + tenant_id=tenant_a, + info_type="company_profile", + content={"secret": "Tenant A Secret"}, + visible_roles=[], + version=1, + created_at=now, + updated_at=now, + ) + + synced_files = {} + + async def mock_store(agent_id, path, content, content_type): + synced_files[(agent_id, path)] = json.loads(content.decode("utf-8")) + + db = AsyncMock() + + async def mock_execute(stmt, *args, **kwargs): + sql = str(stmt) + if "FROM agents" in sql: + return _MockResult([agent_a]) + elif "FROM enterprise_info" in sql: + return _MockResult([info_a]) + return _MockResult([]) + + db.execute = AsyncMock(side_effect=mock_execute) + + with patch("app.services.enterprise_sync.store_agent_bytes", mock_store): + # Sync tenant A + count = await enterprise_sync_service.sync_to_all_agents(db, tenant_id=tenant_a) + + assert count == 1 + # Only Agent A receives Tenant A's secret + assert (agent_a.id, "enterprise_info/company_profile.json") in synced_files + assert synced_files[(agent_a.id, "enterprise_info/company_profile.json")]["content"] == {"secret": "Tenant A Secret"} + + +@pytest.mark.asyncio +async def test_api_list_enterprise_info_filters_by_tenant(): + """API list endpoint must only return EnterpriseInfo records for the current user's tenant.""" + from datetime import datetime, timezone + + tenant_a = uuid.uuid4() + user_a = User(id=uuid.uuid4(), tenant_id=tenant_a, role="member") + now = datetime.now(timezone.utc) + info_a = EnterpriseInfo(id=uuid.uuid4(), tenant_id=tenant_a, info_type="rules", content={"a": 1}, version=1, visible_roles=[], created_at=now, updated_at=now) + + db = AsyncMock() + db.execute = AsyncMock(return_value=_MockResult([info_a])) + + result = await enterprise_api.list_enterprise_info(current_user=user_a, db=db) + + assert len(result) == 1 + assert result[0].info_type == "rules" + assert result[0].content == {"a": 1}