"""视觉候选 snapshot 的版本化 canonicalization 和完整性校验。 该模块位于视觉匹配与任务持久化之间，只处理可序列化的任务事实，不读取数据库、 图片文件或运行时路径。snapshot 由后端生成并供 Task、attempt 和 workspace 共同引用。 """ from __future__ import annotations import copy import hashlib import json import math import re from datetime import datetime from pathlib import PurePosixPath from typing import Any, Mapping, Sequence VISUAL_MATCH_SNAPSHOT_PROTOCOL_VERSION = 2 MAX_VISUAL_MATCH_CANDIDATES = 50 SNAPSHOT_SHA256_RE = re.compile(r"[0-9a-f]{64}\Z") IMAGE_SHA256_RE = re.compile(r"[0-9a-fA-F]{64}\Z") SNAPSHOT_FIELDS = frozenset({"protocol_version", "snapshot_sha256", "matched_at", "query", "candidates"}) QUERY_FIELDS = frozenset({"meme_id", "image_sha256", "model", "dimensions", "preprocess_version"}) CANDIDATE_FIELDS = frozenset({"rank", "meme_id", "image_sha256", "size_bytes", "score", "relative_path", "context"}) MANIFEST_FIELDS = frozenset({"protocol_version", "snapshot_sha256", "matched_at", "candidate_count", "query", "candidates"}) SAFE_CONTEXT_FIELDS = frozenset({"title", "summary", "subjects", "visible_text", "references", "meaning", "keywords", "search_queries", "uncertainties"}) CONTEXT_LIST_FIELDS = frozenset({"subjects", "visible_text", "references", "keywords", "search_queries", "uncertainties"}) CONTEXT_URL_RE = re.compile(r"(?i)\b[a-z][a-z0-9+.-]*://") class VisualMatchSnapshotError(ValueError): """snapshot 输入或持久化内容不符合版本化契约时使用的稳定错误。""" def __init__(self, code: str, message: str | None = None) -> None: """保存稳定错误码，供任务服务映射而不暴露底层校验细节。""" super().__init__(message or code) self.code = code def _string(value: object, *, code: str, name: str, maximum: int = 255) -> str: """校验非空短字符串，拒绝控制字符和不受约束的超长字段。""" if not isinstance(value, str) or not value or len(value) > maximum or any(ord(char) < 0x20 for char in value): raise VisualMatchSnapshotError(code, f"snapshot {name} 无效") return value def _sha256(value: object, *, code: str, name: str) -> str: """校验图片或 snapshot 的十六进制 SHA-256 字符串。""" if not isinstance(value, str) or not IMAGE_SHA256_RE.fullmatch(value): raise VisualMatchSnapshotError(code, f"snapshot {name} 无效") return value.lower() def _relative_path(value: object) -> str: """校验 task-relative 候选文件名，拒绝绝对路径、跳转和控制字符。""" if not isinstance(value, str) or not value or len(value) > 512 or "\\" in value or any(ord(char) < 0x20 for char in value): raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "候选相对路径无效") parts = value.split("/") if any(not part or part in {".", ".."} for part in parts): raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "候选相对路径无效") path = PurePosixPath(value) if path.is_absolute() or not path.parts: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "候选相对路径无效") if value == "manifest.json": raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "候选相对路径占用 manifest 文件") return value def _canonical_json(value: Mapping[str, object]) -> bytes: """将已校验 snapshot 转为禁止 NaN 的稳定 JSON 字节序列。""" try: return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False).encode("utf-8") except (TypeError, ValueError, UnicodeError) as exc: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "snapshot 不是可序列化 JSON") from exc def _snapshot_digest(value: Mapping[str, object]) -> str: """计算不包含自身 hash 字段的 snapshot canonical SHA-256。""" payload = copy.deepcopy(dict(value)) payload.pop("snapshot_sha256", None) return hashlib.sha256(_canonical_json(payload)).hexdigest() def _canonical_context(value: object) -> dict[str, object]: """保留可供研究参考的固定语境字段，并移除 URL 与内部扩展。""" if not isinstance(value, Mapping): raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "候选 context 无效") result: dict[str, object] = {} for key in SAFE_CONTEXT_FIELDS: if key not in value: continue raw = value[key] if key in CONTEXT_LIST_FIELDS: if not isinstance(raw, list): raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "候选 context 列表无效") cleaned: list[str] = [] for item in raw: if not isinstance(item, str) or any(ord(char) < 0x20 or ord(char) == 0x7F for char in item): raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "候选 context 文本无效") if not CONTEXT_URL_RE.search(item): cleaned.append(item) result[key] = cleaned continue if raw is not None and not isinstance(raw, str): raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "候选 context 文本无效") if isinstance(raw, str) and (any(ord(char) < 0x20 or ord(char) == 0x7F for char in raw) or CONTEXT_URL_RE.search(raw)): continue result[key] = raw try: copied = copy.deepcopy(result) _canonical_json(copied) except (TypeError, ValueError, UnicodeError) as exc: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "候选 context 不是可序列化 JSON") from exc return copied def _candidate(value: Mapping[str, object], *, rank: int) -> dict[str, object]: """校验并复制一个候选，屏蔽 storage key、URL 等非 snapshot 字段。""" if not isinstance(value, Mapping): raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "候选结构无效") candidate_id = _string(value.get("meme_id"), code="visual_match_snapshot_invalid", name="meme_id") image_sha256 = _sha256(value.get("image_sha256"), code="visual_match_snapshot_invalid", name="image_sha256") raw_size = value.get("size_bytes") if not isinstance(raw_size, int) or isinstance(raw_size, bool) or raw_size < 0: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "候选 size_bytes 无效") raw_score = value.get("score") if not isinstance(raw_score, (int, float)) or isinstance(raw_score, bool) or not math.isfinite(float(raw_score)) or not -1.000001 <= float(raw_score) <= 1.000001: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "候选 score 无效") copied_context = _canonical_context(value.get("context", {})) return { "rank": rank, "meme_id": candidate_id, "image_sha256": image_sha256, "size_bytes": raw_size, "score": float(raw_score), "relative_path": _relative_path(value.get("relative_path")), "context": copied_context, } def build_visual_match_snapshot( *, query_meme_id: object, image_sha256: object, model: object, dimensions: object, preprocess_version: object, candidates: Sequence[Mapping[str, object]], matched_at: datetime | str, ) -> dict[str, object]: """构造 protocol v2 snapshot，并返回带完整性 hash 的深拷贝。 输入是后端已经完成 scope、模型和存储校验的查询及候选；输出可直接写入 Task JSONB。候选会按 score 降序、Meme ID 升序重排，重复 Meme ID 会被拒绝。 """ if isinstance(dimensions, bool) or not isinstance(dimensions, int) or dimensions <= 0: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "snapshot dimensions 无效") if isinstance(matched_at, datetime): if matched_at.tzinfo is None: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "snapshot matched_at 缺少时区") matched_value = matched_at.isoformat() else: matched_value = _string(matched_at, code="visual_match_snapshot_invalid", name="matched_at", maximum=64) try: datetime.fromisoformat(matched_value.replace("Z", "+00:00")) except ValueError as exc: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "snapshot matched_at 无效") from exc query = { "meme_id": _string(query_meme_id, code="visual_match_snapshot_invalid", name="query meme_id"), "image_sha256": _sha256(image_sha256, code="visual_match_snapshot_invalid", name="query image_sha256"), "model": _string(model, code="visual_match_snapshot_invalid", name="model"), "dimensions": dimensions, "preprocess_version": _string(preprocess_version, code="visual_match_snapshot_invalid", name="preprocess_version", maximum=128), } prepared: list[dict[str, object]] = [] seen: set[str] = set() seen_paths: set[str] = set() if not isinstance(candidates, Sequence) or len(candidates) > MAX_VISUAL_MATCH_CANDIDATES: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "snapshot 候选数量超出限制") for item in candidates: candidate_id = item.get("meme_id") if isinstance(item, Mapping) else None candidate_key = _string(candidate_id, code="visual_match_snapshot_invalid", name="candidate meme_id") if candidate_key in seen: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "snapshot 候选重复") seen.add(candidate_key) candidate = _candidate(item, rank=0) relative_path = str(candidate["relative_path"]) if relative_path in seen_paths: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "snapshot 候选路径重复") seen_paths.add(relative_path) prepared.append(candidate) prepared.sort(key=lambda item: (-float(item["score"]), str(item["meme_id"]))) for index, item in enumerate(prepared, start=1): item["rank"] = index snapshot: dict[str, object] = { "protocol_version": VISUAL_MATCH_SNAPSHOT_PROTOCOL_VERSION, "query": query, "matched_at": matched_value, "candidates": prepared, } snapshot["snapshot_sha256"] = _snapshot_digest(snapshot) return copy.deepcopy(snapshot) def validate_visual_match_snapshot(value: object, *, expected_sha256: str | None = None) -> dict[str, object]: """完整校验持久 snapshot 并返回不可被调用方修改的副本。 输入来自 JSONB 或 resume attempt；输出保证 protocol、排序、候选字段和 hash 均符合当前契约。``expected_sha256`` 用于 claim 恢复时拒绝错配事实。 """ if not isinstance(value, Mapping) or set(value) != SNAPSHOT_FIELDS or value.get("protocol_version") != VISUAL_MATCH_SNAPSHOT_PROTOCOL_VERSION: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "snapshot protocol 版本无效") query = value.get("query") if not isinstance(query, Mapping) or set(query) != QUERY_FIELDS: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "snapshot query 无效") candidates = value.get("candidates") if not isinstance(candidates, list): raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "snapshot candidates 无效") for candidate in candidates: if not isinstance(candidate, Mapping) or set(candidate) != CANDIDATE_FIELDS: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "snapshot 候选字段无效") normalized_context = _canonical_context(candidate.get("context")) if normalized_context != candidate.get("context"): raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "snapshot 候选 context 含非公开字段") rebuilt = build_visual_match_snapshot( query_meme_id=query.get("meme_id"), image_sha256=query.get("image_sha256"), model=query.get("model"), dimensions=query.get("dimensions"), preprocess_version=query.get("preprocess_version"), candidates=candidates, matched_at=value.get("matched_at"), ) supplied = value.get("snapshot_sha256") if not isinstance(supplied, str) or not SNAPSHOT_SHA256_RE.fullmatch(supplied) or supplied != rebuilt["snapshot_sha256"]: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "snapshot hash 无效") if expected_sha256 is not None and supplied != expected_sha256: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "snapshot hash 与任务事实不一致") return rebuilt def visual_match_snapshot_summary(value: object, *, expected_sha256: str | None = None) -> dict[str, object]: """返回公开任务 DTO 可使用的脱敏 snapshot 摘要。""" snapshot = validate_visual_match_snapshot(value, expected_sha256=expected_sha256) return { "protocol_version": snapshot["protocol_version"], "snapshot_sha256": snapshot["snapshot_sha256"], "matched_at": snapshot["matched_at"], "candidate_count": len(snapshot["candidates"]), } def visual_match_snapshot_manifest(value: object, *, expected_sha256: str | None = None) -> dict[str, object]: """返回可写入候选目录的 manifest，保留语境但不带物理存储信息。""" snapshot = validate_visual_match_snapshot(value, expected_sha256=expected_sha256) return { "protocol_version": snapshot["protocol_version"], "snapshot_sha256": snapshot["snapshot_sha256"], "matched_at": snapshot["matched_at"], "candidate_count": len(snapshot["candidates"]), "query": copy.deepcopy(snapshot["query"]), "candidates": copy.deepcopy(snapshot["candidates"]), } def validate_visual_match_snapshot_manifest(value: object, *, expected_sha256: str | None = None) -> dict[str, object]: """校验候选目录 manifest，并返回与 snapshot 一致的脱敏副本。 输入来自宿主物化器写入的普通文件；除摘要字段外，manifest 必须完整保留 同一 protocol v2 snapshot 的查询和候选事实，供 Runner 在启动外部进程前复核。 """ if not isinstance(value, Mapping) or set(value) != MANIFEST_FIELDS: raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "候选 manifest 字段无效") candidates = value.get("candidates") candidate_count = value.get("candidate_count") if not isinstance(candidates, list) or not isinstance(candidate_count, int) or isinstance(candidate_count, bool) or candidate_count != len(candidates): raise VisualMatchSnapshotError("visual_match_snapshot_invalid", "候选 manifest 数量无效") snapshot = { "protocol_version": value.get("protocol_version"), "snapshot_sha256": value.get("snapshot_sha256"), "matched_at": value.get("matched_at"), "query": value.get("query"), "candidates": candidates, } validated = validate_visual_match_snapshot(snapshot, expected_sha256=expected_sha256) return visual_match_snapshot_manifest(validated, expected_sha256=expected_sha256) __all__ = [ "VISUAL_MATCH_SNAPSHOT_PROTOCOL_VERSION", "MAX_VISUAL_MATCH_CANDIDATES", "VisualMatchSnapshotError", "build_visual_match_snapshot", "validate_visual_match_snapshot", "visual_match_snapshot_summary", "visual_match_snapshot_manifest", "validate_visual_match_snapshot_manifest", ]