hls_m3u8/src/media_playlist.rs

use std::collections::HashSet;
use std::fmt;
use std::str::FromStr;
use std::time::Duration;

use derive_builder::Builder;
use shorthand::ShortHand;

use crate::line::{Line, Lines, Tag};
use crate::media_segment::MediaSegment;
use crate::tags::{
    ExtM3u, ExtXDiscontinuitySequence, ExtXEndList, ExtXIFramesOnly, ExtXIndependentSegments,
    ExtXKey, ExtXMediaSequence, ExtXPlaylistType, ExtXStart, ExtXTargetDuration, ExtXVersion,
};
use crate::types::{EncryptionMethod, ProtocolVersion};
use crate::utils::tag;
use crate::{Error, RequiredVersion};

/// Media playlist.
#[derive(ShortHand, Debug, Clone, Builder, PartialEq, PartialOrd)]
#[builder(build_fn(validate = "Self::validate"))]
#[builder(setter(into, strip_option))]
#[shorthand(enable(must_use, collection_magic, get_mut))]
pub struct MediaPlaylist {
    /// The [`ExtXTargetDuration`] tag of the playlist.
    ///
    /// # Note
    ///
    /// This field is required.
    #[shorthand(enable(copy))]
    target_duration: ExtXTargetDuration,
    /// Sets the [`ExtXMediaSequence`] tag.
    ///
    /// # Note
    ///
    /// This field is optional.
    #[builder(default)]
    media_sequence: Option<ExtXMediaSequence>,
    /// Sets the [`ExtXDiscontinuitySequence`] tag.
    ///
    /// # Note
    ///
    /// This field is optional.
    #[builder(default)]
    discontinuity_sequence: Option<ExtXDiscontinuitySequence>,
    /// Sets the [`ExtXPlaylistType`] tag.
    ///
    /// # Note
    ///
    /// This field is optional.
    #[builder(default)]
    playlist_type: Option<ExtXPlaylistType>,
    /// Sets the [`ExtXIFramesOnly`] tag.
    ///
    /// # Note
    ///
    /// This field is optional.
    #[builder(default)]
    i_frames_only: Option<ExtXIFramesOnly>,
    /// Sets the [`ExtXIndependentSegments`] tag.
    ///
    /// # Note
    ///
    /// This field is optional.
    #[builder(default)]
    independent_segments: Option<ExtXIndependentSegments>,
    /// Sets the [`ExtXStart`] tag.
    ///
    /// # Note
    ///
    /// This field is optional.
    #[builder(default)]
    start: Option<ExtXStart>,
    /// Sets the [`ExtXEndList`] tag.
    ///
    /// # Note
    ///
    /// This field is optional.
    #[builder(default)]
    end_list: Option<ExtXEndList>,
    /// A list of all [`MediaSegment`]s.
    ///
    /// # Note
    ///
    /// This field is required.
    segments: Vec<MediaSegment>,
    /// The allowable excess duration of each media segment in the
    /// associated playlist.
    ///
    /// # Error
    ///
    /// If there is a media segment of which duration exceeds
    /// `#EXT-X-TARGETDURATION + allowable_excess_duration`,
    /// the invocation of `MediaPlaylistBuilder::build()` method will fail.
    ///
    ///
    /// # Note
    ///
    /// This field is optional and the default value is
    /// `Duration::from_secs(0)`.
    #[builder(default = "Duration::from_secs(0)")]
    allowable_excess_duration: Duration,
    /// A list of unknown tags.
    ///
    /// # Note
    ///
    /// This field is optional.
    #[builder(default)]
    unknown_tags: Vec<String>,
}

impl MediaPlaylistBuilder {
    fn validate(&self) -> Result<(), String> {
        if let Some(target_duration) = &self.target_duration {
            self.validate_media_segments(target_duration.duration())
                .map_err(|e| e.to_string())?;
        }

        Ok(())
    }

    fn validate_media_segments(&self, target_duration: Duration) -> crate::Result<()> {
        let mut last_range_uri = None;

        if let Some(segments) = &self.segments {
            for s in segments {
                // CHECK: `#EXT-X-TARGETDURATION`
                let segment_duration = s.inf().duration();
                let rounded_segment_duration = {
                    if segment_duration.subsec_nanos() < 500_000_000 {
                        Duration::from_secs(segment_duration.as_secs())
                    } else {
                        Duration::from_secs(segment_duration.as_secs() + 1)
                    }
                };

                let max_segment_duration = {
                    if let Some(value) = &self.allowable_excess_duration {
                        target_duration + *value
                    } else {
                        target_duration
                    }
                };

                if rounded_segment_duration > max_segment_duration {
                    return Err(Error::custom(format!(
                        "Too large segment duration: actual={:?}, max={:?}, target_duration={:?}, uri={:?}",
                        segment_duration,
                        max_segment_duration,
                        target_duration,
                        s.uri()
                    )));
                }

                // CHECK: `#EXT-X-BYTE-RANGE`
                if let Some(range) = s.byte_range() {
                    if range.start().is_none() {
                        let last_uri = last_range_uri.ok_or_else(Error::invalid_input)?;
                        if last_uri != s.uri() {
                            return Err(Error::invalid_input());
                        }
                    } else {
                        last_range_uri = Some(s.uri());
                    }
                } else {
                    last_range_uri = None;
                }
            }
        }

        Ok(())
    }

    /// Adds a media segment to the resulting playlist.
    pub fn push_segment<VALUE: Into<MediaSegment>>(&mut self, value: VALUE) -> &mut Self {
        if let Some(segments) = &mut self.segments {
            segments.push(value.into());
        } else {
            self.segments = Some(vec![value.into()]);
        }
        self
    }

    /// Parse the rest of the [`MediaPlaylist`] from an m3u8 file.
    pub fn parse(&mut self, input: &str) -> crate::Result<MediaPlaylist> {
        parse_media_playlist(input, self)
    }
}

impl RequiredVersion for MediaPlaylistBuilder {
    fn required_version(&self) -> ProtocolVersion {
        required_version![
            self.target_duration,
            self.media_sequence,
            self.discontinuity_sequence,
            self.playlist_type,
            self.i_frames_only,
            self.independent_segments,
            self.start,
            self.end_list,
            self.segments
        ]
    }
}

impl MediaPlaylist {
    /// Returns a builder for [`MediaPlaylist`].
    #[must_use]
    #[inline]
    pub fn builder() -> MediaPlaylistBuilder { MediaPlaylistBuilder::default() }
}

impl RequiredVersion for MediaPlaylist {
    fn required_version(&self) -> ProtocolVersion {
        required_version![
            self.target_duration,
            self.media_sequence,
            self.discontinuity_sequence,
            self.playlist_type,
            self.i_frames_only,
            self.independent_segments,
            self.start,
            self.end_list,
            self.segments
        ]
    }
}

impl fmt::Display for MediaPlaylist {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "{}", ExtM3u)?;

        if self.required_version() != ProtocolVersion::V1 {
            writeln!(f, "{}", ExtXVersion::new(self.required_version()))?;
        }

        writeln!(f, "{}", self.target_duration)?;

        if let Some(value) = &self.media_sequence {
            writeln!(f, "{}", value)?;
        }

        if let Some(value) = &self.discontinuity_sequence {
            writeln!(f, "{}", value)?;
        }

        if let Some(value) = &self.playlist_type {
            writeln!(f, "{}", value)?;
        }

        if let Some(value) = &self.i_frames_only {
            writeln!(f, "{}", value)?;
        }

        if let Some(value) = &self.independent_segments {
            writeln!(f, "{}", value)?;
        }

        if let Some(value) = &self.start {
            writeln!(f, "{}", value)?;
        }

        // most likely only 1 ExtXKey will be in the HashSet:
        let mut available_keys = HashSet::with_capacity(1);

        for segment in &self.segments {
            for key in segment.keys() {
                // the key is new:
                if available_keys.insert(key) {
                    let mut remove_key = None;

                    // an old key might be removed:
                    for k in &available_keys {
                        if k.key_format() == key.key_format() && &key != k {
                            remove_key = Some(*k);
                            break;
                        }
                    }

                    if let Some(k) = remove_key {
                        // this should always be true:
                        let res = available_keys.remove(k);
                        debug_assert!(res);
                    }

                    writeln!(f, "{}", key)?;
                }
            }

            write!(f, "{}", segment)?;
        }

        if let Some(value) = &self.end_list {
            writeln!(f, "{}", value)?;
        }

        for value in &self.unknown_tags {
            writeln!(f, "{}", value)?;
        }

        Ok(())
    }
}

fn parse_media_playlist(
    input: &str,
    builder: &mut MediaPlaylistBuilder,
) -> crate::Result<MediaPlaylist> {
    let input = tag(input, "#EXTM3U")?;

    let mut segment = MediaSegment::builder();
    let mut segments = vec![];

    let mut has_partial_segment = false;
    let mut has_discontinuity_tag = false;
    let mut unknown_tags = vec![];

    let mut available_keys: Vec<crate::tags::ExtXKey> = vec![];

    for line in Lines::from(input) {
        match line? {
            Line::Tag(tag) => {
                match tag {
                    Tag::ExtInf(t) => {
                        has_partial_segment = true;
                        segment.inf(t);
                    }
                    Tag::ExtXByteRange(t) => {
                        has_partial_segment = true;
                        segment.byte_range(t);
                    }
                    Tag::ExtXDiscontinuity(t) => {
                        has_discontinuity_tag = true;
                        has_partial_segment = true;
                        segment.discontinuity(t);
                    }
                    Tag::ExtXKey(key) => {
                        has_partial_segment = true;

                        // An ExtXKey applies to every MediaSegment and to every Media
                        // Initialization Section declared by an ExtXMap tag, that appears
                        // between it and the next ExtXKey tag in the Playlist file with the
                        // same KEYFORMAT attribute (or the end of the Playlist file).

                        let mut is_new_key = true;

                        for old_key in &mut available_keys {
                            if old_key.key_format() == key.key_format() {
                                *old_key = key.clone();
                                is_new_key = false;
                                // there are no keys with the same key_format in available_keys
                                // so the loop can stop here:
                                break;
                            }
                        }

                        if is_new_key {
                            available_keys.push(key);
                        }
                    }
                    Tag::ExtXMap(mut t) => {
                        has_partial_segment = true;

                        t.set_keys(available_keys.clone());
                        segment.map(t);
                    }
                    Tag::ExtXProgramDateTime(t) => {
                        has_partial_segment = true;
                        segment.program_date_time(t);
                    }
                    Tag::ExtXDateRange(t) => {
                        has_partial_segment = true;
                        segment.date_range(t);
                    }
                    Tag::ExtXTargetDuration(t) => {
                        builder.target_duration(t);
                    }
                    Tag::ExtXMediaSequence(t) => {
                        builder.media_sequence(t);
                    }
                    Tag::ExtXDiscontinuitySequence(t) => {
                        if segments.is_empty() {
                            return Err(Error::invalid_input());
                        }

                        if has_discontinuity_tag {
                            return Err(Error::invalid_input());
                        }

                        builder.discontinuity_sequence(t);
                    }
                    Tag::ExtXEndList(t) => {
                        builder.end_list(t);
                    }
                    Tag::ExtXPlaylistType(t) => {
                        builder.playlist_type(t);
                    }
                    Tag::ExtXIFramesOnly(t) => {
                        builder.i_frames_only(t);
                    }
                    Tag::ExtXMedia(_)
                    | Tag::VariantStream(_)
                    | Tag::ExtXSessionData(_)
                    | Tag::ExtXSessionKey(_) => {
                        return Err(Error::unexpected_tag(tag));
                    }
                    Tag::ExtXIndependentSegments(t) => {
                        builder.independent_segments(t);
                    }
                    Tag::ExtXStart(t) => {
                        builder.start(t);
                    }
                    Tag::ExtXVersion(_) => {}
                    Tag::Unknown(_) => {
                        // [6.3.1. General Client Responsibilities]
                        // > ignore any unrecognized tags.
                        unknown_tags.push(tag.to_string());
                    }
                }
            }
            Line::Uri(uri) => {
                segment.uri(uri);
                segment.keys(available_keys.clone());
                segments.push(segment.build().map_err(Error::builder)?);
                segment = MediaSegment::builder();
                has_partial_segment = false;
            }
            _ => {}
        }
    }

    if has_partial_segment {
        return Err(Error::invalid_input());
    }

    builder.unknown_tags(unknown_tags);
    builder.segments(segments);
    builder.build().map_err(Error::builder)
}

impl FromStr for MediaPlaylist {
    type Err = Error;

    fn from_str(input: &str) -> Result<Self, Self::Err> {
        parse_media_playlist(input, &mut Self::builder())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use pretty_assertions::assert_eq;

    #[test]
    fn too_large_segment_duration_test() {
        let playlist = concat!(
            "#EXTM3U\n",
            "#EXT-X-TARGETDURATION:8\n",
            "#EXT-X-VERSION:3\n",
            "#EXTINF:9.009,\n",
            "http://media.example.com/first.ts\n",
            "#EXTINF:9.509,\n",
            "http://media.example.com/second.ts\n",
            "#EXTINF:3.003,\n",
            "http://media.example.com/third.ts\n",
            "#EXT-X-ENDLIST\n"
        );

        // Error (allowable segment duration = target duration = 8)
        assert!(playlist.parse::<MediaPlaylist>().is_err());

        // Error (allowable segment duration = 9)
        assert!(MediaPlaylist::builder()
            .allowable_excess_duration(Duration::from_secs(1))
            .parse(playlist)
            .is_err());

        // Ok (allowable segment duration = 10)
        MediaPlaylist::builder()
            .allowable_excess_duration(Duration::from_secs(2))
            .parse(playlist)
            .unwrap();
    }

    #[test]
    fn test_empty_playlist() {
        let playlist = "";
        assert!(playlist.parse::<MediaPlaylist>().is_err());
    }
}