381 lines
13 KiB
Python
381 lines
13 KiB
Python
"""이미지 업로드 입력 검증과 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,
|
|
)
|