av/container/mts/payload.go

package mts

import (
	"errors"
	"sort"

	"github.com/Comcast/gots/packet"
	"github.com/Comcast/gots/pes"
)

// Extract extracts the media, PTS, stream ID and meta for an MPEG-TS clip given
// by p, and returns as a Clip. The MPEG-TS must contain only complete packets.
func Extract(p []byte) (*Clip, error) {
	l := len(p)
	// Check that clip is divisible by 188, i.e. contains a series of full MPEG-TS clips.
	if l%PacketSize != 0 {
		return nil, errors.New("MTS clip is not of valid size")
	}

	var (
		frameStart  int               // Index used to indicate the start of current frame in backing slice.
		clip        = &Clip{}         // The data that will be returned.
		meta        map[string]string // Holds the most recently extracted meta.
		lenOfFrame  int               // Len of current frame.
		dataLen     int               // Len of data from MPEG-TS packet.
		curPTS      uint64            // Holds the current PTS.
		curStreamID uint8             // Holds current StreamID (shouldn't change)
		firstPUSI   = true            // Indicates that we have not yet received a PUSI.
		err         error
	)

	// This will hold a copy of all the media in the MPEG-TS clip.
	clip.backing = make([]byte, 0, l/PacketSize)

	// Go through the MPEGT-TS packets.
	var pkt packet.Packet
	for i := 0; i < l; i += PacketSize {
		// We will use comcast/gots Packet type, so copy in.
		copy(pkt[:], p[i:i+PacketSize])

		switch pkt.PID() {
		case PatPid: // Do nothing.
		case PmtPid:
			meta, err = ExtractMeta(pkt[:])
			if err != nil {
				return nil, err
			}
		default: // Must be media.
			// Get the MPEG-TS payload.
			payload, err := pkt.Payload()
			if err != nil {
				return nil, err
			}

			// If PUSI is true then we know it's the start of a new frame, and we have
			// a PES header in the MTS payload.

			if pkt.PayloadUnitStartIndicator() {
				_pes, err := pes.NewPESHeader(payload)
				if err != nil {
					return nil, err
				}

				// Extract the PTS and ID, then add a new frame to the clip.
				curPTS = _pes.PTS()
				curStreamID = _pes.StreamId()
				clip.frames = append(clip.frames, Frame{
					PTS:  curPTS,
					ID:   curStreamID,
					Meta: meta,
				})

				// Append the data to the underlying buffer and get appended lenghth.
				clip.backing = append(clip.backing, _pes.Data()...)
				dataLen = len(_pes.Data())

				// If we haven't hit the first PUSI, then we know we have a full frame
				// and can add this data to the frame pertaining to the finish frame.
				if !firstPUSI {
					clip.frames[len(clip.frames)-2].Media = clip.backing[frameStart:lenOfFrame]
					clip.frames[len(clip.frames)-2].idx = frameStart
					frameStart = lenOfFrame
				}
				firstPUSI = false
			} else {
				// We're not at the start of the frame, so we don't have a PES header.
				// We can append the MPEG-TS data directly to the underlying buf.
				dataLen = len(payload)
				clip.backing = append(clip.backing, payload...)
			}
			lenOfFrame += dataLen
		}
	}
	// We're finished up with media frames, so give the final Frame it's data.
	clip.frames[len(clip.frames)-1].Media = clip.backing[frameStart:lenOfFrame]
	clip.frames[len(clip.frames)-1].idx = frameStart
	return clip, nil
}

// Clip represents a clip of media, i.e. a sequence of media frames.
type Clip struct {
	frames  []Frame
	backing []byte
}

// Frame describes a media frame that may be extracted from a PES packet.
type Frame struct {
	Media []byte            // Contains the media from the frame.
	PTS   uint64            // PTS from PES packet (this gives time relative from start of stream).
	ID    uint8             // StreamID from the PES packet, identifying media codec.
	Meta  map[string]string // Contains metadata from PMT relevant to this frame.
	idx   int               // Index in the backing slice.
}

// Bytes returns the concatentated media bytes from each frame in the Clip c.
func (c *Clip) Bytes() []byte {
	if c.backing == nil {
		panic("the clip backing array cannot be nil")
	}
	return c.backing
}

// Errors used in BytesForPTSInterval.
var (
	errPTSLowerBound = errors.New("PTS 'from' cannot be found")
	errPTSUpperBound = errors.New("PTS 'to' cannot be found")
	errPTSRange      = errors.New("PTS interval invalid")
)

// BytesForPTSInterval returns the media data between PTS' from and to. If from
// sits between two PTS, the Frame posessing lower PTS will be considered the start.
// The Frame before the Frame corresponding to to will be considered the final
// Frame.
func (c *Clip) BytesForPTSInterval(from, to uint64) ([]byte, error) {
	// First check that the interval makes sense.
	if from >= to {
		return nil, errPTSRange
	}

	// Use binary search to find 'from'.
	n := len(c.frames) - 1
	idx := sort.Search(
		n,
		func(i int) bool {
			if from < c.frames[i+1].PTS {
				return true
			}
			return false
		},
	)
	if idx == n {
		return nil, errPTSLowerBound
	}

	// Now get the start index for the backing slice from this Frame.
	start := c.frames[idx].idx

	// Now use binary search again to find 'to'.
	off := idx + 1
	n = n - (off)
	idx = sort.Search(
		n,
		func(i int) bool {
			if to <= c.frames[i+off].PTS {
				return true
			}
			return false
		},
	)
	if idx == n {
		return nil, errPTSUpperBound
	}

	// Now get the end index for the backing slice from this Frame, and return
	// segment from backing slice corresponding to start and end.
	end := c.frames[idx+off-1].idx
	return c.backing[start : end+len(c.frames[idx+off].Media)], nil
}

// Errors that maybe returned from BytesForMetaInterval.
var (
	errMetaRange      = errors.New("invalid meta range")
	errMetaLowerBound = errors.New("meta 'from' cannot be found")
	errMetaUpperBound = errors.New("meta 'to' cannot be found")
)

func (c *Clip) BytesForMetaInterval(key, from, to string) ([]byte, error) {
	// First check that the interval makes sense.
	if from == to {
		return nil, errMetaRange
	}

	var start, end int

	// Try and find from.
	for i := 0; i < len(c.frames); i++ {
		f := c.frames[i]
		if f.Meta[key] == from {
			start = f.idx

			// Now try and find to.
			for ; i < len(c.frames); i++ {
				f = c.frames[i]
				if f.Meta[key] == to {
					end = f.idx
					return c.backing[start : end+len(f.Media)], nil
				}
			}
			return nil, errMetaUpperBound
		}
	}
	return nil, errMetaLowerBound
}