feat(validator): add magic-byte file validator for partial-download detection
Detects truncated/incomplete downloads by checking format-specific head and tail bytes (JPEG SOI/EOI, PNG signature/IEND, GIF header/3B, WEBP RIFF size). Catches the production failure mode where a file landed missing its 2-byte JPEG EOI marker and PIL.ImageFile.load raised "image file is truncated (6 bytes not processed)" downstream. Reads only 16 bytes from each end — O(1) per file. Unknown extensions (JSON sidecars, ugoira zips, etc.) pass through. Hooks into the download flow in a follow-up commit. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,173 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,178 @@
|
||||
"""Unit tests for file_validator.
|
||||
|
||||
Covers the actual production failure mode (JPEG missing FF D9 EOI) and
|
||||
the symmetric checks for the other formats we validate. Unknown formats
|
||||
must pass through — we don't want to police gallery-dl's own JSON
|
||||
sidecars or zip ugoira archives.
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from app.services.file_validator import (
|
||||
GIF_HEAD_89A,
|
||||
JPEG_HEAD,
|
||||
JPEG_TAIL,
|
||||
PNG_HEAD,
|
||||
PNG_TAIL,
|
||||
WEBP_FORMAT,
|
||||
WEBP_RIFF,
|
||||
is_validatable,
|
||||
validate_file,
|
||||
)
|
||||
|
||||
|
||||
def _write(tmp_path: Path, name: str, data: bytes) -> Path:
|
||||
p = tmp_path / name
|
||||
p.write_bytes(data)
|
||||
return p
|
||||
|
||||
|
||||
# --- JPEG ---
|
||||
|
||||
|
||||
def test_jpeg_well_formed_passes(tmp_path):
|
||||
body = JPEG_HEAD + b"\x00" * 64 + JPEG_TAIL
|
||||
p = _write(tmp_path, "ok.jpg", body)
|
||||
result = validate_file(p)
|
||||
assert result.ok
|
||||
assert result.format == "jpeg"
|
||||
|
||||
|
||||
def test_jpeg_missing_eoi_fails(tmp_path):
|
||||
"""The exact production failure: file ends 6 bytes short, no FF D9.
|
||||
|
||||
Downstream PIL.ImageFile.load raised "image file is truncated (6 bytes
|
||||
not processed)" — we must catch this at write time.
|
||||
"""
|
||||
body = JPEG_HEAD + b"\x00" * 64 # No JPEG_TAIL
|
||||
p = _write(tmp_path, "truncated.jpg", body)
|
||||
result = validate_file(p)
|
||||
assert not result.ok
|
||||
assert result.format == "jpeg"
|
||||
assert "EOI" in result.reason
|
||||
|
||||
|
||||
def test_jpeg_missing_soi_fails(tmp_path):
|
||||
p = _write(tmp_path, "bogus.jpg", b"NOTAJPEG" * 8 + JPEG_TAIL)
|
||||
result = validate_file(p)
|
||||
assert not result.ok
|
||||
assert "SOI" in result.reason
|
||||
|
||||
|
||||
def test_jpeg_extension_jpeg_also_works(tmp_path):
|
||||
body = JPEG_HEAD + b"\x00" * 32 + JPEG_TAIL
|
||||
p = _write(tmp_path, "ok.jpeg", body)
|
||||
assert validate_file(p).ok
|
||||
|
||||
|
||||
# --- PNG ---
|
||||
|
||||
|
||||
def test_png_well_formed_passes(tmp_path):
|
||||
body = PNG_HEAD + b"\x00" * 32 + PNG_TAIL
|
||||
p = _write(tmp_path, "ok.png", body)
|
||||
result = validate_file(p)
|
||||
assert result.ok
|
||||
assert result.format == "png"
|
||||
|
||||
|
||||
def test_png_missing_iend_fails(tmp_path):
|
||||
body = PNG_HEAD + b"\x00" * 32 # No IEND chunk
|
||||
p = _write(tmp_path, "truncated.png", body)
|
||||
result = validate_file(p)
|
||||
assert not result.ok
|
||||
assert "IEND" in result.reason
|
||||
|
||||
|
||||
def test_png_missing_signature_fails(tmp_path):
|
||||
p = _write(tmp_path, "bogus.png", b"\x00" * 32 + PNG_TAIL)
|
||||
result = validate_file(p)
|
||||
assert not result.ok
|
||||
assert "signature" in result.reason
|
||||
|
||||
|
||||
# --- GIF ---
|
||||
|
||||
|
||||
def test_gif_well_formed_passes(tmp_path):
|
||||
body = GIF_HEAD_89A + b"\x00" * 32 + b"\x3b"
|
||||
p = _write(tmp_path, "ok.gif", body)
|
||||
result = validate_file(p)
|
||||
assert result.ok
|
||||
assert result.format == "gif"
|
||||
|
||||
|
||||
def test_gif_missing_trailer_fails(tmp_path):
|
||||
body = GIF_HEAD_89A + b"\x00" * 32 # No 0x3B
|
||||
p = _write(tmp_path, "truncated.gif", body)
|
||||
result = validate_file(p)
|
||||
assert not result.ok
|
||||
assert "trailer" in result.reason
|
||||
|
||||
|
||||
# --- WEBP ---
|
||||
|
||||
|
||||
def test_webp_well_formed_passes(tmp_path):
|
||||
payload = b"\x00" * 32
|
||||
declared = (4 + len(payload)).to_bytes(4, "little") # WEBP marker + payload
|
||||
body = WEBP_RIFF + declared + WEBP_FORMAT + payload
|
||||
p = _write(tmp_path, "ok.webp", body)
|
||||
result = validate_file(p)
|
||||
assert result.ok, f"unexpected failure: {result.reason}"
|
||||
assert result.format == "webp"
|
||||
|
||||
|
||||
def test_webp_size_mismatch_fails(tmp_path):
|
||||
"""RIFF declares more bytes than the file actually has — classic truncation."""
|
||||
payload = b"\x00" * 32
|
||||
lying_size = (4 + len(payload) + 100).to_bytes(4, "little")
|
||||
body = WEBP_RIFF + lying_size + WEBP_FORMAT + payload
|
||||
p = _write(tmp_path, "truncated.webp", body)
|
||||
result = validate_file(p)
|
||||
assert not result.ok
|
||||
assert "size mismatch" in result.reason
|
||||
|
||||
|
||||
# --- Unknown formats / edge cases ---
|
||||
|
||||
|
||||
def test_zero_byte_file_fails(tmp_path):
|
||||
p = _write(tmp_path, "empty.jpg", b"")
|
||||
result = validate_file(p)
|
||||
assert not result.ok
|
||||
assert "zero-byte" in result.reason
|
||||
|
||||
|
||||
def test_missing_file_fails(tmp_path):
|
||||
p = tmp_path / "ghost.jpg"
|
||||
result = validate_file(p)
|
||||
assert not result.ok
|
||||
assert "not found" in result.reason
|
||||
|
||||
|
||||
def test_unknown_extension_passes_through(tmp_path):
|
||||
"""JSON sidecars from gallery-dl's metadata postprocessor and other
|
||||
non-image artifacts must not be quarantined."""
|
||||
p = _write(tmp_path, "metadata.json", b'{"ok": true}')
|
||||
result = validate_file(p)
|
||||
assert result.ok
|
||||
assert result.format is None
|
||||
|
||||
|
||||
def test_zip_extension_passes_through(tmp_path):
|
||||
"""Ugoira archives (Pixiv) — gallery-dl writes .zip, not our concern."""
|
||||
p = _write(tmp_path, "anim.zip", b"PK\x03\x04random")
|
||||
result = validate_file(p)
|
||||
assert result.ok
|
||||
|
||||
|
||||
def test_is_validatable_flags_known_formats():
|
||||
assert is_validatable(Path("a.jpg"))
|
||||
assert is_validatable(Path("a.JPG")) # case-insensitive
|
||||
assert is_validatable(Path("a.png"))
|
||||
assert is_validatable(Path("a.gif"))
|
||||
assert is_validatable(Path("a.webp"))
|
||||
assert not is_validatable(Path("a.json"))
|
||||
assert not is_validatable(Path("a.mp4"))
|
||||
Reference in New Issue
Block a user