436 lines
15 KiB
Python
436 lines
15 KiB
Python
"""Inbound media cache + outbound chunked streaming.
|
|
|
|
Inbound: ``media.upload`` (chunked binary frames) -> ``cache_*_from_bytes``
|
|
-> a ``media_ref`` the adapter attaches to the ``MessageEvent``. Enforces
|
|
size limit + sha256 + MIME re-sniff (the client's declared kind is not
|
|
trusted). Chunks are streamed to a temp file so RAM stays bounded during
|
|
the transfer; the file is read once at the end for the hermes cache
|
|
helpers.
|
|
|
|
Outbound: ``send_*`` -> register the (already delivery-validated) file in
|
|
the media registry, mint a ``media_id``, emit ``media.offer {media_id,
|
|
kind, mime, size, filename}``; serve bytes on ``media.pull`` as chunked
|
|
binary frames. Delivery-path security via ``validate_media_delivery_path``
|
|
(re-checked at pull time, so only allowed files are servable).
|
|
|
|
Reuses hermes ``cache_image/audio/video/document_from_bytes`` + the
|
|
``_looks_like_image`` / ``sniff_container`` magic-byte sniffers. Temp files
|
|
live under ``get_hermes_home()/"iris"/media/tmp``.
|
|
|
|
Milestone M4.
|
|
"""
|
|
|
|
import contextlib
|
|
import hashlib
|
|
import logging
|
|
import os
|
|
import re
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
from gateway.platforms.base import (
|
|
_looks_like_image,
|
|
cache_audio_from_bytes,
|
|
cache_document_from_bytes,
|
|
cache_image_from_bytes,
|
|
cache_video_from_bytes,
|
|
)
|
|
from tools.audio_container import sniff_container
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# One WS binary frame carries at most this many media bytes (docs/07 §7.5).
|
|
DEFAULT_CHUNK_BYTES = 256 * 1024
|
|
|
|
KINDS = ("image", "audio", "video", "document", "voice")
|
|
|
|
_SHA256_RE = re.compile(r"^[0-9a-f]{64}$")
|
|
|
|
# Magic-byte containers that are unambiguously audio (vs video-in-same-box).
|
|
_AUDIO_CONTAINERS = {"m4a", "ogg", "flac", "wav", "mp3", "aac"}
|
|
_VIDEO_CONTAINERS = {"mp4", "webm"}
|
|
# Longest file extension we trust from the client (e.g. ".webm").
|
|
_MAX_EXT_LEN = 6
|
|
|
|
# Extension -> MIME for outbound offers (the app picks a player/viewer from it).
|
|
_EXT_TO_MIME: dict[str, str] = {
|
|
".jpg": "image/jpeg",
|
|
".jpeg": "image/jpeg",
|
|
".png": "image/png",
|
|
".webp": "image/webp",
|
|
".gif": "image/gif",
|
|
".bmp": "image/bmp",
|
|
".heic": "image/heic",
|
|
".mp4": "video/mp4",
|
|
".mov": "video/quicktime",
|
|
".webm": "video/webm",
|
|
".mkv": "video/x-matroska",
|
|
".avi": "video/x-msvideo",
|
|
".3gp": "video/3gpp",
|
|
".ogg": "audio/ogg",
|
|
".opus": "audio/ogg",
|
|
".m4a": "audio/mp4",
|
|
".mp3": "audio/mpeg",
|
|
".wav": "audio/wav",
|
|
".flac": "audio/flac",
|
|
".aac": "audio/aac",
|
|
".pdf": "application/pdf",
|
|
".txt": "text/plain",
|
|
".md": "text/markdown",
|
|
".csv": "text/csv",
|
|
".html": "text/html",
|
|
".json": "application/json",
|
|
".zip": "application/zip",
|
|
".doc": "application/msword",
|
|
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
|
".xls": "application/vnd.ms-excel",
|
|
".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
|
".ppt": "application/vnd.ms-powerpoint",
|
|
".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
|
}
|
|
|
|
# MIME -> extension for inbound caching (the cache helpers take an ext).
|
|
_MIME_TO_EXT: dict[str, str] = {
|
|
"image/jpeg": ".jpg",
|
|
"image/png": ".png",
|
|
"image/webp": ".webp",
|
|
"image/gif": ".gif",
|
|
"image/bmp": ".bmp",
|
|
"image/heic": ".heic",
|
|
"video/mp4": ".mp4",
|
|
"video/quicktime": ".mov",
|
|
"video/webm": ".webm",
|
|
"video/x-matroska": ".mkv",
|
|
"video/x-msvideo": ".avi",
|
|
"audio/ogg": ".ogg",
|
|
"audio/mp4": ".m4a",
|
|
"audio/mpeg": ".mp3",
|
|
"audio/wav": ".wav",
|
|
"audio/flac": ".flac",
|
|
"audio/aac": ".aac",
|
|
}
|
|
|
|
|
|
class MediaError(Exception):
|
|
"""Media flow failure carrying a wire ``error`` code."""
|
|
|
|
def __init__(self, code: str, message: str):
|
|
super().__init__(message)
|
|
self.code = code
|
|
self.message = message
|
|
|
|
|
|
def mime_for_path(path: str) -> str:
|
|
ext = os.path.splitext(path)[1].lower()
|
|
return _EXT_TO_MIME.get(ext, "application/octet-stream")
|
|
|
|
|
|
def ext_for_mime(mime: str, filename: str, default: str) -> str:
|
|
"""Best-effort extension for the hermes cache helpers.
|
|
|
|
Prefers the MIME map, falls back to the filename's own extension, then
|
|
*default*.
|
|
"""
|
|
base = (mime or "").split(";")[0].strip().lower()
|
|
ext = _MIME_TO_EXT.get(base)
|
|
if ext:
|
|
return ext
|
|
file_ext = os.path.splitext(filename or "")[1].lower()
|
|
if file_ext and len(file_ext) <= _MAX_EXT_LEN:
|
|
return file_ext
|
|
return default
|
|
|
|
|
|
def reclassify_kind(head: bytes, declared_kind: str) -> str:
|
|
"""Re-sniff the real kind from magic bytes (don't trust the client).
|
|
|
|
Images are detected first (RIFF/WEBP would otherwise look like a
|
|
container); then AV containers. Video declared by the client is kept
|
|
when no magic matches (some codecs have no leading magic we check);
|
|
anything else falls back to ``document``.
|
|
"""
|
|
if _looks_like_image(head):
|
|
return "image"
|
|
container = sniff_container(head)
|
|
if container in _AUDIO_CONTAINERS:
|
|
return "voice" if declared_kind == "voice" else "audio"
|
|
if container in _VIDEO_CONTAINERS:
|
|
return "video"
|
|
return declared_kind if declared_kind == "video" else "document"
|
|
|
|
|
|
@dataclass
|
|
class MediaEntry:
|
|
"""One servable/cached media object (inbound ref or outbound offer)."""
|
|
|
|
media_id: str # mu_* (inbound ref, app-chosen) or md_* (outbound, minted)
|
|
kind: str
|
|
mime: str
|
|
size: int
|
|
filename: str
|
|
path: str # absolute local path
|
|
ts: float = 0.0
|
|
|
|
|
|
class UploadSession:
|
|
"""One in-flight ``media.upload`` (per connection; one at a time).
|
|
|
|
Chunks are appended to a temp file (bounded RAM); the sha256 is
|
|
accumulated incrementally. The declared size is enforced as bytes
|
|
arrive so an over-limit transfer is rejected early.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
media_ref: str,
|
|
kind: str,
|
|
mime: str,
|
|
filename: str,
|
|
declared_size: int,
|
|
request_id: int | None,
|
|
max_bytes: int,
|
|
tmp_dir: Path,
|
|
):
|
|
self.media_ref = media_ref
|
|
self.kind = kind
|
|
self.mime = mime
|
|
self.filename = filename
|
|
self.declared_size = declared_size
|
|
self.request_id = request_id
|
|
self.max_bytes = max_bytes
|
|
self.received = 0
|
|
self.failed = False
|
|
self.error_code = ""
|
|
self.error_message = ""
|
|
self._sha = hashlib.sha256()
|
|
fd, self.tmp_path = tempfile.mkstemp(prefix="upl_", dir=str(tmp_dir))
|
|
self._fh = os.fdopen(fd, "wb")
|
|
|
|
def feed(self, chunk: bytes) -> None:
|
|
if self.failed:
|
|
return
|
|
self.received += len(chunk)
|
|
if self.max_bytes and self.received > self.max_bytes:
|
|
self._fail(
|
|
"media_too_large",
|
|
f"upload exceeds limit ({self.received} > {self.max_bytes} bytes)",
|
|
)
|
|
return
|
|
try:
|
|
self._fh.write(chunk)
|
|
self._fh.flush()
|
|
self._sha.update(chunk)
|
|
except OSError as e:
|
|
self._fail("internal", f"temp file write failed: {e}")
|
|
|
|
def _fail(self, code: str, message: str) -> None:
|
|
self.failed = True
|
|
self.error_code = code
|
|
self.error_message = message
|
|
logger.warning("iris: upload %s failed: %s", self.media_ref, message)
|
|
|
|
def digest(self) -> str:
|
|
return self._sha.hexdigest()
|
|
|
|
def read_bytes(self) -> bytes:
|
|
self._fh.close()
|
|
with open(self.tmp_path, "rb") as f:
|
|
return f.read()
|
|
|
|
def close(self) -> None:
|
|
"""Discard the session and remove the temp file."""
|
|
# Best-effort cleanup: a file that is already gone (or a handle that
|
|
# is already closed) needs no further handling.
|
|
with contextlib.suppress(Exception):
|
|
self._fh.close()
|
|
with contextlib.suppress(OSError):
|
|
os.unlink(self.tmp_path)
|
|
|
|
|
|
class MediaStore:
|
|
"""In-memory media registry + upload session tracking.
|
|
|
|
Inbound refs (``mu_*``, app-chosen) are short-lived: created on
|
|
``media.upload.end``, consumed by the next ``message.send``. Outbound
|
|
offers (``md_*``, minted here) live until pulled (or process restart —
|
|
the outbox replays the offer frame, but the id->path map is in-memory
|
|
for M4).
|
|
"""
|
|
|
|
def __init__(self, hermes_home: Path):
|
|
self._tmp_dir = hermes_home / "iris" / "media" / "tmp"
|
|
self._tmp_dir.mkdir(parents=True, exist_ok=True)
|
|
self._lock = threading.Lock()
|
|
# (device_id, media_ref) -> UploadSession (one active per device)
|
|
self._uploads: dict[tuple[str, str], UploadSession] = {}
|
|
self._inbound: dict[str, MediaEntry] = {}
|
|
self._outbound: dict[str, MediaEntry] = {}
|
|
|
|
# ── Inbound uploads ───────────────────────────────────────────────────
|
|
|
|
def create_upload(
|
|
self,
|
|
device_id: str,
|
|
media_ref: str,
|
|
kind: str,
|
|
mime: str,
|
|
filename: str,
|
|
declared_size: int,
|
|
request_id: int | None,
|
|
max_bytes: int,
|
|
) -> UploadSession:
|
|
with self._lock:
|
|
for dev, _ref in self._uploads:
|
|
if dev == device_id:
|
|
raise MediaError(
|
|
"unsupported", "an upload is already in progress on this connection"
|
|
)
|
|
if media_ref in self._inbound:
|
|
raise MediaError("unsupported", f"media_ref {media_ref} already used")
|
|
sess = UploadSession(
|
|
media_ref,
|
|
kind,
|
|
mime,
|
|
filename,
|
|
declared_size,
|
|
request_id,
|
|
max_bytes,
|
|
self._tmp_dir,
|
|
)
|
|
self._uploads[(device_id, media_ref)] = sess
|
|
return sess
|
|
|
|
def get_upload(self, device_id: str, media_ref: str | None = None) -> UploadSession | None:
|
|
with self._lock:
|
|
if media_ref is not None:
|
|
return self._uploads.get((device_id, media_ref))
|
|
for (dev, _ref), sess in self._uploads.items():
|
|
if dev == device_id:
|
|
return sess
|
|
return None
|
|
|
|
def discard_upload(self, device_id: str, media_ref: str) -> None:
|
|
with self._lock:
|
|
sess = self._uploads.pop((device_id, media_ref), None)
|
|
if sess is not None:
|
|
sess.close()
|
|
|
|
def complete_upload(self, device_id: str, media_ref: str, sha256_hex: str) -> MediaEntry:
|
|
"""Verify + cache a finished upload. Raises ``MediaError`` on failure."""
|
|
with self._lock:
|
|
sess = self._uploads.pop((device_id, media_ref), None)
|
|
if sess is None:
|
|
raise MediaError("not_found", f"no active upload for media_ref {media_ref}")
|
|
try:
|
|
if sess.failed:
|
|
raise MediaError(sess.error_code, sess.error_message)
|
|
if sess.received != sess.declared_size:
|
|
raise MediaError(
|
|
"internal",
|
|
f"size mismatch (declared {sess.declared_size}, received {sess.received})",
|
|
)
|
|
if not _SHA256_RE.match(sha256_hex or ""):
|
|
raise MediaError("unsupported", "sha256 must be 64 hex chars")
|
|
if sess.digest() != sha256_hex:
|
|
raise MediaError("internal", "sha256 mismatch")
|
|
|
|
data = sess.read_bytes()
|
|
kind = reclassify_kind(data[:65536], sess.kind)
|
|
try:
|
|
if kind == "image":
|
|
path = cache_image_from_bytes(
|
|
data, ext_for_mime(sess.mime, sess.filename, ".jpg")
|
|
)
|
|
elif kind in ("audio", "voice"):
|
|
path = cache_audio_from_bytes(
|
|
data, ext_for_mime(sess.mime, sess.filename, ".ogg")
|
|
)
|
|
elif kind == "video":
|
|
path = cache_video_from_bytes(
|
|
data, ext_for_mime(sess.mime, sess.filename, ".mp4")
|
|
)
|
|
else:
|
|
path = cache_document_from_bytes(data, sess.filename or "document")
|
|
except ValueError as e:
|
|
# hermes cap (gateway.max_inbound_media_bytes) or a
|
|
# non-image payload masquerading as an image.
|
|
if "too large" in str(e):
|
|
raise MediaError("media_too_large", str(e)) from e
|
|
raise MediaError("unsupported", str(e)) from e
|
|
|
|
entry = MediaEntry(
|
|
media_id=media_ref,
|
|
kind=kind,
|
|
mime=mime_for_path(path),
|
|
size=len(data),
|
|
filename=os.path.basename(path),
|
|
path=path,
|
|
ts=time.time(),
|
|
)
|
|
with self._lock:
|
|
self._inbound[media_ref] = entry
|
|
logger.info(
|
|
"iris: upload %s cached as %s (%s, %d bytes)",
|
|
media_ref,
|
|
kind,
|
|
path,
|
|
len(data),
|
|
)
|
|
return entry
|
|
finally:
|
|
sess.close()
|
|
|
|
def get_inbound(self, media_ref: str) -> MediaEntry | None:
|
|
with self._lock:
|
|
return self._inbound.get(media_ref)
|
|
|
|
def pop_inbound(self, media_ref: str) -> MediaEntry | None:
|
|
with self._lock:
|
|
return self._inbound.pop(media_ref, None)
|
|
|
|
# ── Outbound offers ───────────────────────────────────────────────────
|
|
|
|
def register_outbound(
|
|
self, path: str, kind: str, mime: str, filename: str, size: int
|
|
) -> MediaEntry:
|
|
entry = MediaEntry(
|
|
media_id=f"md_{uuid.uuid4().hex[:12]}",
|
|
kind=kind,
|
|
mime=mime,
|
|
size=size,
|
|
filename=filename or os.path.basename(path),
|
|
path=path,
|
|
ts=time.time(),
|
|
)
|
|
with self._lock:
|
|
self._outbound[entry.media_id] = entry
|
|
return entry
|
|
|
|
def get_outbound(self, media_id: str) -> MediaEntry | None:
|
|
with self._lock:
|
|
return self._outbound.get(media_id)
|
|
|
|
# ── Cleanup ───────────────────────────────────────────────────────────
|
|
|
|
def discard_device(self, device_id: str) -> None:
|
|
"""Drop all in-flight uploads for a disconnected device."""
|
|
with self._lock:
|
|
stale = [k for k in self._uploads if k[0] == device_id]
|
|
sessions = [self._uploads.pop(k) for k in stale]
|
|
for sess in sessions:
|
|
sess.close()
|
|
|
|
def prune_outbound(self, max_age_s: float = 24 * 3600) -> int:
|
|
"""Drop offers older than *max_age_s* (their files stay in the hermes
|
|
cache; only the id->path map is pruned)."""
|
|
now = time.time()
|
|
with self._lock:
|
|
stale = [k for k, e in self._outbound.items() if now - e.ts > max_age_s]
|
|
for k in stale:
|
|
del self._outbound[k]
|
|
return len(stale)
|