From 26fb229264fb02bef277e0716f923a79bd2d111c Mon Sep 17 00:00:00 2001
From: coletdjnz <coletdjnz@protonmail.com>
Date: Thu, 26 Dec 2024 21:32:46 +1300
Subject: [PATCH] feat(sabr) add SabrContextUpdate pb

---
 utils/mitmproxy_sabrdump.py                     |  8 +++++++-
 utils/read_sabr_response.py                     |  8 +++++++-
 .../extractor/_ytse/downloader/sabr.py          |  2 ++
 .../extractor/_ytse/protos/__init__.py          |  1 +
 .../_ytse/protos/_sabr_context_update.py        | 17 +++++++++++++++++
 5 files changed, 34 insertions(+), 2 deletions(-)
 create mode 100644 yt_dlp_plugins/extractor/_ytse/protos/_sabr_context_update.py

diff --git a/utils/mitmproxy_sabrdump.py b/utils/mitmproxy_sabrdump.py
index d1ddbbe..6dca2b6 100644
--- 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
index bf58aec..614dccd 100644
--- 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
index d7c838d..6b08fc4 100644
--- 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
index 2453a31..01f1dc0 100644
--- 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
new file mode 100644
index 0000000..7f8e621
--- /dev/null
+++ 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