GallerySubscriber/backend/app/services/file_validator.py

"""Detect partial/truncated downloads via cheap magic-byte checks.

Real-world trigger: a Patreon JPEG landed in the library missing its
2-byte EOI marker (`FF D9`). PIL.ImageFile.load downstream raised
"image file is truncated (6 bytes not processed)". gallery-dl already
streams to .part and renames atomically, but TCP RST mid-transfer or
servers lying about Content-Length can still produce structurally
valid-but-incomplete files.

This validator runs at write time (and on the existing library via a
sweep task) and catches that class of failure. It only reads the head
and tail of each file — O(1) per file regardless of size — and is
format-aware for the formats we actually download. Unknown formats
(metadata sidecars, archives, etc.) pass through.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Optional


# Per-format magic-byte signatures. Each format defines what the head
# and tail must look like for the file to be structurally complete.
# We don't decode pixel data here — we just confirm the framing bytes
# the encoder is required to write. Adding a new format means one entry.
JPEG_HEAD = b"\xff\xd8\xff"
JPEG_TAIL = b"\xff\xd9"

PNG_HEAD = b"\x89PNG\r\n\x1a\n"
# Last 12 bytes of any complete PNG: 0-length IEND chunk + CRC.
PNG_TAIL = b"\x00\x00\x00\x00IEND\xaeB`\x82"

GIF_HEAD_87A = b"GIF87a"
GIF_HEAD_89A = b"GIF89a"
GIF_TAIL = b"\x3b"

WEBP_RIFF = b"RIFF"
WEBP_FORMAT = b"WEBP"


@dataclass(frozen=True)
class ValidationResult:
    """Outcome of validating one file.

    `ok=True` means either the file passed format-specific checks OR the
    format isn't one we know how to validate (we don't quarantine on
    ignorance). `format` is the detected/inferred format string for
    logging and stats; `reason` is populated on failure with a short
    operator-readable explanation.
    """
    ok: bool
    format: Optional[str] = None
    reason: Optional[str] = None
    size: int = 0


# Extensions we attempt to validate. Anything not in this set returns
# ok=True with format=None — gallery-dl writes JSON sidecars, ugoira
# zips, etc. that we don't want to police here.
VALIDATED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp"}


def is_validatable(path: Path) -> bool:
    """Return True if the path's extension is one we validate."""
    return path.suffix.lower() in VALIDATED_EXTENSIONS


def validate_file(path: Path) -> ValidationResult:
    """Check whether a file on disk is a complete, well-framed image.

    Returns ok=True for unknown formats and unreadable/missing files
    are surfaced as ok=False so callers can quarantine + investigate.
    """
    try:
        size = path.stat().st_size
    except FileNotFoundError:
        return ValidationResult(ok=False, reason="file not found")
    except OSError as e:
        return ValidationResult(ok=False, reason=f"stat failed: {e}")

    if size == 0:
        return ValidationResult(ok=False, format="empty", reason="zero-byte file", size=0)

    suffix = path.suffix.lower()
    if suffix not in VALIDATED_EXTENSIONS:
        return ValidationResult(ok=True, format=None, size=size)

    try:
        head, tail = _read_head_tail(path, head_n=16, tail_n=16)
    except OSError as e:
        return ValidationResult(ok=False, reason=f"read failed: {e}", size=size)

    if suffix in (".jpg", ".jpeg"):
        return _check_jpeg(head, tail, size)
    if suffix == ".png":
        return _check_png(head, tail, size)
    if suffix == ".gif":
        return _check_gif(head, tail, size)
    if suffix == ".webp":
        return _check_webp(head, tail, size)

    # Unreachable given VALIDATED_EXTENSIONS guard above, but keeps the
    # fall-through safe if the set diverges from the dispatch.
    return ValidationResult(ok=True, format=None, size=size)


def _read_head_tail(path: Path, head_n: int, tail_n: int) -> tuple[bytes, bytes]:
    """Read the first head_n and last tail_n bytes of a file in one open."""
    with path.open("rb") as f:
        head = f.read(head_n)
        f.seek(0, 2)  # SEEK_END
        size = f.tell()
        f.seek(max(0, size - tail_n))
        tail = f.read(tail_n)
    return head, tail


def _check_jpeg(head: bytes, tail: bytes, size: int) -> ValidationResult:
    if not head.startswith(JPEG_HEAD):
        return ValidationResult(
            ok=False, format="jpeg", reason="missing JPEG SOI marker (FF D8 FF)", size=size
        )
    if not tail.endswith(JPEG_TAIL):
        # The exact failure mode the user hit: file truncated before EOI.
        return ValidationResult(
            ok=False, format="jpeg", reason="missing JPEG EOI marker (FF D9)", size=size
        )
    return ValidationResult(ok=True, format="jpeg", size=size)


def _check_png(head: bytes, tail: bytes, size: int) -> ValidationResult:
    if not head.startswith(PNG_HEAD):
        return ValidationResult(
            ok=False, format="png", reason="missing PNG signature", size=size
        )
    if not tail.endswith(PNG_TAIL):
        return ValidationResult(
            ok=False, format="png", reason="missing PNG IEND chunk", size=size
        )
    return ValidationResult(ok=True, format="png", size=size)


def _check_gif(head: bytes, tail: bytes, size: int) -> ValidationResult:
    if not (head.startswith(GIF_HEAD_87A) or head.startswith(GIF_HEAD_89A)):
        return ValidationResult(
            ok=False, format="gif", reason="missing GIF header", size=size
        )
    if not tail.endswith(GIF_TAIL):
        return ValidationResult(
            ok=False, format="gif", reason="missing GIF trailer (3B)", size=size
        )
    return ValidationResult(ok=True, format="gif", size=size)


def _check_webp(head: bytes, tail: bytes, size: int) -> ValidationResult:
    # WEBP is a RIFF container: "RIFF" + 4-byte little-endian size + "WEBP".
    # The container size field equals (file size - 8). If the file got
    # truncated, the declared size won't match the on-disk size.
    if len(head) < 12 or not head.startswith(WEBP_RIFF) or head[8:12] != WEBP_FORMAT:
        return ValidationResult(
            ok=False, format="webp", reason="missing RIFF/WEBP header", size=size
        )
    declared = int.from_bytes(head[4:8], "little")
    if declared != size - 8:
        return ValidationResult(
            ok=False,
            format="webp",
            reason=f"RIFF size mismatch (declared {declared}, actual {size - 8})",
            size=size,
        )
    return ValidationResult(ok=True, format="webp", size=size)