"""이미지 업로드 입력 검증과 continuation 소유권 검사 유틸리티.""" import asyncio import time from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import asynccontextmanager from enum import StrEnum from pathlib import Path from typing import Literal from uuid import UUID from fastapi import HTTPException, UploadFile, status from sqlalchemy import select, text from app.database.session import AsyncSessionLocal, engine from app.home.models import Image from app.home.schemas.home_schema import ImageUploadResultItem from app.utils.logger import get_logger from config import azure_blob_settings logger = get_logger("image_upload") _image_upload_lock_slots = asyncio.Semaphore( azure_blob_settings.IMAGE_UPLOAD_MAX_CONCURRENT_LOCKS ) ALLOWED_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".heic", ".heif"} _HEIF_BRANDS = { b"heic", b"heix", b"hevc", b"hevx", b"heim", b"heis", b"mif1", b"msf1", } _IMAGE_SIGNATURE_BYTES = 64 class ImageUploadLockTimeoutError(TimeoutError): """동일 task 이미지 변경 락을 제한 시간 내 획득하지 못했습니다.""" class BlobReferenceState(StrEnum): """불명확한 DB commit 뒤 Blob URL 참조 확인 결과.""" ALL = "all" NONE = "none" MIXED = "mixed" UNKNOWN = "unknown" def should_cleanup_failed_upload_blobs( *, commit_started: bool, reference_state: BlobReferenceState = BlobReferenceState.UNKNOWN, ) -> bool: """DB commit 시도 후에는 참조가 없다고 확정된 경우에만 Blob을 삭제합니다.""" return not commit_started or reference_state == BlobReferenceState.NONE def classify_blob_references( expected_urls: set[str], found_urls: set[str] ) -> BlobReferenceState: """예상 Blob URL과 독립 조회 결과를 보존 우선 상태로 분류합니다.""" if not expected_urls: return BlobReferenceState.NONE matched_urls = expected_urls & found_urls if matched_urls == expected_urls: return BlobReferenceState.ALL if not matched_urls: return BlobReferenceState.NONE return BlobReferenceState.MIXED async def inspect_blob_references( task_id: str, blob_urls: set[str], ) -> BlobReferenceState: """독립 세션에서 이번 요청 Blob URL의 DB 반영 여부를 확인합니다. commit 응답 유실 직후의 짧은 가시성 경합을 피하려고 NONE 결과만 세 번 재확인합니다. 존재/혼재/조회 실패는 즉시 보존 쪽으로 판정합니다. """ if not blob_urls: return BlobReferenceState.NONE for attempt in range(3): try: async with AsyncSessionLocal() as session: result = await session.execute( select(Image.img_url).where( Image.task_id == task_id, Image.img_url.in_(blob_urls), ) ) found_urls = set(result.scalars().all()) state = classify_blob_references(blob_urls, found_urls) if state != BlobReferenceState.NONE: return state if attempt < 2: await asyncio.sleep(0.1 * (attempt + 1)) except Exception as exc: logger.error( f"[inspect_blob_references] DB verification failed - task_id: " f"{task_id}, {type(exc).__name__}: {exc}" ) return BlobReferenceState.UNKNOWN return BlobReferenceState.NONE async def compensate_failed_upload_blobs( *, task_id: str, blob_urls: set[str], commit_started: bool, cleanup: Callable[[], Awaitable[None]], ) -> BlobReferenceState: """DB 쓰기 실패 후 안전하다고 확인된 Blob만 보상 삭제합니다.""" reference_state = BlobReferenceState.UNKNOWN if commit_started: reference_state = await inspect_blob_references(task_id, blob_urls) if should_cleanup_failed_upload_blobs( commit_started=commit_started, reference_state=reference_state, ): await cleanup() return reference_state def validate_task_image_count(*, existing_count: int, incoming_count: int) -> int: """한 task에 누적 가능한 이미지 수를 검증하고 예상 총 개수를 반환합니다.""" total_count = existing_count + incoming_count max_task_images = azure_blob_settings.IMAGE_UPLOAD_MAX_TASK_IMAGES if total_count > max_task_images: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"한 작업에는 이미지를 최대 {max_task_images}개까지 추가할 수 있습니다.", ) return total_count @asynccontextmanager async def _image_upload_lock_slot( lock_name: str, timeout_seconds: int ) -> AsyncIterator[None]: """DB 락 연결 슬롯 대기에도 동일한 제한 시간을 적용합니다.""" try: await asyncio.wait_for( _image_upload_lock_slots.acquire(), timeout=timeout_seconds, ) except TimeoutError as exc: raise ImageUploadLockTimeoutError(lock_name) from exc try: yield finally: _image_upload_lock_slots.release() @asynccontextmanager async def image_upload_task_lock(task_id: str) -> AsyncIterator[None]: """동일 task 요청 전체를 MySQL advisory lock으로 직렬화합니다. yield dependency가 Azure 업로드부터 최종 태깅까지 connection을 점유하므로, worker별 semaphore로 main DB pool의 나머지 connection을 보존합니다. """ lock_name = f"image_upload:{task_id}" timeout_seconds = azure_blob_settings.IMAGE_UPLOAD_LOCK_TIMEOUT_SECONDS lock_started = time.perf_counter() async with _image_upload_lock_slot(lock_name, timeout_seconds): async with engine.connect() as connection: lock_result = await connection.execute( text("SELECT GET_LOCK(:lock_name, :timeout_seconds)"), { "lock_name": lock_name, "timeout_seconds": timeout_seconds, }, ) if lock_result.scalar_one_or_none() != 1: raise ImageUploadLockTimeoutError(lock_name) logger.info( f"[image_upload_task_lock] ACQUIRED - task_id: {task_id}, " f"wait_ms: {(time.perf_counter() - lock_started) * 1000:.1f}" ) try: yield finally: release_task = asyncio.create_task( connection.execute( text("SELECT RELEASE_LOCK(:lock_name)"), {"lock_name": lock_name}, ) ) try: release_result = await asyncio.shield(release_task) if release_result.scalar_one_or_none() != 1: raise RuntimeError(f"RELEASE_LOCK failed: {lock_name}") except asyncio.CancelledError: # shield 바깥 task가 다시 취소돼도 release query는 끝까지 기다립니다. try: await release_task except BaseException as release_exc: await connection.invalidate(release_exc) raise except BaseException as exc: # 락이 남은 connection이 pool로 복귀하지 않도록 폐기합니다. await connection.invalidate(exc) logger.error( f"[image_upload_task_lock] RELEASE_LOCK failed - " f"{type(exc).__name__}: {exc}" ) else: logger.info( f"[image_upload_task_lock] RELEASED - task_id: {task_id}, " f"held_ms: {(time.perf_counter() - lock_started) * 1000:.1f}" ) def is_valid_image_extension(filename: str | None) -> bool: """파일명의 확장자가 지원 이미지 확장자인지 확인합니다.""" if not filename: return False return Path(filename).suffix.lower() in ALLOWED_IMAGE_EXTENSIONS def _detect_image_format(header: bytes) -> str | None: """신뢰할 수 없는 파일명/MIME 대신 파일 시그니처로 형식을 판별합니다.""" if header.startswith(b"\xff\xd8\xff"): return "jpeg" if header.startswith(b"\x89PNG\r\n\x1a\n"): return "png" if len(header) >= 12 and header.startswith(b"RIFF") and header[8:12] == b"WEBP": return "webp" if len(header) >= 12 and header[4:8] == b"ftyp": brands = {header[8:12]} brands.update( header[index : index + 4] for index in range(16, len(header) - 3, 4) ) if brands & _HEIF_BRANDS: return "heif" return None def _extension_matches_format(extension: str, detected_format: str) -> bool: """동일 포맷의 별칭을 고려해 확장자와 시그니처 일치 여부를 확인합니다.""" expected_formats = { ".jpg": "jpeg", ".jpeg": "jpeg", ".png": "png", ".webp": "webp", ".heic": "heif", ".heif": "heif", } return expected_formats.get(extension) == detected_format async def inspect_upload_file(file: UploadFile) -> tuple[str, str, int]: """UploadFile을 상수 메모리로 검사하고 실제 바이트 크기를 반환합니다.""" original_name = file.filename or "" if len(original_name) > 255: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="파일명은 255자를 초과할 수 없습니다.", ) extension = Path(original_name).suffix.lower() max_file_size = azure_blob_settings.IMAGE_UPLOAD_MAX_FILE_SIZE_BYTES validation_chunk_size = min( azure_blob_settings.AZURE_BLOB_UPLOAD_BLOCK_SIZE_BYTES, 1024 * 1024, ) total_size = 0 header = bytearray() await file.seek(0) try: while chunk := await file.read(validation_chunk_size): total_size += len(chunk) if total_size > max_file_size: raise HTTPException( status_code=status.HTTP_413_CONTENT_TOO_LARGE, detail=( f"파일 '{original_name}'이 최대 크기 " f"{max_file_size // (1024 * 1024)} MiB를 초과합니다." ), ) if len(header) < _IMAGE_SIGNATURE_BYTES: remaining = _IMAGE_SIGNATURE_BYTES - len(header) header.extend(chunk[:remaining]) finally: await file.seek(0) if total_size == 0: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"빈 파일은 업로드할 수 없습니다: {original_name}", ) detected_format = _detect_image_format(bytes(header)) if detected_format is None or not _extension_matches_format( extension, detected_format ): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=( f"파일 내용과 확장자가 일치하는 지원 이미지가 아닙니다: {original_name}" ), ) return original_name, extension, total_size def normalize_continuation_task_id(task_id: str) -> str: """continuation task_id를 canonical UUID7 문자열로 검증합니다.""" try: parsed = UUID(task_id) except (ValueError, AttributeError): parsed = None if ( parsed is None or len(task_id) != 36 or parsed.version != 7 or str(parsed) != task_id.lower() ): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="task_id는 올바른 UUID7 형식이어야 합니다.", ) return str(parsed) def _blob_url_prefix(user_uuid: str, task_id: str) -> str: """현재 사용자와 task에 허용된 Azure 이미지 URL prefix를 반환합니다.""" base_url = azure_blob_settings.AZURE_BLOB_BASE_URL.rstrip("/") return f"{base_url}/{user_uuid}/{task_id}/image/" def assert_continuation_owner( images: list[Image], user_uuid: str, task_id: str ) -> None: """Image에 owner 컬럼이 없어 Blob 경로로 continuation 소유권을 검증합니다.""" if not images: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="이어 올릴 이미지 작업을 찾을 수 없습니다.", ) base_prefix = f"{azure_blob_settings.AZURE_BLOB_BASE_URL.rstrip('/')}/" owner_prefix = _blob_url_prefix(user_uuid, task_id) internal_urls = [ image.img_url for image in images if image.img_url.startswith(base_prefix) ] if not internal_urls or any( not image_url.startswith(owner_prefix) for image_url in internal_urls ): raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="이 이미지 업로드 작업을 이어서 수정할 권한이 없습니다.", ) def image_result_item(image: Image) -> ImageUploadResultItem: """DB Image를 기존 응답 아이템으로 변환합니다.""" base_prefix = f"{azure_blob_settings.AZURE_BLOB_BASE_URL.rstrip('/')}/" source: Literal["url", "blob"] = ( "blob" if image.img_url.startswith(base_prefix) else "url" ) return ImageUploadResultItem( id=image.id, img_name=image.img_name, img_url=image.img_url, img_order=image.img_order, source=source, )