Spaces:
Running
Running
| from __future__ import annotations | |
| import asyncio | |
| import hashlib | |
| import hmac | |
| from pathlib import Path | |
| from typing import Any | |
| from sqlalchemy import select | |
| from sqlalchemy.exc import IntegrityError | |
| from app.security.database import SecurityDatabase | |
| from app.security.models import CanonicalMediaAsset | |
| class CanonicalAssetNotFoundError(Exception): | |
| """A requested output was never issued to the caller's workspace.""" | |
| class CanonicalAssetService: | |
| """Persists workspace ownership for MediaRouter-produced files. | |
| The database stores an immutable, validated locator rather than a client | |
| filesystem path. File resolution remains the responsibility of | |
| ``CleanupService`` so every consumer receives the same traversal checks. | |
| """ | |
| def __init__(self, database: SecurityDatabase) -> None: | |
| self.database = database | |
| async def register_output( | |
| self, | |
| *, | |
| workspace_id: str, | |
| user_id: str | None, | |
| request_id: str, | |
| path: Path, | |
| mime_type: str, | |
| metadata: dict[str, Any] | None = None, | |
| project_id: str | None = None, | |
| ) -> CanonicalMediaAsset: | |
| if path.name != str(path.name) or not path.is_file(): | |
| raise CanonicalAssetNotFoundError("Generated output is unavailable.") | |
| digest = await asyncio.to_thread(self._sha256, path) | |
| record = CanonicalMediaAsset( | |
| workspace_id=workspace_id, | |
| request_id=request_id, | |
| filename=path.name, | |
| mime_type=mime_type, | |
| file_size=path.stat().st_size, | |
| sha256=digest, | |
| metadata_json=dict(metadata or {}), | |
| created_by_user_id=user_id, | |
| project_id=project_id, | |
| ) | |
| try: | |
| async with self.database.session() as session: | |
| session.add(record) | |
| await session.commit() | |
| await session.refresh(record) | |
| return record | |
| except IntegrityError: | |
| async with self.database.session() as session: | |
| existing = await session.scalar( | |
| select(CanonicalMediaAsset).where( | |
| CanonicalMediaAsset.request_id == request_id, | |
| CanonicalMediaAsset.filename == path.name, | |
| ) | |
| ) | |
| if existing is None: | |
| raise | |
| # Output IDs are globally unique. A second workspace must | |
| # never be allowed to claim the same path after a race. | |
| if ( | |
| existing.workspace_id != workspace_id | |
| or existing.project_id != project_id | |
| or existing.mime_type != mime_type | |
| ): | |
| raise CanonicalAssetNotFoundError( | |
| "Generated output is not owned by this workspace." | |
| ) | |
| await self.verify_file(existing, path) | |
| return existing | |
| async def discard_output( | |
| self, | |
| *, | |
| workspace_id: str, | |
| asset_id: str, | |
| request_id: str, | |
| filename: str, | |
| ) -> bool: | |
| """Remove a just-created canonical output after cancellation wins. | |
| Immutable locator fields must all match so this internal compensation | |
| cannot delete an unrelated asset selected only by an opaque ID. | |
| """ | |
| async with self.database.session() as session: | |
| record = await session.scalar( | |
| select(CanonicalMediaAsset) | |
| .where( | |
| CanonicalMediaAsset.id == asset_id, | |
| CanonicalMediaAsset.workspace_id == workspace_id, | |
| CanonicalMediaAsset.request_id == request_id, | |
| CanonicalMediaAsset.filename == filename, | |
| ) | |
| .with_for_update() | |
| ) | |
| if record is None: | |
| return False | |
| await session.delete(record) | |
| await session.commit() | |
| return True | |
| async def get_owned( | |
| self, *, workspace_id: str, request_id: str, filename: str | |
| ) -> CanonicalMediaAsset: | |
| async with self.database.session() as session: | |
| record = await session.scalar( | |
| select(CanonicalMediaAsset).where( | |
| CanonicalMediaAsset.workspace_id == workspace_id, | |
| CanonicalMediaAsset.request_id == request_id, | |
| CanonicalMediaAsset.filename == filename, | |
| ) | |
| ) | |
| if record is None: | |
| raise CanonicalAssetNotFoundError("Media asset was not found in this workspace.") | |
| return record | |
| async def get_owned_by_id( | |
| self, *, workspace_id: str, user_id: str, asset_id: str | |
| ) -> CanonicalMediaAsset: | |
| """Resolve a canonical asset reference without accepting a path. | |
| Generation (and future first-party services) receive only the opaque | |
| canonical asset ID. The workspace predicate remains mandatory even | |
| though the table is also protected by PostgreSQL RLS. | |
| """ | |
| async with self.database.tenant_session( | |
| workspace_id=workspace_id, user_id=user_id | |
| ) as session: | |
| record = await session.scalar( | |
| select(CanonicalMediaAsset).where( | |
| CanonicalMediaAsset.id == asset_id, | |
| CanonicalMediaAsset.workspace_id == workspace_id, | |
| ) | |
| ) | |
| if record is None: | |
| raise CanonicalAssetNotFoundError("Media asset was not found in this workspace.") | |
| return record | |
| async def verify_file(self, record: CanonicalMediaAsset, path: Path) -> None: | |
| if not path.is_file() or path.name != record.filename: | |
| raise CanonicalAssetNotFoundError("Media asset is no longer readable.") | |
| stat = path.stat() | |
| if stat.st_size != record.file_size: | |
| raise CanonicalAssetNotFoundError("Media asset changed after it was registered.") | |
| digest = await asyncio.to_thread(self._sha256, path) | |
| if not hmac.compare_digest(digest, record.sha256): | |
| raise CanonicalAssetNotFoundError("Media asset changed after it was registered.") | |
| def _sha256(path: Path) -> str: | |
| digest = hashlib.sha256() | |
| with path.open("rb") as stream: | |
| while chunk := stream.read(1024 * 1024): | |
| digest.update(chunk) | |
| return digest.hexdigest() | |