From c8a0b7df071218304d0753910a3a81c9d3fd2e09 Mon Sep 17 00:00:00 2001
From: scruzin <anoble@gmail.com>
Date: Thu, 11 Jul 2019 17:33:16 +0930
Subject: [PATCH 1/3] Fix IndexPid.

---
 container/mts/mpegts.go | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/container/mts/mpegts.go b/container/mts/mpegts.go
index c43df44c..2a5f9483 100644
--- a/container/mts/mpegts.go
+++ b/container/mts/mpegts.go
@@ -221,7 +221,7 @@ func LastPid(d []byte, pid uint16) (pkt []byte, i int, err error) {
 // along with optional metadata if present. Commonly used to find a
 // PAT immediately followed by a PMT.
 func IndexPid(d []byte, pids ...uint16) (idx int, m map[string]string, err error) {
-	prev := 0
+	idx = -1
 	for _, pid := range pids {
 		pkt, i, _err := FindPid(d, pid)
 		if err != nil {
@@ -231,16 +231,13 @@ func IndexPid(d []byte, pids ...uint16) (idx int, m map[string]string, err error
 		if pid == PmtPid {
 			m, _ = metaFromPMT(pkt)
 		}
-		if prev == 0 {
+		if idx == -1 {
 			idx = i
-			prev = i
-			continue
-		}
-		if i != prev+PacketSize {
+		} else if i != 0 {
 			err = errors.New("PIDs not consecutive")
 			return
 		}
-		prev = i
+		d = d[i+PacketSize:]
 	}
 	return
 }

From c717595adcd95325d471028a50e90c746bb01bf0 Mon Sep 17 00:00:00 2001
From: scruzin <anoble@gmail.com>
Date: Thu, 11 Jul 2019 18:16:33 +0930
Subject: [PATCH 2/3] Added length check to IndexPid.

---
 container/mts/mpegts.go | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/container/mts/mpegts.go b/container/mts/mpegts.go
index 2a5f9483..28ad802b 100644
--- a/container/mts/mpegts.go
+++ b/container/mts/mpegts.go
@@ -178,8 +178,9 @@ func FindPat(d []byte) ([]byte, int, error) {
 
 // Errors used by FindPid.
 var (
-	errInvalidLen   = errors.New("MPEG-TS data not of valid length")
-	errCouldNotFind = errors.New("could not find packet with given PID")
+	errInvalidLen     = errors.New("MPEG-TS data not of valid length")
+	errCouldNotFind   = errors.New("could not find packet with given PID")
+	errNotConsecutive = errors.New("could not find consecutive PIDs")
 )
 
 // FindPid will take a clip of MPEG-TS and try to find a packet with given PID - if one
@@ -223,10 +224,12 @@ func LastPid(d []byte, pid uint16) (pkt []byte, i int, err error) {
 func IndexPid(d []byte, pids ...uint16) (idx int, m map[string]string, err error) {
 	idx = -1
 	for _, pid := range pids {
-		pkt, i, _err := FindPid(d, pid)
+		if len(d) < PacketSize {
+			return idx, m, errInvalidLen
+		}
+		pkt, i, err := FindPid(d, pid)
 		if err != nil {
-			err = errors.Wrap(_err, "could not find PID")
-			return
+			return idx, m, errCouldNotFind
 		}
 		if pid == PmtPid {
 			m, _ = metaFromPMT(pkt)
@@ -234,8 +237,7 @@ func IndexPid(d []byte, pids ...uint16) (idx int, m map[string]string, err error
 		if idx == -1 {
 			idx = i
 		} else if i != 0 {
-			err = errors.New("PIDs not consecutive")
-			return
+			return idx, m, errNotConsecutive
 		}
 		d = d[i+PacketSize:]
 	}

From 01351a308bc0f83561e92c3956bad94b25e278d9 Mon Sep 17 00:00:00 2001
From: scruzin <anoble@gmail.com>
Date: Thu, 11 Jul 2019 19:29:46 +0930
Subject: [PATCH 3/3] Added tests for IndexPid.

---
 container/mts/mpegts_test.go | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/container/mts/mpegts_test.go b/container/mts/mpegts_test.go
index 3a603a20..9ce93b5b 100644
--- a/container/mts/mpegts_test.go
+++ b/container/mts/mpegts_test.go
@@ -491,6 +491,26 @@ func TestSegmentForMeta(t *testing.T) {
 		if !reflect.DeepEqual(want, got) {
 			t.Errorf("did not get expected result for test %v\nGot: %v\nWant: %v\n", testn, got, want)
 		}
+
+		// Now test IndexPid.
+		i, m, err := IndexPid(clip.Bytes(), PatPid, PmtPid)
+		if err != nil {
+			t.Fatalf("IndexPid failed with error: %v", err)
+		}
+		if i != 0 {
+			t.Fatalf("IndexPid unexpected index; got %d, expected 0", i)
+		}
+		if m["n"] != "1" {
+			t.Fatalf("IndexPid unexpected metadata; got %s, expected 1", m["n"])
+		}
+	}
+
+	// Finally, test IndexPid error handling.
+	for _, d := range [][]byte{[]byte{}, make([]byte, PacketSize/2), make([]byte, PacketSize)} {
+		_, _, err := IndexPid(d, PatPid, PmtPid)
+		if err == nil {
+			t.Fatalf("IndexPid expected error")
+		}
 	}
 }