From 4cb35a17effae5e9c71f905076eb969f5abac74b Mon Sep 17 00:00:00 2001
From: Byungwan Jun <unipro.kr@gmail.com>
Date: Wed, 20 Nov 2024 23:07:14 +0900
Subject: [PATCH] Support multiple #EXT-X-KEY tags and adjust order

---
 src/parser.rs   |  8 ++++----
 src/playlist.rs | 31 +++++++++++++++----------------
 tests/lib.rs    | 23 ++++++++++++++++-------
 3 files changed, 35 insertions(+), 27 deletions(-)

diff --git a/src/parser.rs b/src/parser.rs
index 10344ac..82d7d36 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -391,7 +391,7 @@ fn media_playlist_tag(i: &[u8]) -> IResult<&[u8], MediaPlaylistTag> {
 fn media_playlist_from_tags(mut tags: Vec<MediaPlaylistTag>) -> MediaPlaylist {
     let mut media_playlist = MediaPlaylist::default();
     let mut next_segment = MediaSegment::empty();
-    let mut encryption_key = None;
+    let mut encryption_keys = vec![];
     let mut map = None;
 
     while let Some(tag) = tags.pop() {
@@ -435,7 +435,7 @@ fn media_playlist_from_tags(mut tags: Vec<MediaPlaylistTag>) -> MediaPlaylist {
                     next_segment.discontinuity = true;
                 }
                 SegmentTag::Key(k) => {
-                    encryption_key = Some(k);
+                    encryption_keys.push(k);
                 }
                 SegmentTag::Map(m) => {
                     map = Some(m);
@@ -450,12 +450,12 @@ fn media_playlist_from_tags(mut tags: Vec<MediaPlaylistTag>) -> MediaPlaylist {
                     next_segment.unknown_tags.push(t);
                 }
                 SegmentTag::Uri(u) => {
-                    next_segment.key = encryption_key.clone();
+                    next_segment.keys = encryption_keys;
                     next_segment.map = map.clone();
                     next_segment.uri = u;
                     media_playlist.segments.push(next_segment);
                     next_segment = MediaSegment::empty();
-                    encryption_key = None;
+                    encryption_keys = vec![];
                     map = None;
                 }
                 _ => (),
diff --git a/src/playlist.rs b/src/playlist.rs
index 8b3fa0c..ed0c925 100644
--- a/src/playlist.rs
+++ b/src/playlist.rs
@@ -838,7 +838,7 @@ pub struct MediaSegment {
     /// `#EXT-X-DISCONTINUITY`
     pub discontinuity: bool,
     /// `#EXT-X-KEY:<attribute-list>`
-    pub key: Option<Key>,
+    pub keys: Vec<Key>,
     /// `#EXT-X-MAP:<attribute-list>`
     pub map: Option<Map>,
     /// `#EXT-X-PROGRAM-DATE-TIME:<YYYY-MM-DDThh:mm:ssZ>`
@@ -855,24 +855,19 @@ impl MediaSegment {
     }
 
     pub(crate) fn write_to<T: Write>(&self, w: &mut T) -> std::io::Result<()> {
-        if let Some(ref byte_range) = self.byte_range {
-            write!(w, "#EXT-X-BYTERANGE:")?;
-            byte_range.write_value_to(w)?;
-            writeln!(w)?;
-        }
         if self.discontinuity {
             writeln!(w, "#EXT-X-DISCONTINUITY")?;
         }
-        if let Some(ref key) = self.key {
-            write!(w, "#EXT-X-KEY:")?;
-            key.write_attributes_to(w)?;
-            writeln!(w)?;
-        }
         if let Some(ref map) = self.map {
             write!(w, "#EXT-X-MAP:")?;
             map.write_attributes_to(w)?;
             writeln!(w)?;
         }
+        for key in &self.keys {
+            write!(w, "#EXT-X-KEY:")?;
+            key.write_attributes_to(w)?;
+            writeln!(w)?;
+        }
         if let Some(ref v) = self.program_date_time {
             writeln!(
                 w,
@@ -885,10 +880,6 @@ impl MediaSegment {
             v.write_attributes_to(w)?;
             writeln!(w)?;
         }
-        for unknown_tag in &self.unknown_tags {
-            writeln!(w, "{}", unknown_tag)?;
-        }
-
         match WRITE_OPT_FLOAT_PRECISION.load(Ordering::Relaxed) {
             MAX => {
                 write!(w, "#EXTINF:{},", self.duration)?;
@@ -897,12 +888,20 @@ impl MediaSegment {
                 write!(w, "#EXTINF:{:.*},", n, self.duration)?;
             }
         };
-
         if let Some(ref v) = self.title {
             writeln!(w, "{}", v)?;
         } else {
             writeln!(w)?;
         }
+        if let Some(ref byte_range) = self.byte_range {
+            write!(w, "#EXT-X-BYTERANGE:")?;
+            byte_range.write_value_to(w)?;
+            writeln!(w)?;
+        }
+        for unknown_tag in &self.unknown_tags {
+            writeln!(w, "{}", unknown_tag)?;
+        }
+
 
         writeln!(w, "{}", self.uri)
     }
diff --git a/tests/lib.rs b/tests/lib.rs
index a388365..9ed53b5 100644
--- a/tests/lib.rs
+++ b/tests/lib.rs
@@ -373,13 +373,22 @@ fn create_and_parse_media_playlist_full() {
                 offset: Some(4559),
             }),
             discontinuity: true,
-            key: Some(Key {
-                method: KeyMethod::None,
-                uri: Some("https://secure.domain.com".into()),
-                iv: Some("0xb059217aa2649ce170b734".into()),
-                keyformat: Some("xXkeyformatXx".into()),
-                keyformatversions: Some("xXFormatVers".into()),
-            }),
+            keys: vec![
+                Key {
+                    method: KeyMethod::None,
+                    uri: Some("https://secure.domain.com".into()),
+                    iv: Some("0xb059217aa2649ce170b734".into()),
+                    keyformat: Some("xXkeyformatXx".into()),
+                    keyformatversions: Some("xXFormatVers".into()),
+                },
+                Key {
+                    method: KeyMethod::AES128,
+                    uri: Some("skd://c2VjdXJlLmRvbWFpbi5jb20=".into()),
+                    iv: Some("0xb059217aa2649ce170b734".into()),
+                    keyformat: Some("com.apple.streamingkeydelivery".into()),
+                    keyformatversions: Some("1".into()),
+                },
+            ],
             map: Some(Map {
                 uri: "www.map-uri.com".into(),
                 byte_range: Some(ByteRange {