add EXT-X-KEY tag support to playlist parser

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()
 	}
 

pkg/playlist/media_key.go

@@ -0,0 +1,120 @@
+package playlist
+
+import (
+	"fmt"
+
+	"github.com/bluenviron/gohlslib/v2/pkg/playlist/primitives"
+)
+
+// MediaKeyMethod is the encryption method used for the media segments.
+type MediaKeyMethod string
+
+// standard encryption methods
+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
+}
+
+// Equal checks if two MediaKey objects are equal.
+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
+}

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

pkg/playlist/media_test.go

@@ -225,6 +225,150 @@ main.mp4
 			Endlist: true,
 		},
 	},
+	{
+		"key-basic",
+		`#EXTM3U
+#EXT-X-VERSION:3
+#EXT-X-TARGETDURATION:6
+#EXT-X-MEDIA-SEQUENCE:0
+#EXT-X-KEY:METHOD=AES-128,URI="key.bin"
+#EXTINF:6.00000,
+segment1.ts
+#EXTINF:6.00000,
+segment2.ts`,
+		`#EXTM3U
+#EXT-X-VERSION:3
+#EXT-X-TARGETDURATION:6
+#EXT-X-MEDIA-SEQUENCE:0
+#EXT-X-KEY:METHOD=AES-128,URI="key.bin"
+#EXTINF:6.00000,
+segment1.ts
+#EXTINF:6.00000,
+segment2.ts
+`,
+		Media{
+			Version:        3,
+			TargetDuration: 6,
+			Segments: []*MediaSegment{
+				{
+					Duration: 6 * time.Second,
+					URI:      "segment1.ts",
+					Key: &MediaKey{
+						Method: MediaKeyMethodAES128,
+						URI:    "key.bin",
+					},
+				},
+				{
+					Duration: 6 * time.Second,
+					URI:      "segment2.ts",
+					Key: &MediaKey{
+						Method: MediaKeyMethodAES128,
+						URI:    "key.bin",
+					},
+				},
+			},
+		},
+	},
+	{
+		"key-with-iv",
+		`#EXTM3U
+#EXT-X-VERSION:3
+#EXT-X-TARGETDURATION:6
+#EXT-X-MEDIA-SEQUENCE:0
+#EXT-X-KEY:METHOD=AES-128,URI="key.bin",IV=0x1234567890abcdef1234567890abcdef
+#EXTINF:6.00000,
+segment1.ts
+`,
+		`#EXTM3U
+#EXT-X-VERSION:3
+#EXT-X-TARGETDURATION:6
+#EXT-X-MEDIA-SEQUENCE:0
+#EXT-X-KEY:METHOD=AES-128,URI="key.bin",IV=0x1234567890abcdef1234567890abcdef
+#EXTINF:6.00000,
+segment1.ts
+`,
+		Media{
+			Version:        3,
+			TargetDuration: 6,
+			Segments: []*MediaSegment{
+				{
+					Duration: 6 * time.Second,
+					URI:      "segment1.ts",
+					Key: &MediaKey{
+						Method: MediaKeyMethodAES128,
+						URI:    "key.bin",
+						IV:     "0x1234567890abcdef1234567890abcdef",
+					},
+				},
+			},
+		},
+	},
+	{
+		"key-with-format",
+		`#EXTM3U
+#EXT-X-VERSION:5
+#EXT-X-TARGETDURATION:6
+#EXT-X-MEDIA-SEQUENCE:0
+#EXT-X-KEY:METHOD=SAMPLE-AES,URI="key.bin",KEYFORMAT="com.apple.streamingkeydelivery",KEYFORMATVERSIONS="1"
+#EXTINF:6.00000,
+segment1.ts
+`,
+		`#EXTM3U
+#EXT-X-VERSION:5
+#EXT-X-TARGETDURATION:6
+#EXT-X-MEDIA-SEQUENCE:0
+#EXT-X-KEY:METHOD=SAMPLE-AES,URI="key.bin",KEYFORMAT="com.apple.streamingkeydelivery",KEYFORMATVERSIONS="1"
+#EXTINF:6.00000,
+segment1.ts
+`,
+		Media{
+			Version:        5,
+			TargetDuration: 6,
+			Segments: []*MediaSegment{
+				{
+					Duration: 6 * time.Second,
+					URI:      "segment1.ts",
+					Key: &MediaKey{
+						Method:            MediaKeyMethodSampleAES,
+						URI:               "key.bin",
+						KeyFormat:         "com.apple.streamingkeydelivery",
+						KeyFormatVersions: "1",
+					},
+				},
+			},
+		},
+	},
+	{
+		"key-none",
+		`#EXTM3U
+#EXT-X-VERSION:3
+#EXT-X-TARGETDURATION:6
+#EXT-X-MEDIA-SEQUENCE:0
+#EXT-X-KEY:METHOD=NONE
+#EXTINF:6.00000,
+segment1.ts`,
+		`#EXTM3U
+#EXT-X-VERSION:3
+#EXT-X-TARGETDURATION:6
+#EXT-X-MEDIA-SEQUENCE:0
+#EXT-X-KEY:METHOD=NONE
+#EXTINF:6.00000,
+segment1.ts
+`,
+		Media{
+			Version:        3,
+			TargetDuration: 6,
+			Segments: []*MediaSegment{
+				{
+					Duration: 6 * time.Second,
+					URI:      "segment1.ts",
+					Key: &MediaKey{
+						Method: MediaKeyMethodNone,
+					},
+				},
+			},
+		},
+	},
 }
 
 func TestMediaUnmarshal(t *testing.T) {

pkg/playlist/testdata/fuzz/FuzzMediaUnmarshal/32c9b70f1dcd40a8

@@ -0,0 +1,2 @@
+go test fuzz v1
+string("#EXTM3U\n#EXT-X-KEY:0")

pkg/playlist/testdata/fuzz/FuzzMediaUnmarshal/bdfd84d77054b242

@@ -0,0 +1,2 @@
+go test fuzz v1
+string("#EXTM3U\n#EXT-X-KEY:METHOD=AES-128")

pkg/playlist/testdata/fuzz/FuzzMediaUnmarshal/eba654116c9d19f6

@@ -0,0 +1,2 @@
+go test fuzz v1
+string("#EXTM3U\n#EXT-X-KEY:METHOD=")