av/container/mts/mpegts.go

760 lines
24 KiB
Go

/*
NAME
mpegts.go - provides a data structure intended to encapsulate the properties
of an MPEG-TS packet and also functions to allow manipulation of these packets.
DESCRIPTION
See Readme.md
AUTHORS
Saxon A. Nelson-Milton <saxon.milton@gmail.com>
Trek Hopton <trek@ausocean.org>
LICENSE
mpegts.go is Copyright (C) 2017-2021 the Australian Ocean Lab (AusOcean)
It is free software: you can redistribute it and/or modify them
under the terms of the GNU General Public License as published by the
Free Software Foundation, either version 3 of the License, or (at your
option) any later version.
It is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
for more details.
You should have received a copy of the GNU General Public License
along with revid in gpl.txt. If not, see [GNU licenses](http://www.gnu.org/licenses).
*/
// Package mts provides MPEGT-TS (mts) encoding and related functions.
package mts
import (
"fmt"
"unsafe"
"github.com/Comcast/gots/packet"
gotspsi "github.com/Comcast/gots/psi"
"github.com/pkg/errors"
"bitbucket.org/ausocean/av/container/mts/meta"
"bitbucket.org/ausocean/av/container/mts/psi"
)
const PacketSize = 188
// Standard program IDs for program specific information MPEG-TS packets.
const (
SdtPid = 17
PatPid = 0
PmtPid = 4096
)
// HeadSize is the size of an MPEG-TS packet header.
const HeadSize = 4
// Consts relating to adaptation field.
const (
AdaptationIdx = 4 // Index to the adaptation field (index of AFL).
AdaptationControlIdx = 3 // Index to octet with adaptation field control.
AdaptationFieldsIdx = AdaptationIdx + 1 // Adaptation field index is the index of the adaptation fields.
DefaultAdaptationSize = 2 // Default size of the adaptation field.
AdaptationControlMask = 0x30 // Mask for the adaptation field control in octet 3.
DefaultAdaptationBodySize = 1 // Default size of the adaptation field body.
DiscontinuityIndicatorMask = 0x80 // Mask for the discontinuity indicator at the discontinuity indicator idk.
DiscontinuityIndicatorIdx = AdaptationIdx + 1 // The index at which the discontinuity indicator is found in an MTS packet.
)
// TODO: make this better - currently doesn't make sense.
const (
HasPayload = 0x1
HasAdaptationField = 0x2
)
/*
Packet encapsulates the fields of an MPEG-TS packet. Below is
the formatting of an MPEG-TS packet for reference!
MPEG-TS Packet Formatting
============================================================================
| octet no | bit 0 | bit 1 | bit 2 | bit 3 | bit 4 | bit 5 | bit 6 | bit 7 |
============================================================================
| octet 0 | sync byte (0x47) |
----------------------------------------------------------------------------
| octet 1 | TEI | PUSI | Prior | PID |
----------------------------------------------------------------------------
| octet 2 | PID cont. |
----------------------------------------------------------------------------
| octet 3 | TSC | AFC | CC |
----------------------------------------------------------------------------
| octet 4 | AFL |
----------------------------------------------------------------------------
| octet 5 | DI | RAI | ESPI | PCRF | OPCRF | SPF | TPDF | AFEF |
----------------------------------------------------------------------------
| optional | PCR (48 bits => 6 bytes) |
----------------------------------------------------------------------------
| - | PCR cont. |
----------------------------------------------------------------------------
| - | PCR cont. |
----------------------------------------------------------------------------
| - | PCR cont. |
----------------------------------------------------------------------------
| - | PCR cont. |
----------------------------------------------------------------------------
| - | PCR cont. |
----------------------------------------------------------------------------
| optional | OPCR (48 bits => 6 bytes) |
----------------------------------------------------------------------------
| - | OPCR cont. |
----------------------------------------------------------------------------
| - | OPCR cont. |
----------------------------------------------------------------------------
| - | OPCR cont. |
----------------------------------------------------------------------------
| - | OPCR cont. |
----------------------------------------------------------------------------
| - | OPCR cont. |
----------------------------------------------------------------------------
| optional | SC |
----------------------------------------------------------------------------
| optional | TPDL |
----------------------------------------------------------------------------
| optional | TPD (variable length) |
----------------------------------------------------------------------------
| - | ... |
----------------------------------------------------------------------------
| optional | Extension (variable length) |
----------------------------------------------------------------------------
| - | ... |
----------------------------------------------------------------------------
| optional | Stuffing (variable length) |
----------------------------------------------------------------------------
| - | ... |
----------------------------------------------------------------------------
| optional | Payload (variable length) |
----------------------------------------------------------------------------
| - | ... |
----------------------------------------------------------------------------
*/
type Packet struct {
TEI bool // Transport Error Indicator
PUSI bool // Payload Unit Start Indicator
Priority bool // Tranposrt priority indicator
PID uint16 // Packet identifier
TSC byte // Transport Scrambling Control
AFC byte // Adaption Field Control
CC byte // Continuity Counter
DI bool // Discontinouty indicator
RAI bool // random access indicator
ESPI bool // Elementary stream priority indicator
PCRF bool // PCR flag
OPCRF bool // OPCR flag
SPF bool // Splicing point flag
TPDF bool // Transport private data flag
AFEF bool // Adaptation field extension flag
PCR uint64 // Program clock reference
OPCR uint64 // Original program clock reference
SC byte // Splice countdown
TPDL byte // Tranposrt private data length
TPD []byte // Private data
Ext []byte // Adaptation field extension
Payload []byte // Mpeg ts Payload
}
// FindPmt will take a clip of MPEG-TS and try to find a PMT table - if one
// is found, then it is returned along with its index, otherwise nil, -1 and an error is returned.
func FindPmt(d []byte) ([]byte, int, error) {
return FindPid(d, PmtPid)
}
// FindPat will take a clip of MPEG-TS and try to find a PAT table - if one
// is found, then it is returned along with its index, otherwise nil, -1 and an error is returned.
func FindPat(d []byte) ([]byte, int, error) {
return FindPid(d, PatPid)
}
// 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")
)
// FindPid will take a clip of MPEG-TS and try to find a packet with given PID - if one
// is found, then it is returned along with its index, otherwise nil, -1 and an error is returned.
func FindPid(d []byte, pid uint16) (pkt []byte, i int, err error) {
if len(d) < PacketSize {
return nil, -1, ErrInvalidLen
}
for i = 0; i < len(d); 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
}
// 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
}
// Errors used by FindPSI.
var (
ErrMultiplePrograms = errors.New("more than one program not supported")
ErrNoPrograms = errors.New("no programs in PAT")
ErrNotConsecutive = errors.New("could not find consecutive PIDs")
)
// FindPSI finds the index of a PAT in an a slice of MPEG-TS and returns, along
// with a map of meta from the PMT and the stream PIDs and their types.
func FindPSI(d []byte) (int, map[uint16]uint8, map[string]string, error) {
if len(d) < PacketSize {
return -1, nil, nil, ErrInvalidLen
}
// Find the PAT if it exists.
pkt, i, err := FindPid(d, PatPid)
if err != nil {
return -1, nil, nil, errors.Wrap(err, "error finding PAT")
}
// Let's take this opportunity to check what programs are in this MPEG-TS
// stream, and therefore the PID of the PMT, from which we can get metadata.
// NB: currently we only support one program.
progs, err := Programs(pkt)
if err != nil {
return i, nil, nil, errors.Wrap(err, "cannot get programs from PAT")
}
if len(progs) == 0 {
return i, nil, nil, ErrNoPrograms
}
if len(progs) > 1 {
return i, nil, nil, ErrMultiplePrograms
}
pmtPID := pmtPIDs(progs)[0]
// Now we can look for the PMT. We want to adjust d so that we're not looking
// at the same data twice.
d = d[i+PacketSize:]
pkt, pmtIdx, err := FindPid(d, pmtPID)
if err != nil {
return i, nil, nil, errors.Wrap(err, "error finding PMT")
}
// Check that the PMT comes straight after the PAT.
if pmtIdx != 0 {
return i, nil, nil, ErrNotConsecutive
}
// Now we can try to get meta from the PMT.
meta, _ := metaFromPMT(pkt)
// Now to get the elementary streams defined for this program.
streams, err := Streams(pkt)
if err != nil {
return i, nil, meta, errors.Wrap(err, "could not get streams from PMT")
}
streamMap := make(map[uint16]uint8)
for _, s := range streams {
streamMap[s.ElementaryPid()] = s.StreamType()
}
return i, streamMap, meta, nil
}
var (
ErrStreamMap = errors.New("stream map is empty")
)
// FirstMediaPID returns the first PID and it's type in the given streamMap.
func FirstMediaPID(streamMap map[uint16]uint8) (p uint16, t uint8, err error) {
for p, t = range streamMap {
return
}
err = ErrStreamMap
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 {
currentPktLen := 6 + asInt(p.PCRF)*6
if len(data) > PacketSize-currentPktLen {
p.Payload = make([]byte, PacketSize-currentPktLen)
} else {
p.Payload = make([]byte, len(data))
}
return copy(p.Payload, data)
}
// Bytes interprets the fields of the ts packet instance and outputs a
// corresponding byte slice
func (p *Packet) Bytes(buf []byte) []byte {
if buf == nil || cap(buf) < PacketSize {
buf = make([]byte, PacketSize)
}
if p.OPCRF {
panic("original program clock reference field unsupported")
}
if p.SPF {
panic("splicing countdown unsupported")
}
if p.TPDF {
panic("transport private data unsupported")
}
if p.AFEF {
panic("adaptation field extension unsupported")
}
buf = buf[:6]
buf[0] = 0x47
buf[1] = (asByte(p.TEI)<<7 | asByte(p.PUSI)<<6 | asByte(p.Priority)<<5 | byte((p.PID&0xFF00)>>8))
buf[2] = byte(p.PID & 0x00FF)
buf[3] = (p.TSC<<6 | p.AFC<<4 | p.CC)
var maxPayloadSize int
if p.AFC&0x2 != 0 {
maxPayloadSize = PacketSize - 6 - asInt(p.PCRF)*6
} else {
maxPayloadSize = PacketSize - 4
}
stuffingLen := maxPayloadSize - len(p.Payload)
if p.AFC&0x2 != 0 {
buf[4] = byte(1 + stuffingLen + asInt(p.PCRF)*6)
buf[5] = (asByte(p.DI)<<7 | asByte(p.RAI)<<6 | asByte(p.ESPI)<<5 | asByte(p.PCRF)<<4 | asByte(p.OPCRF)<<3 | asByte(p.SPF)<<2 | asByte(p.TPDF)<<1 | asByte(p.AFEF))
} else {
buf = buf[:4]
}
for i := 40; p.PCRF && i >= 0; i -= 8 {
buf = append(buf, byte((p.PCR<<15)>>uint(i)))
}
for i := 0; i < stuffingLen; i++ {
buf = append(buf, 0xff)
}
curLen := len(buf)
buf = buf[:PacketSize]
copy(buf[curLen:], p.Payload)
return buf
}
func asInt(b bool) int {
if b {
return 1
}
return 0
}
func asByte(b bool) byte {
if b {
return 1
}
return 0
}
// GotsPacket takes a byte slice and returns a Packet as defined in github.com/comcast/gots/packet.
// TODO: replace this with a type conversion as described here: https://github.com/golang/go/issues/395
func GotsPacket(b []byte) *packet.Packet {
if len(b) != packet.PacketSize {
panic("invalid packet size")
}
return *(**packet.Packet)(unsafe.Pointer(&b))
}
type Option func(p *packet.Packet)
// addAdaptationField adds an adaptation field to p, and applys the passed options to this field.
// TODO: this will probably break if we already have adaptation field.
func addAdaptationField(p *packet.Packet, options ...Option) error {
if packet.ContainsAdaptationField((*packet.Packet)(p)) {
return errors.New("Adaptation field is already present in packet")
}
// Create space for adaptation field.
copy(p[HeadSize+DefaultAdaptationSize:], p[HeadSize:len(p)-DefaultAdaptationSize])
// TODO: seperate into own function
// Update adaptation field control.
p[AdaptationControlIdx] &= 0xff ^ AdaptationControlMask
p[AdaptationControlIdx] |= AdaptationControlMask
// Default the adaptationfield.
resetAdaptation(p)
// Apply and options that have bee passed.
for _, option := range options {
option(p)
}
return nil
}
// resetAdaptation sets fields in ps adaptation field to 0 if the adaptation field
// exists, otherwise an error is returned.
func resetAdaptation(p *packet.Packet) error {
if !packet.ContainsAdaptationField((*packet.Packet)(p)) {
return errors.New("No adaptation field in this packet")
}
p[AdaptationIdx] = DefaultAdaptationBodySize
p[AdaptationIdx+1] = 0x00
return nil
}
// DiscontinuityIndicator returns an Option that will set p's discontinuity
// indicator according to f.
func DiscontinuityIndicator(f bool) Option {
return func(p *packet.Packet) {
set := byte(DiscontinuityIndicatorMask)
if !f {
set = 0x00
}
p[DiscontinuityIndicatorIdx] &= 0xff ^ DiscontinuityIndicatorMask
p[DiscontinuityIndicatorIdx] |= DiscontinuityIndicatorMask & set
}
}
// Error used by GetPTSRange.
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 _pts int64
// Get the first PTS for the given PID.
var i int
for {
if i >= len(clip) {
return pts, errNoPTS
}
pkt, _i, err := FindPid(clip[i:], pid)
if err != nil {
return pts, errors.Wrap(err, fmt.Sprintf("could not find packet of PID: %d", pid))
}
_pts, err = GetPTS(pkt)
if err == nil {
break
}
i += _i + PacketSize
}
pts[0] = uint64(_pts)
pts[1] = pts[0] // Until we have find a second PTS.
// 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
}
_pts, err = GetPTS(pkt)
if err == nil {
break
}
i = _i
}
pts[1] = uint64(_pts)
return
}
var (
errNoPesPayload = errors.New("no PES payload")
errNoPesPTS = errors.New("no PES PTS")
errInvalidPesHeader = errors.New("invalid PES header")
errInvalidPesPayload = errors.New("invalid PES payload")
)
// GetPTS returns a PTS from a packet that has PES payload, or an error otherwise.
func GetPTS(pkt []byte) (pts int64, err error) {
// Check the Payload Unit Start Indicator.
if pkt[1]&0x040 == 0 {
err = errNoPesPayload
return
}
// Compute start of PES payload and check its length.
start := HeadSize
if pkt[3]&0x20 != 0 {
// Adaptation field is present, so adjust start of payload accordingly.
start += 1 + int(pkt[4])
}
pes := pkt[start:]
if len(pes) < 14 {
err = errInvalidPesHeader
return
}
// Check the PTS DTS indicator.
if pes[7]&0xc0 == 0 {
err = errNoPesPTS
return
}
pts = extractPTS(pes[9:14])
return
}
// extractTime extracts a PTS from the given data.
func extractPTS(d []byte) int64 {
return (int64((d[0]>>1)&0x07) << 30) | (int64(d[1]) << 22) | (int64((d[2]>>1)&0x7f) << 15) | (int64(d[3]) << 7) | int64((d[4]>>1)&0x7f)
}
var errNoMeta = errors.New("PMT does not contain meta")
// ExtractMeta returns a map of metadata from the first PMT's metaData
// 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) {
pmt, _, err := FindPid(d, PmtPid)
if err != nil {
return nil, err
}
return metaFromPMT(pmt)
}
// 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[HeadSize:])
// 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
// meta data described by key, from and to.
func TrimToMetaRange(d []byte, key, from, to string) ([]byte, error) {
if len(d)%PacketSize != 0 {
return nil, errors.New("MTS clip is not of valid size")
}
if from == to {
return nil, errors.New("'from' and 'to' cannot be identical")
}
var (
start = -1 // Index of the start of the segment in d.
end = -1 // Index of the end of segment in d.
off int // Index of remaining slice of d to check after each PMT found.
)
for {
// Find the next PMT.
pmt, idx, err := FindPid(d[off:], PmtPid)
if err != nil {
switch -1 {
case start:
return nil, errMetaLowerBound
case end:
return nil, errMetaUpperBound
default:
panic("should not have got error from FindPid")
}
}
off += idx + PacketSize
meta, err := ExtractMeta(pmt)
switch err {
case nil: // do nothing
case errNoMeta:
continue
default:
return nil, err
}
if start == -1 {
if meta[key] == from {
start = off - PacketSize
}
} else if meta[key] == to {
end = off
return d[start:end], nil
}
}
}
// SegmentForMeta returns segments of MTS slice d that correspond to a value of
// meta for key and val. Therefore, any sequence of packets corresponding to
// key and val will be appended to the returned [][]byte.
func SegmentForMeta(d []byte, key, val string) ([][]byte, error) {
var (
pkt packet.Packet // We copy data to this so that we can use comcast gots stuff.
segmenting bool // If true we are currently in a segment corresponsing to given meta.
res [][]byte // The resultant [][]byte holding the segments.
start int // The start index of the current segment.
)
// Go through packets.
for i := 0; i < len(d); i += PacketSize {
copy(pkt[:], d[i:i+PacketSize])
if pkt.PID() == PmtPid {
_meta, err := ExtractMeta(pkt[:])
switch err {
// If there's no meta or a problem with meta, we consider this the end
// of the segment.
case errNoMeta, meta.ErrUnexpectedMetaFormat:
if segmenting {
res = append(res, d[start:i])
segmenting = false
}
continue
case nil: // do nothing.
default:
return nil, err
}
// If we've got the meta of interest in the PMT and we're not segmenting
// then start segmenting. If we don't have the meta of interest in the PMT
// and we are segmenting then we want to stop and append the segment to result.
if _meta[key] == val && !segmenting {
start = i
segmenting = true
} else if _meta[key] != val && segmenting {
res = append(res, d[start:i])
segmenting = false
}
}
}
// We've reached the end of the entire MTS clip so if we're segmenting we need
// to append current segment to res.
if segmenting {
res = append(res, d[start:])
}
return res, nil
}
// PID returns the packet identifier for the given packet.
func PID(p []byte) (uint16, error) {
if len(p) < PacketSize {
return 0, errors.New("packet length less than 188")
}
return uint16(p[1]&0x1f)<<8 | uint16(p[2]), nil
}
// Programs returns a map of program numbers and corresponding PMT PIDs for a
// given MPEG-TS PAT packet.
func Programs(p []byte) (map[uint16]uint16, error) {
pat, err := gotspsi.NewPAT(p)
if err != nil {
return nil, err
}
return pat.ProgramMap(), nil
}
// Streams returns elementary streams defined in a given MPEG-TS PMT packet.
func Streams(p []byte) ([]gotspsi.PmtElementaryStream, error) {
payload, err := Payload(p)
if err != nil {
return nil, errors.Wrap(err, "cannot get packet payload")
}
pmt, err := gotspsi.NewPMT(payload)
if err != nil {
return nil, err
}
return pmt.ElementaryStreams(), nil
}
// MediaStreams retrieves the PmtElementaryStreams from the given PSI. This
// function currently assumes that PSI contain a PAT followed by a PMT directly
// after. We also assume that this MPEG-TS stream contains just one program,
// but this program may contain different streams, i.e. a video stream + audio
// stream.
func MediaStreams(p []byte) ([]gotspsi.PmtElementaryStream, error) {
if len(p) < 2*PacketSize {
return nil, errors.New("PSI is not two packets or more long")
}
pat := p[:PacketSize]
pmt := p[PacketSize : 2*PacketSize]
pid, _ := PID(pat)
if pid != PatPid {
return nil, errors.New("first packet is not a PAT")
}
m, err := Programs(pat)
if err != nil {
return nil, errors.Wrap(err, "could not get programs from PAT")
}
if len(m) == 0 {
return nil, ErrNoPrograms
}
if len(m) > 1 {
return nil, ErrMultiplePrograms
}
pid, _ = PID(pmt)
if pid != pmtPIDs(m)[0] {
return nil, errors.New("second packet is not desired PMT")
}
s, err := Streams(pmt)
if err != nil {
return nil, errors.Wrap(err, "could not get streams from PMT")
}
return s, nil
}
// pmtPIDs returns PMT PIDS from a map containing program number as keys and
// corresponding PMT PIDs as values.
func pmtPIDs(m map[uint16]uint16) []uint16 {
r := make([]uint16, 0, len(m))
for _, v := range m {
r = append(r, v)
}
return r
}
// Errors used by Payload.
var ErrNoPayload = errors.New("no payload")
// Payload returns the payload of an MPEG-TS packet p.
// NB: this is not a copy of the payload in the interest of performance.
// TODO: offer function that will do copy if we have interests in safety.
func Payload(p []byte) ([]byte, error) {
c := byte((p[3] & 0x30) >> 4)
if c == 2 {
return nil, ErrNoPayload
}
// Check if there is an adaptation field.
off := 4
if p[3]&0x20 == 1 {
off = int(5 + p[4])
}
return p[off:], nil
}