av/container/mts/mpegts.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

AUTHOR
  Saxon A. Nelson-Milton <saxon.milton@gmail.com>

LICENSE
  mpegts.go is Copyright (C) 2017 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 (
	"errors"

	"bitbucket.org/ausocean/av/container/mts/meta"
	"bitbucket.org/ausocean/av/container/mts/psi"
	"github.com/Comcast/gots/packet"
	"github.com/Comcast/gots/pes"
)

const PacketSize = 188

// Program ID for various types of ts packets.
const (
	SdtPid   = 17
	PatPid   = 0
	PmtPid   = 4096
	VideoPid = 256
)

// StreamID is the id of the first stream.
const StreamID = 0xe0

// 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)
}

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
}

// 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
}

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 and 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
	}
}

// GetPTSRange retreives the first and last PTS of an MPEGTS clip.
func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) {
	// Find the first packet with PID pidType.
	pkt, _, err := FindPid(clip, pid)
	if err != nil {
		return [2]uint64{}, err
	}

	// Get the payload of the packet, which will be the start of the PES packet.
	var _pkt packet.Packet
	copy(_pkt[:], pkt)
	payload, err := packet.Payload(&_pkt)
	if err != nil {
		return [2]uint64{}, err
	}

	// Get the the first PTS from the PES header.
	_pes, err := pes.NewPESHeader(payload)
	if err != nil {
		return [2]uint64{}, err
	}
	pts[0] = _pes.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
		}
	}

	return [2]uint64{}, errors.New("could only find one access unit in mpegts clip")
}

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) {
	pkt, _, 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
	}
	// Skip the descriptor head.
	m := metaDescriptor[2:]

	return meta.GetAllAsMap(m)
}

// 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 equivalent")
	}

	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
		}
	}
}

func SegmentForMeta(d []byte, key, val string) ([][]byte, error) {
	var (
		pkt        packet.Packet
		segmenting bool
		res        [][]byte
		start      int
	)

	for i := 0; i < len(d); i += PacketSize {
		copy(pkt[:], d[i:i+PacketSize])
		if pkt.PID() == PmtPid {
			_meta, err := ExtractMeta(pkt[:])
			switch err {
			case errNoMeta, meta.ErrUnexpectedMetaFormat:
				if segmenting {
					res = append(res, d[start:i])
					segmenting = false
				}
				continue
			case nil: // do nothing.
			default:
				return nil, err
			}

			if _meta[key] == val && !segmenting {
				start = i
				segmenting = true
			} else if _meta[key] != val && segmenting {
				res = append(res, d[start:i])
				segmenting = false
			}
		}
	}

	if segmenting {
		res = append(res, d[start:len(d)])
	}

	return res, nil
}