255 lines
12 KiB
Python
255 lines
12 KiB
Python
"""Versioned, bounded media envelopes shared by master and endpoint bridge.
|
|
|
|
No provider or ComfyUI bootstrap imports. Network I/O is supplied by the caller,
|
|
which must enforce public-HTTPS resolution, redirect and streaming byte limits.
|
|
"""
|
|
import base64
|
|
import binascii
|
|
import hashlib
|
|
import json
|
|
import re
|
|
from datetime import datetime, timezone
|
|
from urllib.parse import urlsplit
|
|
|
|
PROFILE = "comfy-media-v1"
|
|
VERSION = 1
|
|
MAX_ARTIFACT_BYTES = 256 * 1024 * 1024
|
|
MAX_TOTAL_BYTES = 512 * 1024 * 1024
|
|
MAX_ARTIFACTS = 128
|
|
MAX_ITEMS = 64
|
|
MAX_DESCRIPTOR_BYTES = 128 * 1024
|
|
SOCKET_KINDS = {"images": "IMAGE", "audio": "AUDIO", "masks": "MASK", "texts": "STRING", "video": "VIDEO"}
|
|
# Internal End slots do not extend the public Collector or legacy media sockets.
|
|
VALUE_SOCKETS = tuple(f"value_{i}" for i in range(1, 11))
|
|
ARTIFACT_KINDS = frozenset(SOCKET_KINDS.values()) | {"VALUE"}
|
|
BINDING_FIELDS = {"binding_id", "collector_id", "job_id", "participant_id", "node_id", "socket", "kind"}
|
|
_ID = re.compile(r"[A-Za-z0-9_.:/-]{1,200}\Z")
|
|
_HASH = re.compile(r"[a-f0-9]{64}\Z")
|
|
|
|
|
|
def identifier(value):
|
|
if not isinstance(value, str) or not _ID.fullmatch(value):
|
|
raise ValueError("Invalid media identifier")
|
|
return value
|
|
|
|
|
|
def bounded_int(value, maximum, *, minimum=0):
|
|
if type(value) is not int or not minimum <= value <= maximum:
|
|
raise ValueError("Invalid media count or size")
|
|
return value
|
|
|
|
|
|
def validate_json(value, *, max_bytes=MAX_DESCRIPTOR_BYTES):
|
|
"""Bound metadata depth and size without accepting non-finite JSON numbers."""
|
|
def walk(item, depth=0):
|
|
if depth > 12:
|
|
raise ValueError("Media metadata nesting exceeds limit")
|
|
if isinstance(item, dict):
|
|
if len(item) > 4096 or any(not isinstance(k, str) for k in item):
|
|
raise ValueError("Invalid media metadata object")
|
|
for child in item.values():
|
|
walk(child, depth + 1)
|
|
elif isinstance(item, list):
|
|
if len(item) > 4096:
|
|
raise ValueError("Media metadata list exceeds limit")
|
|
for child in item:
|
|
walk(child, depth + 1)
|
|
elif item is not None and type(item) not in (str, int, float, bool):
|
|
raise ValueError("Invalid media metadata value")
|
|
walk(value)
|
|
try:
|
|
raw = json.dumps(value, allow_nan=False, separators=(",", ":")).encode("utf-8")
|
|
except (ValueError, TypeError, UnicodeError, RecursionError):
|
|
raise ValueError("Invalid media JSON") from None
|
|
if len(raw) > max_bytes:
|
|
raise ValueError("Media JSON exceeds byte limit")
|
|
return raw
|
|
|
|
|
|
def validate_binding(binding):
|
|
if not isinstance(binding, dict) or set(binding) != BINDING_FIELDS:
|
|
raise ValueError("Invalid media binding")
|
|
for field in BINDING_FIELDS - {"socket", "kind"}:
|
|
identifier(binding[field])
|
|
kind = "VALUE" if binding["socket"] in VALUE_SOCKETS else SOCKET_KINDS.get(binding["socket"])
|
|
if kind is None or kind != binding["kind"]:
|
|
raise ValueError("Media socket/type mismatch")
|
|
return binding
|
|
|
|
|
|
def validate_bindings(bindings):
|
|
if not isinstance(bindings, list) or not 1 <= len(bindings) <= MAX_ARTIFACTS:
|
|
raise ValueError("Invalid media bindings count")
|
|
result = {}
|
|
destinations = set()
|
|
for binding in bindings:
|
|
validate_binding(binding)
|
|
bid = binding["binding_id"]
|
|
destination = (binding["job_id"], binding["participant_id"], binding["collector_id"], binding["socket"])
|
|
if bid in result or destination in destinations:
|
|
raise ValueError("Duplicate media binding")
|
|
result[bid] = binding
|
|
destinations.add(destination)
|
|
return result
|
|
|
|
|
|
def _validate_url_transport(transport):
|
|
if set(transport) != {"type", "url", "expires_at"}:
|
|
raise ValueError("Invalid media URL transport")
|
|
url = transport["url"]
|
|
if not isinstance(url, str) or len(url) > 8192:
|
|
raise ValueError("Invalid media URL")
|
|
try:
|
|
parsed = urlsplit(url)
|
|
if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password or parsed.fragment:
|
|
raise ValueError()
|
|
if parsed.port not in (None, 443):
|
|
raise ValueError()
|
|
expiry = datetime.fromisoformat(transport["expires_at"].replace("Z", "+00:00"))
|
|
if expiry.tzinfo is None or expiry <= datetime.now(timezone.utc):
|
|
raise ValueError()
|
|
except (ValueError, TypeError, AttributeError):
|
|
raise ValueError("Invalid or expired media URL") from None
|
|
|
|
|
|
def validate_artifact(artifact, *, max_bytes=MAX_ARTIFACT_BYTES):
|
|
fields = {"id", "kind", "descriptor", "byte_length", "sha256", "transport"}
|
|
if not isinstance(artifact, dict) or set(artifact) != fields:
|
|
raise ValueError("Invalid media artifact")
|
|
identifier(artifact["id"])
|
|
if artifact["kind"] not in ARTIFACT_KINDS:
|
|
raise ValueError("Unsupported media kind")
|
|
if not isinstance(artifact["descriptor"], dict):
|
|
raise ValueError("Invalid media descriptor")
|
|
validate_json(artifact["descriptor"])
|
|
size = bounded_int(artifact["byte_length"], max_bytes)
|
|
digest = artifact["sha256"]
|
|
if not isinstance(digest, str) or not _HASH.fullmatch(digest):
|
|
raise ValueError("Invalid media checksum")
|
|
transport = artifact["transport"]
|
|
if not isinstance(transport, dict):
|
|
raise ValueError("Invalid media transport")
|
|
if transport.get("type") == "inline":
|
|
if set(transport) != {"type", "data_base64"} or not isinstance(transport["data_base64"], str):
|
|
raise ValueError("Invalid inline media transport")
|
|
if len(transport["data_base64"]) != 4 * ((size + 2) // 3):
|
|
raise ValueError("Inline media length mismatch")
|
|
elif transport.get("type") == "url":
|
|
_validate_url_transport(transport)
|
|
else:
|
|
raise ValueError("Unsupported media transport")
|
|
return artifact
|
|
|
|
|
|
def artifact_from_payload(artifact_id, kind, descriptor, payload, *, transport=None, max_bytes=MAX_ARTIFACT_BYTES):
|
|
if not isinstance(payload, bytes):
|
|
raise ValueError("Media payload must be bytes")
|
|
bounded_int(len(payload), max_bytes)
|
|
result = {"id": artifact_id, "kind": kind, "descriptor": descriptor,
|
|
"byte_length": len(payload), "sha256": hashlib.sha256(payload).hexdigest(),
|
|
"transport": transport if transport is not None else {
|
|
"type": "inline", "data_base64": base64.b64encode(payload).decode("ascii")}}
|
|
return validate_artifact(result, max_bytes=max_bytes)
|
|
|
|
|
|
async def resolve_artifact(artifact, *, download=None, max_bytes=MAX_ARTIFACT_BYTES):
|
|
"""download(url, *, max_bytes) must return bounded bytes; never log its URL."""
|
|
validate_artifact(artifact, max_bytes=max_bytes)
|
|
transport = artifact["transport"]
|
|
if transport["type"] == "inline":
|
|
try:
|
|
payload = base64.b64decode(transport["data_base64"], validate=True)
|
|
except (ValueError, binascii.Error):
|
|
raise ValueError("Invalid media base64") from None
|
|
else:
|
|
if download is None:
|
|
raise ValueError("Media URL download is not configured")
|
|
payload = await download(transport["url"], max_bytes=artifact["byte_length"])
|
|
if not isinstance(payload, bytes) or len(payload) != artifact["byte_length"]:
|
|
raise ValueError("Media byte length mismatch")
|
|
if hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
|
|
raise ValueError("Media checksum mismatch")
|
|
return payload
|
|
|
|
|
|
def validate_response(response, *, request_id, bindings, max_total_bytes=MAX_TOTAL_BYTES):
|
|
"""Validate ALL identity/completeness before downloading or delivering anything."""
|
|
expected = validate_bindings(bindings)
|
|
identifier(request_id)
|
|
if not isinstance(response, dict) or set(response) != {"version", "profile", "request_id", "bindings", "artifacts", "items"}:
|
|
raise ValueError("Invalid media response envelope")
|
|
if type(response["version"]) is not int or response["version"] != VERSION or response["profile"] != PROFILE:
|
|
raise ValueError("Unsupported media response version/profile")
|
|
if response["request_id"] != request_id:
|
|
raise ValueError("Media request identity mismatch")
|
|
returned = response["bindings"]
|
|
if not isinstance(returned, list) or len(returned) != len(expected):
|
|
raise ValueError("Missing or extra media bindings")
|
|
counts = {}
|
|
for row in returned:
|
|
if not isinstance(row, dict) or set(row) != BINDING_FIELDS | {"count"}:
|
|
raise ValueError("Invalid result binding")
|
|
binding = {key: row[key] for key in BINDING_FIELDS}
|
|
validate_binding(binding)
|
|
bid = binding["binding_id"]
|
|
if bid in counts or expected.get(bid) != binding:
|
|
raise ValueError("Media binding identity mismatch")
|
|
# Connected outputs must produce at least one value. Empty STRING is a value.
|
|
maximum = 1 if binding["socket"] in VALUE_SOCKETS else MAX_ITEMS
|
|
counts[bid] = bounded_int(row["count"], maximum, minimum=1)
|
|
artifacts = response["artifacts"]
|
|
if not isinstance(artifacts, list) or not 1 <= len(artifacts) <= MAX_ARTIFACTS:
|
|
raise ValueError("Invalid media artifacts count")
|
|
indexed = {}
|
|
total = 0
|
|
for artifact in artifacts:
|
|
validate_artifact(artifact)
|
|
aid = artifact["id"]
|
|
if aid in indexed:
|
|
raise ValueError("Duplicate media artifact")
|
|
total += artifact["byte_length"]
|
|
if total > max_total_bytes:
|
|
raise ValueError("Media response exceeds total byte limit")
|
|
indexed[aid] = artifact
|
|
items = response["items"]
|
|
if not isinstance(items, list) or len(items) != sum(counts.values()):
|
|
raise ValueError("Missing or extra media items")
|
|
ordered = {bid: [None] * count for bid, count in counts.items()}
|
|
used = set()
|
|
for item in items:
|
|
if not isinstance(item, dict) or set(item) != {"binding_id", "item_index", "artifact_id"}:
|
|
raise ValueError("Invalid media item")
|
|
bid, aid = identifier(item["binding_id"]), identifier(item["artifact_id"])
|
|
if bid not in counts or aid not in indexed:
|
|
raise ValueError("Unknown media item identity")
|
|
index = bounded_int(item["item_index"], counts[bid] - 1)
|
|
if ordered[bid][index] is not None or indexed[aid]["kind"] != expected[bid]["kind"]:
|
|
raise ValueError("Duplicate or mistyped media item")
|
|
ordered[bid][index] = indexed[aid]
|
|
used.add(aid)
|
|
if used != set(indexed) or any(value is None for values in ordered.values() for value in values):
|
|
raise ValueError("Incomplete media result manifest")
|
|
return ordered
|
|
|
|
|
|
async def decode_response(response, *, request_id, bindings, download=None, max_total_bytes=MAX_TOTAL_BYTES):
|
|
ordered = validate_response(response, request_id=request_id, bindings=bindings, max_total_bytes=max_total_bytes)
|
|
# Resolve and decode before exposing a single result to the collector queue.
|
|
from .codec import decode_payload
|
|
cache = {}
|
|
decoded = {}
|
|
for bid, artifacts in ordered.items():
|
|
values = []
|
|
for artifact in artifacts:
|
|
aid = artifact["id"]
|
|
if aid not in cache:
|
|
payload = await resolve_artifact(artifact, download=download)
|
|
if artifact["kind"] == "VIDEO":
|
|
from .video_codec import decode_video
|
|
cache[aid] = decode_video(artifact["descriptor"], payload)
|
|
else:
|
|
cache[aid] = decode_payload(artifact["kind"], artifact["descriptor"], payload)
|
|
values.append(cache[aid])
|
|
decoded[bid] = values
|
|
return decoded
|