diff --git a/container/mts/mpegts.go b/container/mts/mpegts.go index 00028478..338fa235 100644 --- a/container/mts/mpegts.go +++ b/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 { - 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 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 + } + + // 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 } + return metaFromPMT(pmt) +} - // Get as PSI type. Need to skip MTS header. - pmt := psi.PSIBytes(pkt[4:]) - _, metaDescriptor := pmt.HasDescriptor(psi.MetadataTag) +// 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:]) - if metaDescriptor == nil { - return nil, errNoMeta + // Get the metadata descriptor. + _, desc := pmt.HasDescriptor(psi.MetadataTag) + if desc == nil { + return m, errNoMeta } - // Skip the descriptor head. - m := metaDescriptor[2:] - - return meta.GetAllAsMap(m) + // 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