From 6ebab44aeee149e232bb531c9bfeb382e4bfa335 Mon Sep 17 00:00:00 2001
From: Lysander <mgician3@gmail.com>
Date: Sat, 14 Dec 2024 18:17:35 +0800
Subject: [PATCH] add EXT-X-KEY tag support to playlist parser

---
 pkg/playlist/media.go         |  19 ++++++
 pkg/playlist/media_key.go     | 117 ++++++++++++++++++++++++++++++++++
 pkg/playlist/media_segment.go |   3 +
 3 files changed, 139 insertions(+)
 create mode 100644 pkg/playlist/media_key.go

diff --git a/pkg/playlist/media.go b/pkg/playlist/media.go
index e7bd3ea..b837acc 100644
--- a/pkg/playlist/media.go
+++ b/pkg/playlist/media.go
@@ -107,6 +107,8 @@ func (m *Media) Unmarshal(buf []byte) error {
 		return err
 	}
 
+	var curKey *MediaKey
+
 	curSegment := &MediaSegment{}
 
 	for {
@@ -224,6 +226,15 @@ func (m *Media) Unmarshal(buf []byte) error {
 				return err
 			}
 
+		case strings.HasPrefix(line, "#EXT-X-KEY:"):
+			line = line[len("#EXT-X-KEY:"):]
+
+			curKey = &MediaKey{}
+			err = curKey.unmarshal(line)
+			if err != nil {
+				return err
+			}
+
 		case strings.HasPrefix(line, "#EXT-X-SKIP:"):
 			line = line[len("#EXT-X-SKIP:"):]
 
@@ -278,6 +289,8 @@ func (m *Media) Unmarshal(buf []byte) error {
 			curSegment.Duration = du
 			curSegment.Title = strings.TrimSpace(parts[1])
 
+			curSegment.Key = curKey
+
 		case strings.HasPrefix(line, "#EXT-X-BYTERANGE:"):
 			line = line[len("#EXT-X-BYTERANGE:"):]
 
@@ -387,7 +400,13 @@ func (m Media) Marshal() ([]byte, error) {
 		ret += m.Skip.marshal()
 	}
 
+	var prevKey *MediaKey
 	for _, seg := range m.Segments {
+		if seg.Key != nil && (prevKey == nil || !seg.Key.Equal(prevKey)) {
+			ret += seg.Key.marshal()
+			prevKey = seg.Key
+		}
+
 		ret += seg.marshal()
 	}
 
diff --git a/pkg/playlist/media_key.go b/pkg/playlist/media_key.go
new file mode 100644
index 0000000..3f23509
--- /dev/null
+++ b/pkg/playlist/media_key.go
@@ -0,0 +1,117 @@
+package playlist
+
+import (
+	"fmt"
+
+	"github.com/bluenviron/gohlslib/v2/pkg/playlist/primitives"
+)
+
+type MediaKeyMethod string
+
+const (
+	MediaKeyMethodNone      = "NONE"
+	MediaKeyMethodAES128    = "AES-128"
+	MediaKeyMethodSampleAes = "SAMPLE-AES"
+)
+
+// MediaKey is a EXT-X-KEY tag.
+type MediaKey struct {
+	// METHOD
+	// required
+	Method MediaKeyMethod
+
+	// URI is required unless METHOD is NONE
+	URI string
+
+	// IV
+	IV string
+
+	// KEYFORMAT
+	KeyFormat string
+
+	// KEYFORMATVERSIONS
+	KeyFormatVersions string
+}
+
+func (t *MediaKey) unmarshal(v string) error {
+	attrs, err := primitives.AttributesUnmarshal(v)
+	if err != nil {
+		return err
+	}
+
+	for key, val := range attrs {
+		switch key {
+		case "METHOD":
+			km := MediaKeyMethod(val)
+			if km != MediaKeyMethodNone &&
+				km != MediaKeyMethodAES128 &&
+				km != MediaKeyMethodSampleAes {
+				return fmt.Errorf("invalid method: %s", val)
+			}
+			t.Method = km
+
+		case "URI":
+			t.URI = val
+
+		case "IV":
+			t.IV = val
+
+		case "KEYFORMAT":
+			t.KeyFormat = val
+
+		case "KEYFORMATVERSIONS":
+			t.KeyFormatVersions = val
+		}
+	}
+
+	switch t.Method {
+	case MediaKeyMethodAES128, MediaKeyMethodSampleAes:
+		if t.URI == "" {
+			return fmt.Errorf("URI is required for method %s", t.Method)
+		}
+	default:
+	}
+
+	return nil
+}
+
+func (t MediaKey) marshal() string {
+	ret := "#EXT-X-KEY:METHOD=" + string(t.Method)
+
+	// If the encryption method is NONE, other attributes MUST NOT be present.
+	if t.Method != MediaKeyMethodNone {
+		ret += ",URI=\"" + t.URI + "\""
+
+		if t.IV != "" {
+			ret += ",IV=" + t.IV
+		}
+
+		if t.KeyFormat != "" {
+			ret += ",KEYFORMAT=\"" + t.KeyFormat + "\""
+		}
+
+		if t.KeyFormatVersions != "" {
+			ret += ",KEYFORMATVERSIONS=\"" + t.KeyFormatVersions + "\""
+		}
+	}
+
+	ret += "\n"
+
+	return ret
+}
+
+func (t *MediaKey) Equal(key *MediaKey) bool {
+	if t == key {
+		return true
+	}
+
+	if key == nil {
+		return false
+	}
+
+	return t.Method == key.Method &&
+		t.URI == key.URI &&
+		t.IV == key.IV &&
+		t.KeyFormat == key.KeyFormat &&
+		t.KeyFormatVersions == key.KeyFormatVersions
+}
diff --git a/pkg/playlist/media_segment.go b/pkg/playlist/media_segment.go
index 4a5c30e..0de36b3 100644
--- a/pkg/playlist/media_segment.go
+++ b/pkg/playlist/media_segment.go
@@ -31,6 +31,9 @@ type MediaSegment struct {
 	// EXT-X-BITRATE
 	Bitrate *int
 
+	// EXT-X-KEY
+	Key *MediaKey
+
 	// EXT-X-BYTERANGE
 	ByteRangeLength *uint64
 	ByteRangeStart  *uint64