feat(youtube): support live streams through SABR

extractor/src/main/java/org/schabi/newpipe/extractor/services/youtube/extractors/YoutubeStreamExtractor.java

@@ -822,19 +822,18 @@ private void ensureStreamsAreCached() throws ExtractionException {
             cachedVideoStreams = new ArrayList<>();
             cachedVideoOnlyStreams = new ArrayList<>();
             final String selectedClient = NewPipe.getYoutubePlayerClient();
-            if (("mweb".equals(selectedClient) || "web".equals(selectedClient))
-                    && streamType != StreamType.LIVE_STREAM
-                    && streamType != StreamType.POST_LIVE_STREAM
-                    && hasSabrStreamingUrl()) {
+            final boolean useSabr = ("mweb".equals(selectedClient)
+                    || "web".equals(selectedClient)) && hasSabrStreamingUrl();
+            if (useSabr) {
                 buildSabrStreams(videoId);
             } else if (!("tv_downgraded".equals(selectedClient)
                     && streamType == StreamType.LIVE_STREAM)) {
                 extractAdaptiveFormats(videoId);
             }
-            if (streamType == StreamType.POST_LIVE_STREAM
+            if (!useSabr && (streamType == StreamType.POST_LIVE_STREAM
                     || (streamType == StreamType.LIVE_STREAM
                         && "tv_downgraded".equals(selectedClient))
-                    || "web_safari".equals(selectedClient)) {
+                    || "web_safari".equals(selectedClient))) {
                 tryExtractHlsStreams(videoId);
             }
             streamsCached = true;

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

@@ -0,0 +1,170 @@
+package org.schabi.newpipe.extractor.services.youtube.sabr;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Locale;
+
+final class SabrMediaDataNormalizer {
+    private static final int MP4_HEADER_SIZE = 8;
+    private static final int MP4_EXTENDED_HEADER_SIZE = 16;
+    private static final long EBML_HEADER_ID = 0x1A45DFA3L;
+    private static final long WEBM_SEGMENT_ID = 0x18538067L;
+    private static final long WEBM_CLUSTER_ID = 0x1F43B675L;
+
+    private SabrMediaDataNormalizer() {
+    }
+
+    @Nullable
+    static SabrMediaDataParts split(@Nullable final String mimeType,
+                                    @Nonnull final byte[] data) {
+        if (mimeType == null) {
+            return null;
+        }
+        final String container = mimeType.split(";", 2)[0].trim().toLowerCase(Locale.ROOT);
+        if (container.endsWith("/mp4")) {
+            return splitMp4(data);
+        }
+        if (container.endsWith("/webm")) {
+            return splitWebm(data);
+        }
+        return null;
+    }
+
+    @Nullable
+    private static SabrMediaDataParts splitMp4(@Nonnull final byte[] data) {
+        int offset = 0;
+        int initializationEnd = -1;
+        while (offset + MP4_HEADER_SIZE <= data.length) {
+            final long size32 = readUnsignedInt(data, offset);
+            final String type = new String(data, offset + 4, 4, StandardCharsets.US_ASCII);
+            final int headerSize = size32 == 1 ? MP4_EXTENDED_HEADER_SIZE : MP4_HEADER_SIZE;
+            if (offset + headerSize > data.length) {
+                return null;
+            }
+            final long size = size32 == 0 ? data.length - (long) offset
+                    : size32 == 1 ? readUnsignedLong(data, offset + MP4_HEADER_SIZE) : size32;
+            if (size < headerSize || size > data.length - (long) offset) {
+                return null;
+            }
+            final int end = offset + (int) size;
+            if ("moov".equals(type)) {
+                initializationEnd = end;
+            } else if ("moof".equals(type) && initializationEnd > 0) {
+                return splitAt(data, initializationEnd);
+            }
+            offset = end;
+        }
+        return null;
+    }
+
+    @Nullable
+    private static SabrMediaDataParts splitWebm(@Nonnull final byte[] data) {
+        final EbmlElement ebml = readEbmlElement(data, 0);
+        if (ebml == null || ebml.id != EBML_HEADER_ID || ebml.size < 0) {
+            return null;
+        }
+        final long segmentOffset = ebml.payloadOffset + ebml.size;
+        if (segmentOffset > Integer.MAX_VALUE) {
+            return null;
+        }
+        final EbmlElement segment = readEbmlElement(data, (int) segmentOffset);
+        if (segment == null || segment.id != WEBM_SEGMENT_ID) {
+            return null;
+        }
+        int offset = segment.payloadOffset;
+        while (offset < data.length) {
+            final EbmlElement element = readEbmlElement(data, offset);
+            if (element == null) {
+                return null;
+            }
+            if (element.id == WEBM_CLUSTER_ID) {
+                return offset <= segment.payloadOffset ? null : splitAt(data, offset);
+            }
+            if (element.size < 0) {
+                return null;
+            }
+            final long next = element.payloadOffset + element.size;
+            if (next <= offset || next > data.length) {
+                return null;
+            }
+            offset = (int) next;
+        }
+        return null;
+    }
+
+    @Nullable
+    private static EbmlElement readEbmlElement(@Nonnull final byte[] data, final int offset) {
+        final int idLength = vintLength(data, offset, 4);
+        if (idLength < 0) {
+            return null;
+        }
+        final long id = readRawValue(data, offset, idLength);
+        final int sizeOffset = offset + idLength;
+        final int sizeLength = vintLength(data, sizeOffset, 8);
+        if (sizeLength < 0) {
+            return null;
+        }
+        final long rawSize = readRawValue(data, sizeOffset, sizeLength);
+        final long marker = 1L << (7 * sizeLength);
+        final long sizeValue = rawSize & (marker - 1);
+        final long size = sizeValue == marker - 1 ? -1 : sizeValue;
+        return new EbmlElement(id, size, sizeOffset + sizeLength);
+    }
+
+    private static int vintLength(@Nonnull final byte[] data, final int offset,
+                                  final int maximum) {
+        if (offset < 0 || offset >= data.length) {
+            return -1;
+        }
+        final int first = data[offset] & 0xFF;
+        if (first == 0) {
+            return -1;
+        }
+        final int length = Integer.numberOfLeadingZeros(first) - 23;
+        return length >= 1 && length <= maximum && offset + length <= data.length ? length : -1;
+    }
+
+    private static long readUnsignedInt(@Nonnull final byte[] data, final int offset) {
+        return (long) (data[offset] & 0xFF) << 24
+                | (long) (data[offset + 1] & 0xFF) << 16
+                | (long) (data[offset + 2] & 0xFF) << 8
+                | data[offset + 3] & 0xFFL;
+    }
+
+    private static long readUnsignedLong(@Nonnull final byte[] data, final int offset) {
+        final long value = readRawValue(data, offset, 8);
+        return value < 0 ? -1 : value;
+    }
+
+    private static long readRawValue(@Nonnull final byte[] data, final int offset,
+                                     final int length) {
+        if (length < 1 || length > 8 || offset < 0 || offset + length > data.length) {
+            return -1;
+        }
+        long value = 0;
+        for (int i = 0; i < length; i++) {
+            value = (value << 8) | (data[offset + i] & 0xFFL);
+        }
+        return value;
+    }
+
+    @Nonnull
+    private static SabrMediaDataParts splitAt(@Nonnull final byte[] data, final int offset) {
+        return new SabrMediaDataParts(Arrays.copyOfRange(data, 0, offset),
+                Arrays.copyOfRange(data, offset, data.length));
+    }
+
+    private static final class EbmlElement {
+        private final long id;
+        private final long size;
+        private final int payloadOffset;
+
+        private EbmlElement(final long id, final long size, final int payloadOffset) {
+            this.id = id;
+            this.size = size;
+            this.payloadOffset = payloadOffset;
+        }
+    }
+}

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

@@ -0,0 +1,26 @@
+package org.schabi.newpipe.extractor.services.youtube.sabr;
+
+import javax.annotation.Nonnull;
+
+final class SabrMediaDataParts {
+    @Nonnull
+    private final byte[] initializationData;
+    @Nonnull
+    private final byte[] mediaData;
+
+    SabrMediaDataParts(@Nonnull final byte[] initializationData,
+                       @Nonnull final byte[] mediaData) {
+        this.initializationData = initializationData;
+        this.mediaData = mediaData;
+    }
+
+    @Nonnull
+    byte[] getInitializationData() {
+        return initializationData;
+    }
+
+    @Nonnull
+    byte[] getMediaData() {
+        return mediaData;
+    }
+}

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

@@ -181,6 +181,26 @@ static SabrMediaHeader decode(@Nonnull final byte[] data) throws SabrProtocolExc
                 sequenceLastModified);
     }
 
+    @Nonnull
+    static SabrMediaHeader initializationFrom(@Nonnull final SabrMediaHeader source,
+                                              final int contentLength) {
+        return new SabrMediaHeader(source.headerId, source.videoId, source.itag,
+                source.lastModified, source.xtags, source.startRange,
+                source.compressionAlgorithm, true, -1, source.bitrateBps, 0, 0,
+                contentLength, 0, 0, source.timeRangeTimescale, source.sequenceLastModified);
+    }
+
+    @Nonnull
+    static SabrMediaHeader mediaFrom(@Nonnull final SabrMediaHeader source,
+                                     final int contentLength) {
+        return new SabrMediaHeader(source.headerId, source.videoId, source.itag,
+                source.lastModified, source.xtags, source.startRange,
+                source.compressionAlgorithm, false, source.sequenceNumber, source.bitrateBps,
+                source.startMs, source.durationMs, contentLength, source.timeRangeStartTicks,
+                source.timeRangeDurationTicks, source.timeRangeTimescale,
+                source.sequenceLastModified);
+    }
+
     @Nonnull
     private static FormatId decodeFormatId(@Nonnull final byte[] data) throws SabrProtocolException {
         int itag = -1;

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

@@ -145,6 +145,10 @@ public SabrMediaSegment onMediaEnd(@Nonnull final byte[] partData)
             return null;
         }
 
+        public boolean hasOpenSegments() {
+            return !openSegments.isEmpty();
+        }
+
         public void abort() {
             for (final OpenSegment segment : openSegments.values()) {
                 segment.abort();

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

@@ -42,6 +42,11 @@ public interface StoppableSegmentConsumer {
         boolean accept(@Nonnull SabrMediaSegment segment) throws SabrProtocolException;
     }
 
+    @FunctionalInterface
+    public interface LiveMetadataConsumer {
+        void accept(@Nonnull SabrLiveMetadata metadata) throws SabrProtocolException;
+    }
+
     /**
      * Streams completed segments directly to {@code segmentConsumer}. When a consumer is supplied,
      * completed segments are not retained by the result, bounding the response reader to one open
@@ -78,7 +83,7 @@ public static Result readUntil(@Nonnull final InputStream in,
                                    @Nullable final SegmentConsumer segmentStartConsumer,
                                    @Nullable final File spoolDirectory)
             throws SabrProtocolException, IOException {
-        return readUntil(in, segmentConsumer, segmentStartConsumer, spoolDirectory,
+        return readUntil(in, segmentConsumer, segmentStartConsumer, null, spoolDirectory,
                 SabrMediaProtocol.builtin());
     }
 
@@ -89,6 +94,29 @@ public static Result readUntil(@Nonnull final InputStream in,
                                    @Nullable final File spoolDirectory,
                                    @Nonnull final SabrMediaProtocol mediaProtocol)
             throws SabrProtocolException, IOException {
+        return readUntil(in, segmentConsumer, segmentStartConsumer, null, spoolDirectory,
+                mediaProtocol);
+    }
+
+    @Nonnull
+    public static Result readUntil(@Nonnull final InputStream in,
+                                   final StoppableSegmentConsumer segmentConsumer,
+                                   @Nullable final SegmentConsumer segmentStartConsumer,
+                                   @Nullable final LiveMetadataConsumer liveMetadataConsumer,
+                                   @Nullable final File spoolDirectory)
+            throws SabrProtocolException, IOException {
+        return readUntil(in, segmentConsumer, segmentStartConsumer, liveMetadataConsumer,
+                spoolDirectory, SabrMediaProtocol.builtin());
+    }
+
+    @Nonnull
+    public static Result readUntil(@Nonnull final InputStream in,
+                                   final StoppableSegmentConsumer segmentConsumer,
+                                   @Nullable final SegmentConsumer segmentStartConsumer,
+                                   @Nullable final LiveMetadataConsumer liveMetadataConsumer,
+                                   @Nullable final File spoolDirectory,
+                                   @Nonnull final SabrMediaProtocol mediaProtocol)
+            throws SabrProtocolException, IOException {
         final List<UmpPart> controlParts = new ArrayList<>();
         final List<String> partSummaries = new ArrayList<>();
         final List<SabrMediaSegment> segments = new ArrayList<>();
@@ -100,6 +128,7 @@ public static Result readUntil(@Nonnull final InputStream in,
         final long[] maxPartBytes = {0};
         final long[] maxMediaPartPayloadBytes = {0};
         final long[] maxSegmentBytes = {0};
+        final boolean[] stopAfterOpenSegments = {false};
         // headerId -> total media bytes seen, so the decoded response passes the same integrity check
         // (getIntegrityIssues -> "missing-media") as the buffered path WITHOUT retaining the bytes.
         final Map<Integer, Long> mediaBytesByHeaderId = new HashMap<>();
@@ -160,16 +189,23 @@ public static Result readUntil(@Nonnull final InputStream in,
                         maxSegmentBytes[0] = Math.max(maxSegmentBytes[0], segment.getLength());
                         if (segmentConsumer == null) {
                             segments.add(segment);
-                        } else {
-                            return segmentConsumer.accept(segment);
+                        } else if (!segmentConsumer.accept(segment)) {
+                            stopAfterOpenSegments[0] = true;
                         }
                     }
+                } else if (type == SabrResponseDecoder.LIVE_METADATA) {
+                    final byte[] payload = readPayloadBytes(payloadStream, size);
+                    controlPayloadBytes[0] += payload.length;
+                    controlParts.add(new UmpPart(type, payload.length, payload));
+                    if (liveMetadataConsumer != null) {
+                        liveMetadataConsumer.accept(SabrLiveMetadata.decode(payload));
+                    }
                 } else {
                     final byte[] payload = readPayloadBytes(payloadStream, size);
                     controlPayloadBytes[0] += payload.length;
                     controlParts.add(new UmpPart(type, payload.length, payload));
                 }
-                return true;
+                return !stopAfterOpenSegments[0] || collector.hasOpenSegments();
             });
         } finally {
             // Any header left open after EOF, cancellation or failure must wake a growing-file