Spaces:
Running
Running
| from __future__ import annotations | |
| import asyncio | |
| import hashlib | |
| import hmac | |
| import os | |
| from datetime import datetime, timedelta, timezone | |
| from pathlib import Path | |
| from typing import Any | |
| from sqlalchemy import select | |
| from sqlalchemy.exc import IntegrityError | |
| from app.core.config import Settings | |
| from app.security.database import SecurityDatabase | |
| from app.security.models import CanonicalMediaAsset, CanonicalMediaVariant | |
| 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, settings: Settings | None = None) -> None: | |
| self.database = database | |
| self.settings = settings | |
| 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 protected_request_ids(self) -> set[str]: | |
| """Return request IDs that still have canonical assets.""" | |
| async with self.database.session() as session: | |
| rows = await session.scalars( | |
| select(CanonicalMediaAsset.request_id).distinct() | |
| ) | |
| return set(rows.all()) | |
| async def cleanup_expired(self, output_dir: Path) -> int: | |
| """Delete expired, workspace-level assets and their files. | |
| Project-attached assets are retained. Missing files are reconciled from | |
| the database so stale asset IDs cannot accumulate indefinitely. | |
| """ | |
| hours = self.settings.asset_retention_hours if self.settings is not None else max( | |
| 1, int(os.getenv("MEDIA_ASSET_RETENTION_HOURS", "24")) | |
| ) | |
| cutoff = datetime.now(timezone.utc) - timedelta(hours=hours) | |
| removed = 0 | |
| async with self.database.session() as session: | |
| records = ( | |
| await session.scalars( | |
| select(CanonicalMediaAsset).where( | |
| CanonicalMediaAsset.created_at < cutoff, | |
| CanonicalMediaAsset.project_id.is_(None), | |
| ) | |
| ) | |
| ).all() | |
| for record in records: | |
| variant_id = await session.scalar( | |
| select(CanonicalMediaVariant.id) | |
| .where(CanonicalMediaVariant.asset_id == record.id) | |
| .limit(1) | |
| ) | |
| if variant_id is not None: | |
| continue | |
| path = (output_dir / record.request_id / record.filename).resolve() | |
| root = (output_dir / record.request_id).resolve() | |
| if path.parent != root: | |
| continue | |
| try: | |
| if path.is_file(): | |
| await asyncio.to_thread(path.unlink) | |
| except OSError: | |
| continue | |
| await session.delete(record) | |
| removed += 1 | |
| if removed: | |
| await session.commit() | |
| return removed | |
| 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() | |