Factored metaFromPMT out of ExtractMeta and added LastPid and IndexPid.

container/mts/mpegts.go

@@ -198,6 +198,52 @@ func FindPid(d []byte, pid uint16) (pkt []byte, i int, err error) {
 	return nil, -1, errCouldNotFind
 }
 
+// LastPid will take a clip of MPEG-TS and try to find a packet
+// with given PID searching in reverse from the end of the clip. If
+// one is found, then it is returned along with its index, otherwise
+// nil, -1 and an error is returned.
+func LastPid(d []byte, pid uint16) (pkt []byte, i int, err error) {
+	if len(d) < PacketSize {
+		return nil, -1, errInvalidLen
+	}
+
+	for i = len(d) - PacketSize; i >= 0; i -= PacketSize {
+		p := (uint16(d[i+1]&0x1f) << 8) | uint16(d[i+2])
+		if p == pid {
+			pkt = d[i : i+PacketSize]
+			return
+		}
+	}
+	return nil, -1, errCouldNotFind
+}
+
+// IndexPid returns the position of one or more consecutive pids,
+// 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
+	for _, pid := range pids {
+		pkt, i, _err := FindPid(d, pid)
+		if err != nil {
+			err = errors.Wrap(_err, "could not find PID")
+			return
+		}
+		if pid == PmtPid {
+			m, _ = metaFromPMT(pkt)
+		}
+		if prev == 0 {
+			idx = i
+			prev = i
+			continue
+		}
+		if i != prev+PacketSize {
+			err = errors.Wrap(err, "PIDs not consecutive")
+			return
+		}
+	}
+	return
+}
+
 // FillPayload takes a channel and fills the packets Payload field until the
 // channel is empty or we've the packet reaches capacity
 func (p *Packet) FillPayload(data []byte) int {
@@ -331,17 +377,18 @@ func DiscontinuityIndicator(f bool) Option {
 var errNoPTS = errors.New("could not find PTS")
 
 // GetPTSRange retreives the first and last PTS of an MPEGTS clip.
+// If there is only one PTS, it is included twice in the pts return value.
 func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) {
 	var _pkt packet.Packet
 	// Find the first packet with PID pidType and PUSI.
 	var i int
 	for {
 		if i >= len(clip) {
-			return [2]uint64{}, errNoPTS
+			return pts, errNoPTS
 		}
 		pkt, _i, err := FindPid(clip[i:], pid)
 		if err != nil {
-			return [2]uint64{}, errors.Wrap(err, fmt.Sprintf("could not find packet of PID: %d", pid))
+			return pts, errors.Wrap(err, fmt.Sprintf("could not find packet of PID: %d", pid))
 		}
 		copy(_pkt[:], pkt)
 		if _pkt.PayloadUnitStartIndicator() {
@@ -353,32 +400,47 @@ func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) {
 	// Get the payload of the packet, which will be the start of the PES packet.
 	payload, err := packet.Payload(&_pkt)
 	if err != nil {
-		return [2]uint64{}, err
+		return pts, err
 	}
 
 	// Get the the first PTS from the PES header.
 	_pes, err := pes.NewPESHeader(payload)
 	if err != nil {
-		return [2]uint64{}, err
+		return pts, err
 	}
 	pts[0] = _pes.PTS()
+	pts[1] = pts[0] // Until we have find a second PTS.
 
-	// Get the final PTS searching from end of clip for access unit start.
-	for i := len(clip) - PacketSize; i >= 0; i -= PacketSize {
-		copy(_pkt[:], clip[i:i+PacketSize])
-		if packet.PayloadUnitStartIndicator(&_pkt) && uint16(_pkt.PID()) == pid {
+	// Get the last PTS searching in reverse from end of the clip.
+	first := i
+	i = len(clip)
+	for {
+		pkt, _i, err := LastPid(clip[:i], pid)
+		if err != nil || i <= first {
+			return pts, nil
+		}
+
+		copy(_pkt[:], pkt)
+		if packet.PayloadUnitStartIndicator(&_pkt) {
+			break
+		}
+
+		i = _i
+	}
+
+	// Get the payload of the packet.
 	payload, err = packet.Payload(&_pkt)
 	if err != nil {
-				return [2]uint64{}, err
-			}
-			_pes, err = pes.NewPESHeader(payload)
-			if err != nil {
-				return [2]uint64{}, err
-			}
-			pts[1] = _pes.PTS()
 		return
 	}
+
+	// Get the the last PTS from the PES header.
+	_pes, err = pes.NewPESHeader(payload)
+	if err != nil {
+		return
 	}
+	pts[1] = _pes.PTS()
+
 	return
 }
 
@@ -388,22 +450,25 @@ var errNoMeta = errors.New("PMT does not contain meta")
 // descriptor, that is found in the MPEG-TS clip d. d must contain a series of
 // complete MPEG-TS packets.
 func ExtractMeta(d []byte) (map[string]string, error) {
-	pkt, _, err := FindPid(d, PmtPid)
+	pmt, _, err := FindPid(d, PmtPid)
 	if err != nil {
 		return nil, err
 	}
-
-	// Get as PSI type. Need to skip MTS header.
-	pmt := psi.PSIBytes(pkt[4:])
-	_, metaDescriptor := pmt.HasDescriptor(psi.MetadataTag)
-
-	if metaDescriptor == nil {
-		return nil, errNoMeta
+	return metaFromPMT(pmt)
 }
-	// Skip the descriptor head.
-	m := metaDescriptor[2:]
 
-	return meta.GetAllAsMap(m)
+// metaFromPMT returns metadata, if any, from a PMT.
+func metaFromPMT(d []byte) (m map[string]string, err error) {
+	// Get as PSI type, skipping the MTS header.
+	pmt := psi.PSIBytes(d[4:])
+
+	// Get the metadata descriptor.
+	_, desc := pmt.HasDescriptor(psi.MetadataTag)
+	if desc == nil {
+		return m, errNoMeta
+	}
+	// Get the metadata as a map, skipping the descriptor head.
+	return meta.GetAllAsMap(desc[2:])
 }
 
 // TrimToMetaRange trims a slice of MPEG-TS to a segment between two points of