1 files changed, 178 insertions, 0 deletions
diff --git a/src/Jellyfin.MediaEncoding.Keyframes/Matroska/Extensions/EbmlReaderExtensions.cs b/src/Jellyfin.MediaEncoding.Keyframes/Matroska/Extensions/EbmlReaderExtensions.cs
new file mode 100644
index 000000000..75d5aafe0
--- /dev/null
+++ b/src/Jellyfin.MediaEncoding.Keyframes/Matroska/Extensions/EbmlReaderExtensions.cs
@@ -0,0 +1,178 @@
+using System;
+using System.Buffers.Binary;
+using Jellyfin.MediaEncoding.Keyframes.Matroska.Models;
+using NEbml.Core;
+
+namespace Jellyfin.MediaEncoding.Keyframes.Matroska.Extensions
+{
+    /// <summary>
+    /// Extension methods for the <see cref="EbmlReader"/> class.
+    /// </summary>
+    internal static class EbmlReaderExtensions
+    {
+        /// <summary>
+        /// Traverses the current container to find the element with <paramref name="identifier"/> identifier.
+        /// </summary>
+        /// <param name="reader">An instance of <see cref="EbmlReader"/>.</param>
+        /// <param name="identifier">The element identifier.</param>
+        /// <returns>A value indicating whether the element was found.</returns>
+        internal static bool FindElement(this EbmlReader reader, ulong identifier)
+        {
+            while (reader.ReadNext())
+            {
+                if (reader.ElementId.EncodedValue == identifier)
+                {
+                    return true;
+                }
+            }
+
+            return false;
+        }
+
+        /// <summary>
+        /// Reads the current position in the file as an unsigned integer converted from binary.
+        /// </summary>
+        /// <param name="reader">An instance of <see cref="EbmlReader"/>.</param>
+        /// <returns>The unsigned integer.</returns>
+        internal static uint ReadUIntFromBinary(this EbmlReader reader)
+        {
+            var buffer = new byte[4];
+            reader.ReadBinary(buffer, 0, 4);
+            return BinaryPrimitives.ReadUInt32BigEndian(buffer);
+        }
+
+        /// <summary>
+        /// Reads from the start of the file to retrieve the SeekHead segment.
+        /// </summary>
+        /// <param name="reader">An instance of <see cref="EbmlReader"/>.</param>
+        /// <returns>Instance of <see cref="SeekHead"/>.</returns>
+        internal static SeekHead ReadSeekHead(this EbmlReader reader)
+        {
+            reader = reader ?? throw new ArgumentNullException(nameof(reader));
+
+            if (reader.ElementPosition != 0)
+            {
+                throw new InvalidOperationException("File position must be at 0");
+            }
+
+            // Skip the header
+            if (!reader.FindElement(MatroskaConstants.SegmentContainer))
+            {
+                throw new InvalidOperationException("Expected a segment container");
+            }
+
+            reader.EnterContainer();
+
+            long? tracksPosition = null;
+            long? cuesPosition = null;
+            long? infoPosition = null;
+            // The first element should be a SeekHead otherwise we'll have to search manually
+            if (!reader.FindElement(MatroskaConstants.SeekHead))
+            {
+                throw new InvalidOperationException("Expected a SeekHead");
+            }
+
+            reader.EnterContainer();
+            while (reader.FindElement(MatroskaConstants.Seek))
+            {
+                reader.EnterContainer();
+                reader.ReadNext();
+                var type = (ulong)reader.ReadUIntFromBinary();
+                switch (type)
+                {
+                    case MatroskaConstants.Tracks:
+                        reader.ReadNext();
+                        tracksPosition = (long)reader.ReadUInt();
+                        break;
+                    case MatroskaConstants.Cues:
+                        reader.ReadNext();
+                        cuesPosition = (long)reader.ReadUInt();
+                        break;
+                    case MatroskaConstants.Info:
+                        reader.ReadNext();
+                        infoPosition = (long)reader.ReadUInt();
+                        break;
+                }
+
+                reader.LeaveContainer();
+
+                if (tracksPosition.HasValue && cuesPosition.HasValue && infoPosition.HasValue)
+                {
+                    break;
+                }
+            }
+
+            reader.LeaveContainer();
+
+            if (!tracksPosition.HasValue || !cuesPosition.HasValue || !infoPosition.HasValue)
+            {
+                throw new InvalidOperationException("SeekHead is missing or does not contain Info, Tracks and Cues positions");
+            }
+
+            return new SeekHead(infoPosition.Value, tracksPosition.Value, cuesPosition.Value);
+        }
+
+        /// <summary>
+        /// Reads from SegmentContainer to retrieve the Info segment.
+        /// </summary>
+        /// <param name="reader">An instance of <see cref="EbmlReader"/>.</param>
+        /// <param name="position">The position of the info segment relative to the Segment container.</param>
+        /// <returns>Instance of <see cref="Info"/>.</returns>
+        internal static Info ReadInfo(this EbmlReader reader, long position)
+        {
+            reader.ReadAt(position);
+
+            double? duration = null;
+            reader.EnterContainer();
+            // Mandatory element
+            reader.FindElement(MatroskaConstants.TimestampScale);
+            var timestampScale = reader.ReadUInt();
+
+            if (reader.FindElement(MatroskaConstants.Duration))
+            {
+                duration = reader.ReadFloat();
+            }
+
+            reader.LeaveContainer();
+
+            return new Info((long)timestampScale, duration);
+        }
+
+        /// <summary>
+        /// Enters the Tracks segment and reads all tracks to find the specified type.
+        /// </summary>
+        /// <param name="reader">Instance of <see cref="EbmlReader"/>.</param>
+        /// <param name="tracksPosition">The relative position of the tracks segment.</param>
+        /// <param name="type">The track type identifier.</param>
+        /// <returns>The first track number with the specified type.</returns>
+        /// <exception cref="InvalidOperationException">Stream type is not found.</exception>
+        internal static ulong FindFirstTrackNumberByType(this EbmlReader reader, long tracksPosition, ulong type)
+        {
+            reader.ReadAt(tracksPosition);
+
+            reader.EnterContainer();
+            while (reader.FindElement(MatroskaConstants.TrackEntry))
+            {
+                reader.EnterContainer();
+                // Mandatory element
+                reader.FindElement(MatroskaConstants.TrackNumber);
+                var trackNumber = reader.ReadUInt();
+
+                // Mandatory element
+                reader.FindElement(MatroskaConstants.TrackType);
+                var trackType = reader.ReadUInt();
+
+                reader.LeaveContainer();
+                if (trackType == MatroskaConstants.TrackTypeVideo)
+                {
+                    reader.LeaveContainer();
+                    return trackNumber;
+                }
+            }
+
+            reader.LeaveContainer();
+
+            throw new InvalidOperationException($"No stream with type {type} found");
+        }
+    }
+}