Fix #36 Media Playlist Set* methods modifying wrong segment

reader.go

@@ -506,7 +506,7 @@ func decodeLineOfMediaPlaylist(p *MediaPlaylist, wv *WV, state *decodingState, l
 		}
 		// If EXT-X-KEY appeared before reference to segment (EXTINF) then it linked to this segment
 		if state.tagKey {
-			p.Segments[(p.tail-1)%p.capacity].Key = &Key{state.xkey.Method, state.xkey.URI, state.xkey.IV, state.xkey.Keyformat, state.xkey.Keyformatversions}
+			p.Segments[p.last()].Key = &Key{state.xkey.Method, state.xkey.URI, state.xkey.IV, state.xkey.Keyformat, state.xkey.Keyformatversions}
 			// First EXT-X-KEY may appeared in the header of the playlist and linked to first segment
 			// but for convenient playlist generation it also linked as default playlist key
 			if p.Key == nil {
@@ -516,7 +516,7 @@ func decodeLineOfMediaPlaylist(p *MediaPlaylist, wv *WV, state *decodingState, l
 		}
 		// If EXT-X-MAP appeared before reference to segment (EXTINF) then it linked to this segment
 		if state.tagMap {
-			p.Segments[(p.tail-1)%p.capacity].Map = &Map{state.xmap.URI, state.xmap.Limit, state.xmap.Offset}
+			p.Segments[p.last()].Map = &Map{state.xmap.URI, state.xmap.Limit, state.xmap.Offset}
 			// First EXT-X-MAP may appeared in the header of the playlist and linked to first segment
 			// but for convenient playlist generation it also linked as default playlist map
 			if p.Map == nil {

writer.go

@@ -250,6 +250,14 @@ func NewMediaPlaylist(winsize uint, capacity uint) (*MediaPlaylist, error) {
 	return p, nil
 }
 
+// last returns the previously written segment's index
+func (p *MediaPlaylist) last() uint {
+	if p.tail == 0 {
+		return p.capacity - 1
+	}
+	return p.tail - 1
+}
+
 // Remove current segment from the head of chunk slice form a media playlist. Useful for sliding playlists.
 // This operation does reset playlist cache.
 func (p *MediaPlaylist) Remove() (err error) {
@@ -584,7 +592,7 @@ func (p *MediaPlaylist) SetKey(method, uri, iv, keyformat, keyformatversions str
 		version(&p.ver, 5)
 	}
 
-	p.Segments[(p.tail-1)%p.capacity].Key = &Key{method, uri, iv, keyformat, keyformatversions}
+	p.Segments[p.last()].Key = &Key{method, uri, iv, keyformat, keyformatversions}
 	return nil
 }
 
@@ -594,7 +602,7 @@ func (p *MediaPlaylist) SetMap(uri string, limit, offset int64) error {
 		return errors.New("playlist is empty")
 	}
 	version(&p.ver, 5) // due section 4
-	p.Segments[(p.tail-1)%p.capacity].Map = &Map{uri, limit, offset}
+	p.Segments[p.last()].Map = &Map{uri, limit, offset}
 	return nil
 }
 
@@ -604,16 +612,16 @@ func (p *MediaPlaylist) SetRange(limit, offset int64) error {
 		return errors.New("playlist is empty")
 	}
 	version(&p.ver, 4) // due section 3.4.1
-	p.Segments[(p.tail-1)%p.capacity].Limit = limit
-	p.Segments[(p.tail-1)%p.capacity].Offset = offset
+	p.Segments[p.last()].Limit = limit
+	p.Segments[p.last()].Offset = offset
 	return nil
 }
 
 func (p *MediaPlaylist) SetSCTE(cue string, id string, time float64) error {
 	if p.count == 0 {
 		return errors.New("playlist is empty")
 	}
-	p.Segments[(p.tail-1)%p.capacity].SCTE = &SCTE{cue, id, time}
+	p.Segments[p.last()].SCTE = &SCTE{cue, id, time}
 	return nil
 }
 
@@ -625,7 +633,7 @@ func (p *MediaPlaylist) SetDiscontinuity() error {
 	if p.count == 0 {
 		return errors.New("playlist is empty")
 	}
-	p.Segments[(p.tail-1)%p.capacity].Discontinuity = true
+	p.Segments[p.last()].Discontinuity = true
 	return nil
 }
 
@@ -638,6 +646,6 @@ func (p *MediaPlaylist) SetProgramDateTime(value time.Time) error {
 	if p.count == 0 {
 		return errors.New("playlist is empty")
 	}
-	p.Segments[(p.tail-1)%p.capacity].ProgramDateTime = value
+	p.Segments[p.last()].ProgramDateTime = value
 	return nil
 }

writer_test.go

@@ -47,6 +47,20 @@ func TestCreateMediaPlaylistWithWrongSize(t *testing.T) {
 	}
 }
 
+// Tests the last method on media playlist
+func TestLastSegmentMediaPlaylist(t *testing.T) {
+	p, _ := NewMediaPlaylist(5, 5)
+	if p.last() != 4 {
+		t.Errorf("last is %v, expected: 4", p.last())
+	}
+	for i := uint(0); i < 5; i++ {
+		_ = p.Append("uri.ts", 4, "")
+		if p.last() != i {
+			t.Errorf("last is: %v, expected: %v", p.last(), i)
+		}
+	}
+}
+
 // Create new media playlist
 // Add two segments to media playlist
 func TestAddSegmentToMediaPlaylist(t *testing.T) {
@@ -268,48 +282,31 @@ func TestLoopSegmentsOfMediaPlaylist(t *testing.T) {
 	//fmt.Println(p.Encode().String())
 }
 
-// Create new media playlist with capacity 30
-// Add 10 segments to media playlist
-// Add encryption key
-// Add another 10 segments to media playlist
-// Add new encryption key
-// Add another 10 segments to media playlist
-// Iterate over segments
+// Create new media playlist with capacity 5
+// Add 5 segments and 5 unique keys
+// Test correct keys set on correct segments
 func TestEncryptionKeysInMediaPlaylist(t *testing.T) {
-	// Create new media playlist with capacity 30
-	p, e := NewMediaPlaylist(5, 15)
-	if e != nil {
-		t.Fatalf("Create media playlist failed: %s", e)
-	}
-	// Add 10 segments to media playlist
-	for i := 0; i < 5; i++ {
-		e = p.Append(fmt.Sprintf("test0-%d.ts", i), 6.0, "")
-		if e != nil {
-			t.Errorf("Add segment #%d to a media playlist failed: %s", i, e)
+	p, _ := NewMediaPlaylist(5, 5)
+	// Add 5 segments and set custom encryption key
+	for i := uint(0); i < 5; i++ {
+		uri := fmt.Sprintf("uri-%d", i)
+		expected := &Key{
+			Method:            "AES-128",
+			URI:               uri,
+			IV:                fmt.Sprintf("%d", i),
+			Keyformat:         "identity",
+			Keyformatversions: "1",
 		}
-	}
-	// Add encryption key
-	p.SetKey("AES-128", "https://example.com/", "0X00000000000000000000000000000000", "key-format1", "version x.x")
-	// Add 10 segments to media playlist
-	for i := 0; i < 5; i++ {
-		e = p.Append(fmt.Sprintf("test1-%d.ts", i), 6.0, "")
-		if e != nil {
-			t.Errorf("Add segment #%d to a media playlist failed: %s", i, e)
+		_ = p.Append(uri+".ts", 4, "")
+		_ = p.SetKey(expected.Method, expected.URI, expected.IV, expected.Keyformat, expected.Keyformatversions)
+
+		if p.Segments[i].Key == nil {
+			t.Fatalf("Key was not set on segment %v", i)
 		}
-	}
-	// Add encryption key
-	p.SetKey("AES-128", "https://example.com/", "0X00000000000000000000000000000001", "key-format2", "version x.x")
-	// Add 10 segments to media playlist
-	for i := 0; i < 5; i++ {
-		e = p.Append(fmt.Sprintf("test2-%d.ts", i), 6.0, "")
-		if e != nil {
-			t.Errorf("Add segment #%d to a media playlist failed: %s", i, e)
+		if *p.Segments[i].Key != *expected {
+			t.Errorf("Key %+v does not match expected %+v", p.Segments[i].Key, expected)
 		}
 	}
-	for i := 0; i < 3; i++ {
-		//fmt.Printf("Iteration %d:\n%s\n", i, p.Encode().String())
-		p.Remove()
-	}
 }
 
 // Create new media playlist