Spaces:
Sleeping
Sleeping
File size: 6,519 Bytes
3493993 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | 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.")
@staticmethod
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()
|