music-assistant-server

1"""
2Handler for chained OGG streams used in internet radio.
3
4FFmpeg cannot handle chained OGG streams (multiple logical bitstreams). This module
5stitches them into a single continuous stream by skipping EOS/BOS boundaries and
6re-sequencing page numbers.
7"""
8
9from __future__ import annotations
10
11import logging
12import struct
13from collections.abc import AsyncGenerator, Callable
14from typing import TYPE_CHECKING, Any
15
16import aiohttp
17from music_assistant_models.errors import ProviderUnavailableError
18
19from music_assistant.constants import MASS_LOGGER_NAME
20
21if TYPE_CHECKING:
22    from music_assistant.mass import MusicAssistant
23
24LOGGER = logging.getLogger(f"{MASS_LOGGER_NAME}.ogg_handler")
25
26# OGG constants
27OGG_SYNC_PATTERN: bytes = b"OggS"
28OGG_HEADER_SIZE: int = 27  # Fixed header size before segment table
29
30# Header type flags (in header_type byte at offset 5)
31OGG_FLAG_CONTINUATION: int = 0x01  # Continuation of previous page
32OGG_FLAG_BOS: int = 0x02  # Beginning of stream
33OGG_FLAG_EOS: int = 0x04  # End of stream
34
35# Ogg FLAC constants
36FLAC_METADATA_HEADER_SIZE: int = (
37    4  # Size of FLAC metadata block header (1 byte type + 3 bytes length)
38)
39FLAC_METADATA_BLOCK_VORBIS_COMMENT: int = 4  # FLAC metadata block type for Vorbis comments
40
41
42class OggPage:
43    """Parsed OGG page with header fields and payload."""
44
45    def __init__(
46        self,
47        raw_data: bytes,
48        header_type: int,
49        granule_position: int,
50        serial_number: int,
51        page_sequence: int,
52        segment_data: bytes,
53    ) -> None:
54        """Initialize OGG page."""
55        self.raw_data = raw_data
56        self.header_type = header_type
57        self.granule_position = granule_position
58        self.serial_number = serial_number
59        self.page_sequence = page_sequence
60        self.segment_data = segment_data
61
62    @property
63    def is_bos(self) -> bool:
64        """Return True if beginning of stream flag is set."""
65        return bool(self.header_type & OGG_FLAG_BOS)
66
67    @property
68    def is_eos(self) -> bool:
69        """Return True if end of stream flag is set."""
70        return bool(self.header_type & OGG_FLAG_EOS)
71
72    @property
73    def is_continuation(self) -> bool:
74        """Return True if continuation flag is set."""
75        return bool(self.header_type & OGG_FLAG_CONTINUATION)
76
77    @property
78    def is_opus_head(self) -> bool:
79        """Return True if page contains OpusHead header."""
80        return self.segment_data.startswith(b"OpusHead")
81
82    @property
83    def is_opus_tags(self) -> bool:
84        """Return True if page contains OpusTags header."""
85        return self.segment_data.startswith(b"OpusTags")
86
87    @property
88    def is_vorbis_id(self) -> bool:
89        """Return True if page contains Vorbis identification header."""
90        return len(self.segment_data) > 7 and self.segment_data[0:7] == b"\x01vorbis"
91
92    @property
93    def is_vorbis_comment(self) -> bool:
94        """Return True if page contains Vorbis comment header."""
95        return len(self.segment_data) > 7 and self.segment_data[0:7] == b"\x03vorbis"
96
97    @property
98    def is_ogg_flac_mapping_page(self) -> bool:
99        """Return True if page starts with the Ogg FLAC mapping header packet."""
100        return len(self.segment_data) > 5 and self.segment_data[0:5] == b"\x7fFLAC"
101
102    def is_header_page(self, is_ogg_flac_stream: bool = False) -> bool:
103        """Return True if page is a header (not audio data)."""
104        return (
105            self.is_opus_head
106            or self.is_opus_tags
107            or self.is_vorbis_id
108            or self.is_vorbis_comment
109            or self.is_ogg_flac_mapping_page
110            # In Ogg FLAC streams pages with granule position 0 contain header packets and not audio data
111            or (is_ogg_flac_stream and self.granule_position == 0)
112        )
113
114
115def parse_ogg_page(data: bytes | bytearray, offset: int = 0) -> tuple[OggPage, int] | None:
116    """Parse a single OGG page from buffer. Returns (OggPage, consumed) or None if incomplete."""
117    if len(data) < offset + OGG_HEADER_SIZE:
118        return None
119
120    if data[offset : offset + 4] != OGG_SYNC_PATTERN:
121        return None
122
123    header_type = data[offset + 5]
124    granule_position = struct.unpack_from("<Q", data, offset + 6)[0]
125    serial_number = struct.unpack_from("<I", data, offset + 14)[0]
126    page_sequence = struct.unpack_from("<I", data, offset + 18)[0]
127    num_segments = data[offset + 26]
128
129    header_size = OGG_HEADER_SIZE + num_segments
130    if len(data) < offset + header_size:
131        return None
132
133    segment_table = data[offset + OGG_HEADER_SIZE : offset + header_size]
134    segment_data_size = sum(segment_table)
135
136    total_page_size = header_size + segment_data_size
137    if len(data) < offset + total_page_size:
138        return None
139
140    segment_data = bytes(data[offset + header_size : offset + total_page_size])
141    raw_data = bytes(data[offset : offset + total_page_size])
142
143    page = OggPage(
144        raw_data=raw_data,
145        header_type=header_type,
146        granule_position=granule_position,
147        serial_number=serial_number,
148        page_sequence=page_sequence,
149        segment_data=segment_data,
150    )
151
152    return (page, offset + total_page_size)
153
154
155_OGG_CRC_TABLE: list[int] = []
156
157
158def _build_ogg_crc_table() -> list[int]:
159    """Build OGG CRC32 lookup table (polynomial 0x04c11db7)."""
160    table: list[int] = []
161    poly = 0x04C11DB7
162    for i in range(256):
163        crc = i << 24
164        for _ in range(8):
165            crc = ((crc << 1) ^ poly) & 0xFFFFFFFF if crc & 0x80000000 else (crc << 1) & 0xFFFFFFFF
166        table.append(crc)
167    return table
168
169
170_OGG_CRC_TABLE = _build_ogg_crc_table()
171
172
173def calculate_ogg_crc(data: bytes) -> int:
174    """Calculate OGG CRC32 checksum for page data (with checksum field zeroed)."""
175    crc = 0
176    for byte in data:
177        crc = ((crc << 8) ^ _OGG_CRC_TABLE[((crc >> 24) ^ byte) & 0xFF]) & 0xFFFFFFFF
178    return crc
179
180
181def rewrite_ogg_page(
182    page: OggPage,
183    new_serial: int | None = None,
184    new_sequence: int | None = None,
185    new_granule: int | None = None,
186    clear_bos: bool = False,
187) -> bytes:
188    """Rewrite an OGG page with modified header fields and recalculated CRC."""
189    data = bytearray(page.raw_data)
190
191    if clear_bos:
192        data[5] = page.header_type & ~OGG_FLAG_BOS
193    if new_granule is not None:
194        struct.pack_into("<Q", data, 6, new_granule)
195    if new_serial is not None:
196        struct.pack_into("<I", data, 14, new_serial)
197    if new_sequence is not None:
198        struct.pack_into("<I", data, 18, new_sequence)
199
200    data[22:26] = b"\x00\x00\x00\x00"
201    crc = calculate_ogg_crc(bytes(data))
202    struct.pack_into("<I", data, 22, crc)
203
204    return bytes(data)
205
206
207def parse_vorbis_comments(data: bytes) -> dict[str, str]:
208    """Parse Vorbis comments structure. Data should exclude magic header bytes."""
209    comments: dict[str, str] = {}
210    try:
211        offset = 0
212        if len(data) < 4:
213            return comments
214        vendor_length = struct.unpack_from("<I", data, offset)[0]
215        offset += 4 + vendor_length
216
217        if len(data) < offset + 4:
218            return comments
219        num_comments = struct.unpack_from("<I", data, offset)[0]
220        offset += 4
221
222        for _ in range(num_comments):
223            if len(data) < offset + 4:
224                break
225            comment_length = struct.unpack_from("<I", data, offset)[0]
226            offset += 4
227            if len(data) < offset + comment_length:
228                break
229            comment_bytes = data[offset : offset + comment_length]
230            offset += comment_length
231            try:
232                comment_str = comment_bytes.decode("utf-8")
233                if "=" in comment_str:
234                    key, value = comment_str.split("=", 1)
235                    comments[key.lower()] = value
236            except UnicodeDecodeError:
237                continue
238    except struct.error, IndexError:
239        pass
240    return comments
241
242
243def parse_flac_vorbis_comment_block(data: bytes) -> dict[str, str]:
244    """Parse Vorbis comments from a native FLAC metadata block."""
245    if len(data) < FLAC_METADATA_HEADER_SIZE:
246        return {}
247    block_length = int.from_bytes(data[1:4], byteorder="big")
248    if len(data) < FLAC_METADATA_HEADER_SIZE + block_length:
249        LOGGER.debug(
250            "Skipping FLAC Vorbis comment block spanning multiple OGG pages: "
251            "need %d bytes, have %d",
252            FLAC_METADATA_HEADER_SIZE + block_length,
253            len(data),
254        )
255        return {}
256    comment_data = data[FLAC_METADATA_HEADER_SIZE : FLAC_METADATA_HEADER_SIZE + block_length]
257    return parse_vorbis_comments(comment_data)
258
259
260def _is_flac_vorbis_comment_block(data: bytes) -> bool:
261    """Return True if Ogg FLAC packet data contains a FLAC Vorbis comment block."""
262    if len(data) < FLAC_METADATA_HEADER_SIZE:
263        return False
264    block_type = data[0] & 0x7F
265    return block_type == FLAC_METADATA_BLOCK_VORBIS_COMMENT
266
267
268def extract_metadata_from_page(
269    page: OggPage, is_ogg_flac_stream: bool = False
270) -> dict[str, str] | None:
271    """Extract metadata from page if it contains supported comment metadata."""
272    if page.is_opus_tags:
273        return parse_vorbis_comments(page.segment_data[8:])
274    if page.is_vorbis_comment:
275        return parse_vorbis_comments(page.segment_data[7:])
276    if is_ogg_flac_stream and _is_flac_vorbis_comment_block(page.segment_data):
277        return parse_flac_vorbis_comment_block(page.segment_data)
278    return None
279
280
281class _ChainedOggState:
282    """State machine for stitching chained OGG streams."""
283
284    def __init__(self, metadata_callback: Callable[[dict[str, str]], Any] | None = None) -> None:
285        self.metadata_callback = metadata_callback
286        self.output_serial: int | None = None
287        self.output_sequence: int = 0
288        self.first_chain: bool = True
289        self.seen_eos: bool = False
290        self.is_ogg_flac_chain: bool = False
291        self.header_pages_sent: int = 0
292        self.last_granule: int = 0
293        self.granule_offset: int = 0
294
295    def process_page(self, page: OggPage) -> bytes | None:
296        """Process page. Returns data to yield or None to skip."""
297        if self.first_chain:
298            return self._process_first_chain_page(page)
299        return self._process_chain_page(page)
300
301    def _handle_metadata(self, page: OggPage) -> None:
302        """Extract and invoke callback for supported in-band metadata pages."""
303        if self.metadata_callback:
304            metadata = extract_metadata_from_page(page, is_ogg_flac_stream=self.is_ogg_flac_chain)
305            if metadata:
306                LOGGER.debug("Extracted metadata: %s", metadata)
307                self.metadata_callback(metadata)
308
309    def _process_first_chain_page(self, page: OggPage) -> bytes | None:
310        """Process page from first chain. Returns data to yield or None to skip."""
311        if page.is_bos:
312            self.output_serial = page.serial_number
313            LOGGER.debug("First chain BOS, serial=%d", self.output_serial)
314            if page.is_ogg_flac_mapping_page:
315                self.is_ogg_flac_chain = True
316            self.output_sequence = page.page_sequence
317            self.header_pages_sent = 1
318            return page.raw_data
319
320        # Ogg FLAC chains can carry additional header pages after the mapping header,
321        # so do not stop at the two-page limit used for Opus/Vorbis setup headers.
322        if page.is_header_page(self.is_ogg_flac_chain) and (
323            self.header_pages_sent < 2 or self.is_ogg_flac_chain
324        ):
325            LOGGER.debug("First chain header page %d", self.header_pages_sent)
326            self._handle_metadata(page)
327            self.output_sequence = page.page_sequence
328            self.header_pages_sent += 1
329            return page.raw_data
330
331        if page.granule_position != 0xFFFFFFFFFFFFFFFF:
332            self.last_granule = page.granule_position
333
334        if page.is_eos:
335            # Skip EOS - FFmpeg cannot handle them
336            LOGGER.debug(
337                "First chain EOS at seq %d, granule %d (skipping)",
338                page.page_sequence,
339                self.last_granule,
340            )
341            self.granule_offset = self.last_granule
342            self.first_chain = False
343            self.seen_eos = True
344            self.is_ogg_flac_chain = False
345            return None
346
347        self.output_sequence += 1
348        if page.page_sequence != self.output_sequence:
349            return rewrite_ogg_page(page, new_sequence=self.output_sequence)
350        return page.raw_data
351
352    def _process_chain_page(self, page: OggPage) -> bytes | None:
353        """Process page from subsequent chains. Returns data to yield or None to skip."""
354        if self.seen_eos and page.is_bos:
355            LOGGER.debug("New chain BOS, serial=%d (skipping)", page.serial_number)
356            self.seen_eos = False
357            self.is_ogg_flac_chain = False
358            if page.is_ogg_flac_mapping_page:
359                self.is_ogg_flac_chain = True
360            return None
361
362        if page.is_header_page(self.is_ogg_flac_chain):
363            LOGGER.debug("Chain header page (skipping)")
364            self._handle_metadata(page)
365            return None
366
367        if page.granule_position != 0xFFFFFFFFFFFFFFFF:
368            self.last_granule = page.granule_position + self.granule_offset
369
370        if page.is_eos:
371            # Skip EOS - FFmpeg cannot handle them
372            LOGGER.debug(
373                "Chain EOS at seq %d, new offset %d (skipping)",
374                page.page_sequence,
375                self.last_granule,
376            )
377            self.granule_offset = self.last_granule
378            self.seen_eos = True
379            return None
380
381        new_granule: int | None = None
382        if page.granule_position != 0xFFFFFFFFFFFFFFFF:
383            new_granule = page.granule_position + self.granule_offset
384
385        self.output_sequence += 1
386        return rewrite_ogg_page(
387            page,
388            new_serial=self.output_serial,
389            new_sequence=self.output_sequence,
390            new_granule=new_granule,
391            clear_bos=page.is_bos,
392        )
393
394
395def _resync_ogg_buffer(buffer: bytearray) -> int:
396    """Find next OGG sync pattern, returning bytes to skip (0 if already synced)."""
397    idx = buffer.find(OGG_SYNC_PATTERN, 1)
398    if idx > 0:
399        LOGGER.warning("Skipping %d bytes of corrupted OGG data", idx)
400        return idx
401    return 0
402
403
404_MAX_BUFFER_SIZE = 65536
405
406
407async def get_chained_ogg_stream(
408    mass: MusicAssistant,
409    url: str,
410    metadata_callback: Callable[[dict[str, str]], Any] | None = None,
411) -> AsyncGenerator[bytes]:
412    """
413    Yield continuous OGG data from a chained stream, stitching chain boundaries.
414
415    :param mass: MusicAssistant instance.
416    :param url: URL of the OGG radio stream.
417    :param metadata_callback: Optional callback invoked on metadata changes.
418    """
419    state = _ChainedOggState(metadata_callback)
420    buffer = bytearray()
421
422    LOGGER.debug("Starting chained OGG stream handler for %s", url)
423
424    try:
425        async for chunk in mass.streams.audio.get_reconnecting_radio_stream(url):
426            buffer.extend(chunk)
427
428            while True:
429                result = parse_ogg_page(buffer, 0)
430                if result is None:
431                    if len(buffer) > _MAX_BUFFER_SIZE:
432                        skip = _resync_ogg_buffer(buffer)
433                        if skip > 0:
434                            buffer = buffer[skip:]
435                            continue
436                        discard = len(buffer) // 2
437                        LOGGER.warning("Buffer overflow, discarding %d bytes", discard)
438                        buffer = buffer[discard:]
439                    break
440
441                page, consumed = result
442                buffer = buffer[consumed:]
443
444                output = state.process_page(page)
445                if output is not None:
446                    yield output
447    except aiohttp.ClientError as err:
448        raise ProviderUnavailableError(f"Failed to fetch OGG stream: {err}") from err
449
450    LOGGER.debug("Chained OGG stream handler ended for %s", url)
451