diff --git a/tests/core/test_get_boxes.py b/tests/core/test_get_boxes.py new file mode 100644 index 00000000..546a41c4 --- /dev/null +++ b/tests/core/test_get_boxes.py @@ -0,0 +1,44 @@ +"""get_boxes must give each thread its own box when tracks read their init segments at the same time.""" + +from __future__ import annotations + +import struct +import threading +from uuid import UUID, uuid4 + +from unshackle.core.utilities import get_boxes + + +def tenc_init(kid: UUID, version: int, crypt_skip: int) -> bytes: + body = struct.pack(">B3s", version, b"\0\0\0") + bytes([0, crypt_skip if version else 0]) + b"\x01\x10" + kid.bytes + return b"\x00" * 8 + struct.pack(">I", 8 + len(body)) + b"tenc" + body + + +def test_get_boxes_finds_the_tenc() -> None: + kid = uuid4() + (box,) = get_boxes(tenc_init(kid, 1, 0x19), b"tenc") + assert box.key_ID == kid + + +def test_get_boxes_is_right_when_tracks_parse_at_the_same_time() -> None: + cases = [(uuid4(), version, crypt_skip) for version, crypt_skip in ((0, 0), (1, 0x19), (1, 0x91), (1, 0x55))] * 2 + failures: list[str] = [] + + def parse(kid: UUID, version: int, crypt_skip: int) -> None: + data = tenc_init(kid, version, crypt_skip) + for _ in range(1500): + try: + boxes = list(get_boxes(data, b"tenc")) + except Exception as e: + failures.append(repr(e)) + continue + if [box.key_ID for box in boxes] != [kid]: + failures.append("wrong or missing tenc") + + threads = [threading.Thread(target=parse, args=case) for case in cases] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=120) + assert not failures, failures[:3] + assert not any(thread.is_alive() for thread in threads) diff --git a/unshackle/core/tracks/subtitle.py b/unshackle/core/tracks/subtitle.py index b362e7c9..de8fb22f 100644 --- a/unshackle/core/tracks/subtitle.py +++ b/unshackle/core/tracks/subtitle.py @@ -25,7 +25,7 @@ from subtitle_filter import Subtitles from unshackle.core import binaries from unshackle.core.config import config from unshackle.core.tracks.track import Track -from unshackle.core.utilities import try_ensure_utf8, with_script +from unshackle.core.utilities import PYMP4_LOCK, try_ensure_utf8, with_script from unshackle.core.utils.subprocess import log_tool_run from unshackle.core.utils.webvtt import merge_segmented_webvtt, rebase_webvtt @@ -737,10 +737,10 @@ class Subtitle(Track): caption_set = pycaption.SRTReader().read(text) elif codec == Subtitle.Codec.fTTML: caption_lists: dict[str, pycaption.CaptionList] = defaultdict(pycaption.CaptionList) + with PYMP4_LOCK: + boxes = MP4.parse_stream(BytesIO(data)) for segment in ( - Subtitle.parse(box.data, Subtitle.Codec.TimedTextMarkupLang) - for box in MP4.parse_stream(BytesIO(data)) - if box.type == b"mdat" + Subtitle.parse(box.data, Subtitle.Codec.TimedTextMarkupLang) for box in boxes if box.type == b"mdat" ): for lang in segment.get_languages(): caption_lists[lang].extend(segment.get_captions(lang)) @@ -923,7 +923,9 @@ class Subtitle(Track): child["type"] = child.get("type") or child.get("format") yield child - for box in flatten_boxes(MP4.parse_stream(BytesIO(data))): + with PYMP4_LOCK: + boxes = MP4.parse_stream(BytesIO(data)) + for box in flatten_boxes(boxes): # init if box.type == b"mdhd": timescale = box.timescale @@ -963,7 +965,8 @@ class Subtitle(Track): if not saw_trun_box: raise ValueError("The TRUN box was not found in the Segmented WebVTT.") - vttc_boxes = MP4.parse_stream(BytesIO(box.data)) + with PYMP4_LOCK: + vttc_boxes = MP4.parse_stream(BytesIO(box.data)) current_time = base_time + period_start for sample, vttc_box in zip(samples, vttc_boxes): diff --git a/unshackle/core/utilities.py b/unshackle/core/utilities.py index 932c461c..1a768d06 100644 --- a/unshackle/core/utilities.py +++ b/unshackle/core/utilities.py @@ -9,6 +9,7 @@ import os import re import socket import sys +import threading import time import traceback import unicodedata @@ -456,6 +457,11 @@ def resolve_sort_langs( return resolved +# construct 2.8 gives each bit-field struct one shared stream, so two threads in a pymp4 +# parse or build of a box with bit fields (tenc, senc, trun, sidx, ...) read each other's bytes +PYMP4_LOCK = threading.Lock() + + def get_boxes(data: bytes, box_type: bytes, as_bytes: bool = False) -> Box: # type: ignore """ Scan a byte array for a wanted MP4/ISOBMFF box, then parse and yield each find. @@ -506,14 +512,13 @@ def get_boxes(data: bytes, box_type: bytes, as_bytes: bool = False) -> Box: # t box_start = pos - 4 try: - box = Box.parse(data[box_start:]) - if as_bytes: - box = Box.build(box) + with PYMP4_LOCK: + box = Box.parse(data[box_start:]) + built = Box.build(box) - yield box + yield built if as_bytes else box - box_size = len(Box.build(box)) - offset = box_start + box_size + offset = box_start + len(built) except IOError: break