fix(youtube): respect active SABR tracks on cold start

extractor/src/main/java/org/schabi/newpipe/extractor/services/youtube/sabr/YoutubeSabrRequestBuilder.java

@@ -34,7 +34,10 @@ static byte[] buildFirstMediaRequest(@Nonnull final YoutubeSabrInfo info,
 
         final SabrProto.Writer request = new SabrProto.Writer();
         request.writeMessage(1, buildClientAbrState(audioFormat, videoFormat, 0, false,
-                ENABLED_TRACK_TYPES_VIDEO_AND_AUDIO, streamState));
+                streamState == null
+                        ? ENABLED_TRACK_TYPES_VIDEO_AND_AUDIO
+                        : streamState.getEnabledTrackTypesBitfield(),
+                streamState));
         request.writeBytes(5, decodeBase64(ustreamerConfig));
         writePreferredFormats(request, info, audioFormat, videoFormat, streamState);
         request.writeMessage(19, streamState == null
@@ -221,19 +224,23 @@ private static void writePreferredFormats(@Nonnull final SabrProto.Writer reques
                 return;
             }
             for (final YoutubeSabrFormat format : info.getFormats()) {
-                if (format.isAudio()) {
+                if (format.isAudio() && streamState.shouldSelectAudioFormat()) {
                     request.writeMessage(16, SabrProto.formatId(format));
                 }
             }
             for (final YoutubeSabrFormat format : info.getFormats()) {
-                if (format.isVideo()) {
+                if (format.isVideo() && streamState.shouldSelectVideoFormat()) {
                     request.writeMessage(17, SabrProto.formatId(format));
                 }
             }
             return;
         }
-        request.writeMessage(16, SabrProto.formatId(audioFormat));
-        request.writeMessage(17, SabrProto.formatId(videoFormat));
+        if (streamState == null || streamState.shouldSelectAudioFormat()) {
+            request.writeMessage(16, SabrProto.formatId(audioFormat));
+        }
+        if (streamState == null || streamState.shouldSelectVideoFormat()) {
+            request.writeMessage(17, SabrProto.formatId(videoFormat));
+        }
     }
 
     private static void writeOfficialWebPreferredFormats(@Nonnull final SabrProto.Writer request,

extractor/src/main/java/org/schabi/newpipe/extractor/services/youtube/sabr/YoutubeSabrSession.java

@@ -23,6 +23,8 @@ public final class YoutubeSabrSession {
     // server can ask us to reload the player response (URLs/config expired on a long watch). re-probe
     // and resume in place instead of killing the session. bounded so a reload loop can't run forever.
     private static final int MAX_RELOADS_PER_SESSION = 2;
+    private static final int INTEGRITY_RELOAD_AFTER_FAILURES = 3;
+    private static final int MAX_RECOVERABLE_INTEGRITY_FAILURES = 8;
     // How many times a stale/rejected PO token may be force-re-minted before giving up (token
     // expiry mid-playback). Bounded so a genuinely-rejected token can't loop forever.
     private static final int MAX_PO_TOKEN_REFRESHES = 2;
@@ -48,6 +50,7 @@ public final class YoutubeSabrSession {
     private int redirectCount;
     private int poTokenRefreshes;
     private int reloads;
+    private int consecutiveIntegrityFailures;
     // Insertion order + total bytes of cached MEDIA segments (init segments are never evicted).
     // Mutated only by the single pump thread in pumpOnce; readers only do concurrent-map gets.
     private final Deque<String> cacheOrder = new ArrayDeque<>();
@@ -111,7 +114,15 @@ public SabrMediaSegment fetchSegment(@Nonnull final SabrSegmentRequest request,
         for (int attempts = 0; attempts < MAX_REQUESTS_PER_SEGMENT; attempts++) {
             final YoutubeSabrProbeResult result = fetchNextResponse(localization);
             final SabrDecodedResponse decoded = result.getDecodedResponse();
-            validateResponseIntegrity(decoded, request);
+            final List<String> integrityIssues = decoded.getIntegrityIssues();
+            if (!integrityIssues.isEmpty()) {
+                if (recoverFromIncompleteMediaResponse(localization, integrityIssues)) {
+                    continue;
+                }
+                throw new SabrProtocolException("SABR media integrity issue while fetching "
+                        + describeRequest(request) + ": " + integrityIssues);
+            }
+            consecutiveIntegrityFailures = 0;
             streamState.ingest(decoded);
             final List<SabrMediaSegment> segments = result.getSegments();
             for (final SabrMediaSegment segment : segments) {
@@ -239,8 +250,12 @@ public List<SabrMediaSegment> pumpOnce(@Nonnull final Localization localization)
         final SabrDecodedResponse decoded = result.getDecodedResponse();
         final List<String> integrityIssues = decoded.getIntegrityIssues();
         if (!integrityIssues.isEmpty()) {
+            if (recoverFromIncompleteMediaResponse(localization, integrityIssues)) {
+                return Collections.emptyList();
+            }
             throw new SabrProtocolException("SABR media integrity issue: " + integrityIssues);
         }
+        consecutiveIntegrityFailures = 0;
         streamState.ingest(decoded);
         final List<SabrMediaSegment> segments = result.getSegments();
         for (final SabrMediaSegment segment : segments) {
@@ -495,16 +510,6 @@ private void failIfKnownOutOfBounds(@Nonnull final SabrSegmentRequest request)
         }
     }
 
-    private void validateResponseIntegrity(@Nonnull final SabrDecodedResponse decoded,
-                                           @Nonnull final SabrSegmentRequest request)
-            throws SabrProtocolException {
-        final List<String> integrityIssues = decoded.getIntegrityIssues();
-        if (!integrityIssues.isEmpty()) {
-            throw new SabrProtocolException("SABR media integrity issue while fetching "
-                    + describeRequest(request) + ": " + integrityIssues);
-        }
-    }
-
     private boolean maybePrepareForDistantMediaSegment(
             @Nonnull final SabrSegmentRequest request) {
         if (request.isInitializationSegment() || requestNumber == 0) {
@@ -521,6 +526,40 @@ private boolean maybePrepareForDistantMediaSegment(
         return true;
     }
 
+    private boolean recoverFromIncompleteMediaResponse(@Nonnull final Localization localization,
+                                                       @Nonnull final List<String> integrityIssues)
+            throws IOException, ExtractionException {
+        if (!isRecoverableIncompleteMediaResponse(integrityIssues)) {
+            return false;
+        }
+        consecutiveIntegrityFailures++;
+        if (consecutiveIntegrityFailures > MAX_RECOVERABLE_INTEGRITY_FAILURES) {
+            return false;
+        }
+        if (consecutiveIntegrityFailures >= INTEGRITY_RELOAD_AFTER_FAILURES
+                && maybeReload(localization)) {
+            consecutiveIntegrityFailures = 0;
+            return true;
+        }
+        sleepBackoff(Math.min(MAX_BACKOFF_MS, 500 * consecutiveIntegrityFailures));
+        return true;
+    }
+
+    private static boolean isRecoverableIncompleteMediaResponse(
+            @Nonnull final List<String> integrityIssues) {
+        if (integrityIssues.isEmpty()) {
+            return false;
+        }
+        for (final String issue : integrityIssues) {
+            if (!issue.startsWith("length-mismatch:")
+                    && !issue.startsWith("missing-media-end:")
+                    && !issue.startsWith("missing-media:")) {
+                return false;
+            }
+        }
+        return true;
+    }
+
     private static void sleepBackoff(final int backoffTimeMs) throws SabrProtocolException {
         // Clamp to [0, MAX_BACKOFF_MS]: a negative (overflowed varint) must not skip the wait, and
         // a huge server backoff must not be honoured verbatim (would stall playback for minutes).