feat(sabr) add SabrContextUpdate pb

coletdjnz · Dec 26, 2024 · 26fb229 · 26fb229
1 parent 89a8d91
commit 26fb229
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 2 deletions.
diff --git a/utils/mitmproxy_sabrdump.py b/utils/mitmproxy_sabrdump.py
@@ -22,7 +22,8 @@
     unknown_fields,
     SelectableFormats,
     PrewarmConnection,
-    AllowedCachedFormats
+    AllowedCachedFormats,
+    SabrContextUpdate
 )
 from yt_dlp_plugins.extractor._ytse.ump import UMPPartType
 
@@ -119,6 +120,11 @@ def response(self, flow: http.HTTPFlow) -> None:
                         f.write(f'Allowed Cached Formats: {acf}\n')
                         write_unknown_fields(f, acf)
 
+                    elif part.part_type == UMPPartType.SABR_CONTEXT_UPDATE:
+                        scu = protobug.loads(part.data, SabrContextUpdate)
+                        f.write(f'Sabr Context Update: {scu}\n')
+                        write_unknown_fields(f, scu)
+
                     elif part.part_type == UMPPartType.MEDIA or part.part_type == UMPPartType.MEDIA_END:
                         f.write(f'Media Header Id: {part.data[0]}\n')
 

diff --git a/utils/read_sabr_response.py b/utils/read_sabr_response.py
@@ -24,7 +24,8 @@
     unknown_fields,
     SelectableFormats,
     PrewarmConnection,
-    AllowedCachedFormats
+    AllowedCachedFormats,
+    SabrContextUpdate
 )
 from yt_dlp_plugins.extractor._ytse.ump import UMPPartType
 
@@ -106,6 +107,11 @@ def print_sabr_parts(fp):
             print(f'Allowed Cached Formats: {acf}')
             write_unknown_fields(f, acf)
 
+        elif part.part_type == UMPPartType.SABR_CONTEXT_UPDATE:
+            scu = protobug.loads(part.data, SabrContextUpdate)
+            print(f'Sabr Context Update: {scu}')
+            write_unknown_fields(f, scu)
+
         elif part.part_type == UMPPartType.MEDIA or part.part_type == UMPPartType.MEDIA_END:
             print(f'Media Header Id: {part.data[0]}')
 

diff --git a/yt_dlp_plugins/extractor/_ytse/downloader/sabr.py b/yt_dlp_plugins/extractor/_ytse/downloader/sabr.py
@@ -291,6 +291,7 @@ def parse_ump_response(self, response):
                 self.process_sabr_seek(part)
             elif part.part_type == UMPPartType.SABR_ERROR:
                 self.process_sabr_error(part)
+            # TODO: implement SABR_CONTEXT_UPDATE
             else:
                 self.write_sabr_debug(f'Unhandled part type', part=part, data=part.data)
                 continue
@@ -372,6 +373,7 @@ def process_media_header(self, part: UMPPart):
             else:
                 # Attempt to keep in sync with livestream, as the segment duration target is not always perfect.
                 # The server seems to care more about the segment index than the duration.
+                # TODO: SABR_CONTEXT_UPDATE should be used to update the buffered range
                 if current_buffered_range.start_time_ms > start_ms:
                     raise DownloadError(f'Buffered range start time mismatch: {current_buffered_range.start_time_ms} > {start_ms}')
                 new_duration = (start_ms - current_buffered_range.start_time_ms) + duration_ms

diff --git a/yt_dlp_plugins/extractor/_ytse/protos/__init__.py b/yt_dlp_plugins/extractor/_ytse/protos/__init__.py
@@ -21,6 +21,7 @@
 from ._selectable_formats import SelectableFormats
 from ._prewarm_connection import PrewarmConnection
 from ._allowed_cached_formats import AllowedCachedFormats
+from ._sabr_context_update import SabrContextUpdate
 
 
 def unknown_fields(obj: typing.Any, path=()) -> typing.Iterable[tuple[tuple[str, ...], dict[int, list]]]:

diff --git a/yt_dlp_plugins/extractor/_ytse/protos/_sabr_context_update.py b/yt_dlp_plugins/extractor/_ytse/protos/_sabr_context_update.py
@@ -0,0 +1,17 @@
+import typing
+import protobug
+from yt_dlp_plugins.extractor._ytse.protos import BufferedRange
+
+
+@protobug.message
+class ContextUpdate:
+    # At least appears to match BufferedRange pb...
+    buffered_ranges: list[BufferedRange] = protobug.field(1, default_factory=list)
+
+
+@protobug.message
+class SabrContextUpdate:
+    unknown_field_1: typing.Optional[protobug.Int32] = protobug.field(1, default=None)  # seen = 2
+    unknown_field_2: typing.Optional[protobug.Int32] = protobug.field(2, default=None)  # seen = 2
+    context_update: ContextUpdate = protobug.field(3, default_factory=ContextUpdate),
+    unknown_field_4: typing.Optional[protobug.Int32] = protobug.field(4, default=None)  # seen = 1