aboutsummaryrefslogtreecommitdiffstats
path: root/packages/meshbay-node/tests/test_stream_seek_audio_alignment.py
blob: 5d78c00cc3c656b5f1c5a05dbccdcd61afada47a (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
"""
A seek must not leave the audio a GOP behind the picture.

`-ss` before `-i` lands on the keyframe at or before the requested position.
`-accurate_seek` is on by default and trims each stream to exactly that
position *where it can*: it cannot trim copied video, which has to begin on a
keyframe, but it does trim the re-encoded audio. The output then begins with
video from the keyframe and audio from the request, with a hole between them
one whole GOP wide — and after the hole, sound and picture belong to different
moments for the rest of the film.

**Every timestamp is correct while this happens**, which is why it survived:
the first PTS of each stream, their durations, their spans and the browser's
own A/V delta through MediaSource all agree, because both streams really are
where the container says. Only the content is displaced. So these tests assert
on the *shape of the audio track* — the hole — rather than on timestamps, which
cannot see it.

Found against a real film with a 10 s keyframe interval: seeking to 609 s
against a keyframe at 599.104 s left 9.979 s of silence. Re-encoded video never
showed it, since video that is re-encoded can begin exactly where it was asked
to — which is why the one part of a library that worked was the part the node
transcodes.
"""

import shutil
import subprocess
from pathlib import Path

import pytest
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from meshbay_common.crypto import generate_gek
from meshbay_common.webcrypto import chunk_key_aes, decrypt_chunk_aes
from meshbay_node.indexer.group_index import GroupIndex
from meshbay_node.transport.webrtc_server import WebRTCPeerSession

from conftest import needs_subprocess, one_root

_HAVE_FFMPEG = shutil.which("ffmpeg") and shutil.which("ffprobe")
pytestmark = [
    pytest.mark.asyncio,
    pytest.mark.skipif(not _HAVE_FFMPEG, reason="ffmpeg/ffprobe not installed"),
    needs_subprocess,
]

FPS = 24
GOP_S = 10                 # keyframe every 10 s, as real films in the wild are
CLIP_S = 40
SEEK_S = 25.0              # mid-GOP: the keyframe before it is at 20 s
EXPECTED_PREROLL = 5.0     # SEEK_S - 20.0


def _make_long_gop_clip(path: Path) -> None:
    """Video with a 10 s keyframe interval and continuous audio.

    The interval is what gives the fault room to appear: with a keyframe every
    frame there is nothing for the audio to be trimmed past.
    """
    subprocess.run(
        ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
         "-f", "lavfi", "-i", f"testsrc=size=320x240:rate={FPS}:duration={CLIP_S}",
         "-f", "lavfi", "-i",
         f"sine=frequency=440:duration={CLIP_S}:sample_rate=48000",
         "-c:v", "libx264", "-preset", "ultrafast",
         "-g", str(GOP_S * FPS), "-keyint_min", str(GOP_S * FPS),
         "-sc_threshold", "0",
         "-c:a", "aac", str(path)],
        check=True, capture_output=True,
    )


def _session(tmp_path: Path, video_path: Path, gek: bytes):
    import blake3
    file_bytes = video_path.read_bytes()
    file_id = blake3.blake3(file_bytes).hexdigest()

    sk_node = Ed25519PrivateKey.generate()
    index = GroupIndex(group_id="g" * 32, sk_node=sk_node, gek=gek)
    from meshbay_common.protocol import IndexEntry
    index.add_entry(IndexEntry(
        id=file_id, name=video_path.name, path=video_path.parent.name,
        size=len(file_bytes), type="video", added_at=0))

    session = WebRTCPeerSession.__new__(WebRTCPeerSession)
    session._ctx = {
        "roots": one_root(video_path.parent),
        "index": index,
        "gek": gek,
        "sk_node": sk_node,
        "max_concurrent_streams": 4,
    }
    session._group_id = None
    session._user_id = "tester"
    session._stream_stopped = False
    session._stream_keepalives = 0
    session.sent = []
    session._send = session.sent.append
    session._audit = lambda *a, **k: None
    return session, file_id


def _reassemble(sent, gek, file_id) -> bytes:
    file_hash = bytes.fromhex(file_id)
    out = b""
    for m in sorted((m for m in sent if m.get("type") == "stream_data"),
                    key=lambda m: m["segment_index"]):
        key = chunk_key_aes(gek, file_hash, m["segment_index"])
        out += decrypt_chunk_aes(key, m["nonce"], m["ct"])
    return out


def _packet_times(path: Path, kind: str) -> list[float]:
    out = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", kind,
         "-show_entries", "packet=pts_time", "-of", "csv=p=0", str(path)],
        capture_output=True, text=True).stdout
    return sorted(float(x.rstrip(",")) for x in out.split() if x.strip())


def _largest_audio_gap(path: Path) -> float:
    t = _packet_times(path, "a")
    return max((t[i + 1] - t[i] for i in range(len(t) - 1)), default=0.0)


async def _stream_to_file(tmp_path: Path, clip: Path, start: float, tag: str) -> Path:
    gek = generate_gek()
    session, file_id = _session(tmp_path, clip, gek)
    await session._stream_video_inner(
        {"file_id": file_id, "start": start, "credits": 0})
    errors = [m for m in session.sent if m.get("type") == "error"]
    assert not errors, f"streaming must not fail: {errors}"
    out = tmp_path / f"seek-{tag}.mp4"
    out.write_bytes(_reassemble(session.sent, gek, file_id))
    return out


async def test_a_seek_mid_gop_leaves_no_hole_in_the_audio(tmp_path):
    """The regression itself.

    Without the fix the audio begins one whole GOP after the video and the gap
    is the distance from the keyframe to the request — here five seconds.
    """
    clip = tmp_path / "longgop.mkv"
    _make_long_gop_clip(clip)

    out = await _stream_to_file(tmp_path, clip, SEEK_S, "midgop")
    gap = _largest_audio_gap(out)

    assert gap < 0.5, (
        f"the audio track has a {gap:.3f}s hole after a seek to {SEEK_S}s. "
        "Video was copied from the keyframe before it while the re-encoded "
        "audio was trimmed to the request, so sound and picture are a GOP "
        "apart for the rest of the stream")


async def test_the_fixture_really_has_a_long_gop(tmp_path):
    """Otherwise the test above proves nothing: with a keyframe every frame
    there is no interval for the audio to be trimmed past, and it would pass
    against the broken code."""
    clip = tmp_path / "longgop.mkv"
    _make_long_gop_clip(clip)

    keys = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "v:0",
         "-skip_frame", "nokey", "-show_entries", "frame=pts_time",
         "-of", "csv=p=0", str(clip)],
        capture_output=True, text=True).stdout.split()
    times = sorted(float(x.rstrip(",")) for x in keys if x.strip())
    gaps = [times[i + 1] - times[i] for i in range(len(times) - 1)]

    assert gaps, "no keyframe interval to measure"
    assert max(gaps) >= GOP_S - 1, (
        f"fixture keyframe interval is {max(gaps):.1f}s, too short to expose "
        "the fault this file exists to guard")
    assert SEEK_S not in times, "the seek must land between keyframes"


async def test_the_stream_carries_as_much_sound_as_picture(tmp_path):
    """The positive form, and it has to be measured on **decoded audio**.

    Comparing the first packet of each stream proves nothing here: the hole
    leaves one audio packet at zero and resumes a GOP later, so both streams
    still *start* at the same timestamp and every timestamp comparison agrees.
    What is actually missing is a GOP of sound, so that is what is counted.
    """
    clip = tmp_path / "longgop.mkv"
    _make_long_gop_clip(clip)

    out = await _stream_to_file(tmp_path, clip, SEEK_S, "together")
    decoded = subprocess.run(
        ["ffmpeg", "-v", "error", "-i", str(out), "-map", "0:a:0",
         "-ac", "1", "-ar", "8000", "-f", "s16le", "-"],
        capture_output=True).stdout
    audio_s = len(decoded) / 2 / 8000
    video_s = max(_packet_times(out, "v")) - min(_packet_times(out, "v"))

    assert audio_s >= video_s - 0.5, (
        f"{video_s:.2f}s of picture arrived with only {audio_s:.2f}s of sound: "
        f"{video_s - audio_s:.2f}s of audio is missing from the head of the "
        "stream, which is the GOP the seek copied video across")


async def test_playing_from_the_beginning_is_untouched(tmp_path):
    """No `-ss` at all, so no seek behaviour to change — and the fault never
    appeared there, which is why a film watched straight through looked fine."""
    clip = tmp_path / "longgop.mkv"
    _make_long_gop_clip(clip)

    out = await _stream_to_file(tmp_path, clip, 0, "zero")

    assert _largest_audio_gap(out) < 0.5


async def test_a_transcoded_video_keeps_accurate_seeking(tmp_path):
    """Re-encoded video can begin exactly where it was asked to, so the
    trimming that breaks a copy is right there. Turning it off everywhere
    would have cost seek precision on the one path that has it.
    """
    clip = tmp_path / "longgop.mkv"
    _make_long_gop_clip(clip)
    gek = generate_gek()
    session, file_id = _session(tmp_path, clip, gek)
    # Force the re-encode branch the way a browser-incompatible source would.
    session._ctx["transcode_incompatible_video"] = True
    import meshbay_node.transport.webrtc_server as ws
    real = ws.BROWSER_INCOMPATIBLE_VIDEO_CODECS
    ws.BROWSER_INCOMPATIBLE_VIDEO_CODECS = frozenset({"h264"})
    try:
        await session._stream_video_inner(
            {"file_id": file_id, "start": SEEK_S, "credits": 0})
    finally:
        ws.BROWSER_INCOMPATIBLE_VIDEO_CODECS = real

    assert not [m for m in session.sent if m.get("type") == "error"]
    out = tmp_path / "seek-transcoded.mp4"
    out.write_bytes(_reassemble(session.sent, gek, file_id))
    assert _largest_audio_gap(out) < 0.5