mirror
/
av
mirror of https://bitbucket.org/ausocean/av.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merged in mts-payload-extract (pull request #206) 

container/mts: MTS payload extraction and further development of MTS utilities.

Approved-by: Alan Noble <anoble@gmail.com>
This commit is contained in:
Saxon Milton 2019-07-05 03:00:29 +00:00
parent 3e3fd39251 36f1e8c78b
commit aeb10e0ab7
8 changed files with 1248 additions and 105 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

86
container/mts/encoder_test.go
View File

@ -29,12 +29,14 @@ import (
"bytes"
"io"
"io/ioutil"
"reflect"
"testing"
"github.com/Comcast/gots/packet"
"github.com/Comcast/gots/pes"
"bitbucket.org/ausocean/av/container/mts/meta"
"bitbucket.org/ausocean/av/container/mts/psi"
)
type nopCloser struct{ io.Writer }
@ -250,3 +252,87 @@ func TestEncodePcm(t *testing.T) {
t.Error("data decoded from mts did not match input data")
}
}
const fps = 25
// TestMetaEncode1 checks that we can externally add a single metadata entry to
// the mts global Meta meta.Data struct and then successfully have the mts encoder
// write this to psi.
func TestMetaEncode1(t *testing.T) {
Meta = meta.New()
var buf bytes.Buffer
e := NewEncoder(nopCloser{&buf}, fps, EncodeH264)
Meta.Add("ts", "12345678")
if err := e.writePSI(); err != nil {
t.Errorf("unexpected error: %v\n", err.Error())
}
out := buf.Bytes()
got := out[PacketSize+4:]
want := []byte{
0x00, 0x02, 0xb0, 0x23, 0x00, 0x01, 0xc1, 0x00, 0x00, 0xe1, 0x00, 0xf0, 0x11,
psi.MetadataTag, // Descriptor tag
0x0f, // Length of bytes to follow
0x00, 0x10, 0x00, 0x0b, 't', 's', '=', '1', '2', '3', '4', '5', '6', '7', '8', // timestamp
0x1b, 0xe1, 0x00, 0xf0, 0x00,
}
want = psi.AddCrc(want)
want = psi.AddPadding(want)
if !bytes.Equal(got, want) {
t.Errorf("unexpected output. \n Got : %v\n, Want: %v\n", got, want)
}
}
// TestMetaEncode2 checks that we can externally add two metadata entries to the
// Meta meta.Data global and then have the mts encoder successfully encode this
// into psi.
func TestMetaEncode2(t *testing.T) {
Meta = meta.New()
var buf bytes.Buffer
e := NewEncoder(nopCloser{&buf}, fps, EncodeH264)
Meta.Add("ts", "12345678")
Meta.Add("loc", "1234,4321,1234")
if err := e.writePSI(); err != nil {
t.Errorf("did not expect error: %v from writePSI", err.Error())
}
out := buf.Bytes()
got := out[PacketSize+4:]
want := []byte{
0x00, 0x02, 0xb0, 0x36, 0x00, 0x01, 0xc1, 0x00, 0x00, 0xe1, 0x00, 0xf0, 0x24,
psi.MetadataTag, // Descriptor tag
0x22, // Length of bytes to follow
0x00, 0x10, 0x00, 0x1e, 't', 's', '=', '1', '2', '3', '4', '5', '6', '7', '8', '\t', // timestamp
'l', 'o', 'c', '=', '1', '2', '3', '4', ',', '4', '3', '2', '1', ',', '1', '2', '3', '4', // location
0x1b, 0xe1, 0x00, 0xf0, 0x00,
}
want = psi.AddCrc(want)
want = psi.AddPadding(want)
if !bytes.Equal(got, want) {
t.Errorf("did not get expected results.\ngot: %v\nwant: %v\n", got, want)
}
}
// TestExtractMeta checks that ExtractMeta can correclty get a map of metadata
// from the first PMT in a clip of MPEG-TS.
func TestExtractMeta(t *testing.T) {
Meta = meta.New()
var buf bytes.Buffer
e := NewEncoder(nopCloser{&buf}, fps, EncodeH264)
Meta.Add("ts", "12345678")
Meta.Add("loc", "1234,4321,1234")
if err := e.writePSI(); err != nil {
t.Errorf("did not expect error: %v", err.Error())
}
out := buf.Bytes()
got, err := ExtractMeta(out)
if err != nil {
t.Errorf("did not expect error: %v", err.Error())
}
want := map[string]string{
"ts": "12345678",
"loc": "1234,4321,1234",
}
if !reflect.DeepEqual(got, want) {
t.Errorf("did not get expected result.\ngot: %v\nwant: %v\n", got, want)
}
}

32
container/mts/meta/meta.go
View File

@ -51,7 +51,7 @@ const (
var (
errKeyAbsent = errors.New("Key does not exist in map")
errInvalidMeta = errors.New("Invalid metadata given")
errUnexpectedMetaFormat = errors.New("Unexpected meta format")
ErrUnexpectedMetaFormat = errors.New("Unexpected meta format")
)
// Metadata provides functionality for the storage and encoding of metadata
@ -209,13 +209,41 @@ func GetAll(d []byte) ([][2]string, error) {
for i, entry := range entries {
kv := strings.Split(entry, "=")
if len(kv) != 2 {
return nil, errUnexpectedMetaFormat
return nil, ErrUnexpectedMetaFormat
}
copy(all[i][:], kv)
}
return all, nil
}
// GetAllAsMap returns a map containing keys and values from a slice d containing
// metadata.
func GetAllAsMap(d []byte) (map[string]string, error) {
err := checkMeta(d)
if err != nil {
return nil, err
}
// Skip the header, which is our data length and version.
d = d[headSize:]
// Each metadata entry (key and value) is seperated by a tab, so split at tabs
// to get individual entries.
entries := strings.Split(string(d), "\t")
// Go through entries and add to all map.
all := make(map[string]string)
for _, entry := range entries {
// Keys and values are seperated by '=', so split and check that len(kv)=2.
kv := strings.Split(entry, "=")
if len(kv) != 2 {
return nil, ErrUnexpectedMetaFormat
}
all[kv[0]] = kv[1]
}
return all, nil
}
// checkHeader checks that a valid metadata header exists in the given data.
func checkMeta(d []byte) error {
if len(d) == 0 || d[0] != 0 || binary.BigEndian.Uint16(d[2:headSize]) != uint16(len(d[headSize:])) {

17
container/mts/meta/meta_test.go
View File

@ -189,6 +189,23 @@ func TestGetAll(t *testing.T) {
}
}
// TestGetAllAsMap checks that GetAllAsMap will correctly return a map of meta
// keys and values from a slice of meta.
func TestGetAllAsMap(t *testing.T) {
tstMeta := append([]byte{0x00, 0x10, 0x00, 0x12}, "loc=a,b,c\tts=12345"...)
want := map[string]string{
"loc": "a,b,c",
"ts": "12345",
}
got, err := GetAllAsMap(tstMeta)
if err != nil {
t.Errorf("Unexpected error: %v\n", err)
}
if !reflect.DeepEqual(got, want) {
t.Errorf("Did not get expected out. \nGot : %v, \nWant: %v\n", got, want)
}
}
// TestKeys checks that we can successfully get keys from some metadata using
// the meta.Keys method.
func TestKeys(t *testing.T) {

100
container/mts/metaEncode_test.go
View File

@ -1,100 +0,0 @@
/*
NAME
metaEncode_test.go
DESCRIPTION
See Readme.md
AUTHOR
Saxon Nelson-Milton <saxon@ausocean.org>
LICENSE
metaEncode_test.go is Copyright (C) 2017-2019 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 http://www.gnu.org/licenses.
*/
package mts
import (
"bytes"
"testing"
"bitbucket.org/ausocean/av/container/mts/meta"
"bitbucket.org/ausocean/av/container/mts/psi"
)
const (
errNotExpectedOut = "Unexpected output. \n Got : %v\n, Want: %v\n"
errUnexpectedErr = "Unexpected error: %v\n"
)
const fps = 25
// TestMetaEncode1 checks that we can externally add a single metadata entry to
// the mts global Meta meta.Data struct and then successfully have the mts encoder
// write this to psi.
func TestMetaEncode1(t *testing.T) {
Meta = meta.New()
var buf bytes.Buffer
e := NewEncoder(nopCloser{&buf}, fps, EncodeH264)
Meta.Add("ts", "12345678")
if err := e.writePSI(); err != nil {
t.Errorf(errUnexpectedErr, err.Error())
}
out := buf.Bytes()
got := out[PacketSize+4:]
want := []byte{
0x00, 0x02, 0xb0, 0x23, 0x00, 0x01, 0xc1, 0x00, 0x00, 0xe1, 0x00, 0xf0, 0x11,
psi.MetadataTag, // Descriptor tag
0x0f, // Length of bytes to follow
0x00, 0x10, 0x00, 0x0b, 't', 's', '=', '1', '2', '3', '4', '5', '6', '7', '8', // timestamp
0x1b, 0xe1, 0x00, 0xf0, 0x00,
}
want = psi.AddCrc(want)
want = psi.AddPadding(want)
if !bytes.Equal(got, want) {
t.Errorf(errNotExpectedOut, got, want)
}
}
// TestMetaEncode2 checks that we can externally add two metadata entries to the
// Meta meta.Data global and then have the mts encoder successfully encode this
// into psi.
func TestMetaEncode2(t *testing.T) {
Meta = meta.New()
var buf bytes.Buffer
e := NewEncoder(nopCloser{&buf}, fps, EncodeH264)
Meta.Add("ts", "12345678")
Meta.Add("loc", "1234,4321,1234")
if err := e.writePSI(); err != nil {
t.Errorf(errUnexpectedErr, err.Error())
}
out := buf.Bytes()
got := out[PacketSize+4:]
want := []byte{
0x00, 0x02, 0xb0, 0x36, 0x00, 0x01, 0xc1, 0x00, 0x00, 0xe1, 0x00, 0xf0, 0x24,
psi.MetadataTag, // Descriptor tag
0x22, // Length of bytes to follow
0x00, 0x10, 0x00, 0x1e, 't', 's', '=', '1', '2', '3', '4', '5', '6', '7', '8', '\t', // timestamp
'l', 'o', 'c', '=', '1', '2', '3', '4', ',', '4', '3', '2', '1', ',', '1', '2', '3', '4', // location
0x1b, 0xe1, 0x00, 0xf0, 0x00,
}
want = psi.AddCrc(want)
want = psi.AddPadding(want)
if !bytes.Equal(got, want) {
t.Errorf(errNotExpectedOut, got, want)
}
}

144
container/mts/mpegts.go
View File

@ -35,6 +35,9 @@ import (
"github.com/Comcast/gots/packet"
"github.com/Comcast/gots/pes"
"github.com/pkg/errors"
"bitbucket.org/ausocean/av/container/mts/meta"
"bitbucket.org/ausocean/av/container/mts/psi"
)
const PacketSize = 188
@ -173,11 +176,17 @@ 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, errors.New("MPEG-TS data not of valid length")
return nil, -1, errInvalidLen
}
for i = 0; i < len(d); i += PacketSize {
p := (uint16(d[i+1]&0x1f) << 8) | uint16(d[i+2])
@ -186,7 +195,7 @@ func FindPid(d []byte, pid uint16) (pkt []byte, i int, err error) {
return
}
}
return nil, -1, fmt.Errorf("could not find packet with pid: %d", pid)
return nil, -1, errCouldNotFind
}
// FillPayload takes a channel and fills the packets Payload field until the
@ -318,6 +327,9 @@ func DiscontinuityIndicator(f bool) Option {
}
}
// Error used by GetPTSRange.
var errNoPTS = errors.New("could not find PTS")
// GetPTSRange retreives the first and last PTS of an MPEGTS clip.
func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) {
var _pkt packet.Packet
@ -370,4 +382,130 @@ func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) {
return
}
var errNoPTS = errors.New("could not find PTS")
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 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:len(d)])
}
return res, nil
}

166
container/mts/mpegts_test.go
View File

@ -30,9 +30,12 @@ package mts
import (
"bytes"
"math/rand"
"reflect"
"strconv"
"testing"
"time"
"bitbucket.org/ausocean/av/container/mts/meta"
"bitbucket.org/ausocean/av/container/mts/pes"
"bitbucket.org/ausocean/av/container/mts/psi"
"github.com/Comcast/gots/packet"
@ -344,3 +347,166 @@ func TestFindPid(t *testing.T) {
t.Errorf("index of found packet is not correct.\nGot: %v, want: %v\n", _got, targetPacketNum)
}
}
// TestTrimToMetaRange checks that TrimToMetaRange can correctly return a segment
// of MPEG-TS corresponding to a meta interval in a slice of MPEG-TS.
func TestTrimToMetaRange(t *testing.T) {
Meta = meta.New()
const (
nPSI = 10
key = "n"
)
var clip bytes.Buffer
for i := 0; i < nPSI; i++ {
Meta.Add(key, strconv.Itoa((i*2)+1))
err := writePSIWithMeta(&clip)
if err != nil {
t.Fatalf("did not expect to get error writing PSI, error: %v", err)
}
}
tests := []struct {
from string
to string
expect []byte
err error
}{
{
from: "3",
to: "9",
expect: clip.Bytes()[3*PacketSize : 10*PacketSize],
err: nil,
},
{
from: "30",
to: "8",
expect: nil,
err: errMetaLowerBound,
},
{
from: "3",
to: "30",
expect: nil,
err: errMetaUpperBound,
},
}
// Run tests.
for i, test := range tests {
got, err := TrimToMetaRange(clip.Bytes(), key, test.from, test.to)
// First check the error.
if err != nil && err != test.err {
t.Errorf("unexpected error: %v for test: %v", err, i)
continue
} else if err != test.err {
t.Errorf("expected to get error: %v for test: %v", test.err, i)
continue
}
// Now check data.
if test.err == nil && !bytes.Equal(test.expect, got) {
t.Errorf("did not get expected data for test: %v\n Got: %v\n, Want: %v\n", i, got, test.expect)
}
}
}
// TestSegmentForMeta checks that SegmentForMeta can correctly segment some MTS
// data based on a given meta key and value.
func TestSegmentForMeta(t *testing.T) {
Meta = meta.New()
const (
nPSI = 10 // The number of PSI pairs to write.
key = "n" // The meta key we will work with.
val = "*" // This is the meta value we will look for.
)
tests := []struct {
metaVals [nPSI]string // This represents the meta value for meta pairs (PAT and PMT)
expectIdxs []rng // This gives the expect index ranges for the segments.
}{
{
metaVals: [nPSI]string{"1", "2", val, val, val, "3", val, val, "4", "4"},
expectIdxs: []rng{
scale(2, 5),
scale(6, 8),
},
},
{
metaVals: [nPSI]string{"1", "2", val, val, val, "", "3", val, val, "4"},
expectIdxs: []rng{
scale(2, 5),
scale(7, 9),
},
},
{
metaVals: [nPSI]string{"1", "2", val, val, val, "", "3", val, val, val},
expectIdxs: []rng{
scale(2, 5),
{((7 * 2) + 1) * PacketSize, (nPSI * 2) * PacketSize},
},
},
{
metaVals: [nPSI]string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
expectIdxs: nil,
},
}
var clip bytes.Buffer
for testn, test := range tests {
// We want a clean buffer for each new test, so reset.
clip.Reset()
// Add meta and write PSI to clip.
for i := 0; i < nPSI; i++ {
if test.metaVals[i] != "" {
Meta.Add(key, test.metaVals[i])
} else {
Meta.Delete(key)
}
err := writePSIWithMeta(&clip)
if err != nil {
t.Fatalf("did not expect to get error writing PSI, error: %v", err)
}
}
// Now we get the expected segments using the index ranges from the test.
var want [][]byte
for _, idxs := range test.expectIdxs {
want = append(want, clip.Bytes()[idxs.start:idxs.end])
}
// Now use the function we're testing to get the segments.
got, err := SegmentForMeta(clip.Bytes(), key, val)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
// Check that segments are OK.
if !reflect.DeepEqual(want, got) {
t.Errorf("did not get expected result for test %v\nGot: %v\nWant: %v\n", testn, got, want)
}
}
}
// rng describes an index range and is used by TestSegmentForMeta.
type rng struct {
start int
end int
}
// scale takes a PSI index (i.e. first PSI is 0, next is 1) and modifies to be
// the index of the first byte of the PSI pair (PAT and PMT) in the byte stream.
// This assumes there are only PSI written consequitively, and is used by
// TestSegmentForMeta.
func scale(x, y int) rng {
return rng{
((x * 2) + 1) * PacketSize,
((y * 2) + 1) * PacketSize,
}
}

306
container/mts/payload.go Normal file
View File

@ -0,0 +1,306 @@
/*
NAME
payload.go
DESCRIPTION
payload.go provides functionality for extracting and manipulating the payload
data from MPEG-TS.
AUTHOR
Saxon A. Nelson-Milton <saxon@ausocean.org>
LICENSE
Copyright (C) 2019 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
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.
// The resultant data is a copy of the original.
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 length.
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 TrimToPTSRange.
var (
errPTSLowerBound = errors.New("PTS 'from' cannot be found")
errPTSUpperBound = errors.New("PTS 'to' cannot be found")
errPTSRange = errors.New("PTS interval invalid")
)
// TrimToPTSRange returns the sub Clip in a PTS range defined by from and to.
// The first Frame in the new Clip will be the Frame for which from corresponds
// exactly with Frame.PTS, or the Frame in which from lies within. The final
// Frame in the Clip will be the previous of that for which to coincides with,
// or the Frame that to lies within.
func (c *Clip) TrimToPTSRange(from, to uint64) (*Clip, 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
startFrameIdx := sort.Search(
n,
func(i int) bool {
if from < c.frames[i+1].PTS {
return true
}
return false
},
)
if startFrameIdx == n {
return nil, errPTSLowerBound
}
// Now get the start index for the backing slice from this Frame.
startBackingIdx := c.frames[startFrameIdx].idx
// Now use binary search again to find 'to'.
off := startFrameIdx + 1
n = n - (off)
endFrameIdx := sort.Search(
n,
func(i int) bool {
if to <= c.frames[i+off].PTS {
return true
}
return false
},
)
if endFrameIdx == n {
return nil, errPTSUpperBound
}
// Now get the end index for the backing slice from this Frame.
endBackingIdx := c.frames[endFrameIdx+off-1].idx
// Now return a new clip. NB: data is not copied.
return &Clip{
frames: c.frames[startFrameIdx : endFrameIdx+1],
backing: c.backing[startBackingIdx : endBackingIdx+len(c.frames[endFrameIdx+off].Media)],
}, nil
}
// Errors that maybe returned from TrimToMetaRange.
var (
errMetaRange = errors.New("invalid meta range")
errMetaLowerBound = errors.New("meta 'from' cannot be found")
errMetaUpperBound = errors.New("meta 'to' cannot be found")
)
// TrimToMetaRange returns a sub Clip with meta range described by from and to
// with key 'key'. The meta values must not be equivalent.
func (c *Clip) TrimToMetaRange(key, from, to string) (*Clip, 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]
startFrameIdx := 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
endFrameIdx := i
return &Clip{
frames: c.frames[startFrameIdx : endFrameIdx+1],
backing: c.backing[start : end+len(f.Media)],
}, nil
}
}
return nil, errMetaUpperBound
}
}
return nil, errMetaLowerBound
}
// SegmentForMeta segments sequences of frames within c possesing meta described
// by key and val and are appended to a []Clip which is subsequently returned.
func (c *Clip) SegmentForMeta(key, val string) []Clip {
var (
segmenting bool // If true we are currently in a segment corresponsing to given meta.
res []Clip // The resultant [][]Clip holding the segments.
start int // The start index of the current segment.
)
// Go through frames of clip.
for i, frame := range c.frames {
// If there is no meta (meta = nil) and we are segmenting, then append the
// current segment to res.
if frame.Meta == nil {
if segmenting {
res = appendSegment(res, c, start, i)
segmenting = false
}
continue
}
// If we've got the meta of interest in current frame and we're not
// segmenting, set this i as start and set segmenting true. If we don't
// have the meta of interest and we are segmenting then we
// want to stop and append the segment to res.
if frame.Meta[key] == val && !segmenting {
start = i
segmenting = true
} else if frame.Meta[key] != val && segmenting {
res = appendSegment(res, c, start, i)
segmenting = false
}
}
// We've reached the end of the entire clip so if we're segmenting we need
// to append current segment to res.
if segmenting {
res = appendSegment(res, c, start, len(c.frames))
}
return res
}
// appendSegment is a helper function used by Clip.SegmentForMeta to append a
// clip to a []Clip.
func appendSegment(segs []Clip, c *Clip, start, end int) []Clip {
return append(segs, Clip{
frames: c.frames[start:end],
backing: c.backing[c.frames[start].idx : c.frames[end-1].idx+len(c.frames[end-1].Media)],
},
)
}

502
container/mts/payload_test.go Normal file
View File

@ -0,0 +1,502 @@
/*
NAME
payload_test.go
DESCRIPTION
payload_test.go provides testing to validate utilities found in payload.go.
AUTHOR
Saxon A. Nelson-Milton <saxon@ausocean.org>
LICENSE
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
import (
"bytes"
"math/rand"
"reflect"
"strconv"
"testing"
"time"
"bitbucket.org/ausocean/av/container/mts/meta"
"bitbucket.org/ausocean/av/container/mts/psi"
)
// TestExtract checks that we can coorectly extract media, pts, id and meta from
// an MPEGTS stream using Extract.
func TestExtract(t *testing.T) {
Meta = meta.New()
const (
psiInterval = 5 // Write PSI at start and after every 5 frames.
numOfFrames = 30 // Total number of frames to write.
maxFrameSize = 1000 // Max frame size to randomly generate.
minFrameSize = 100 // Min frame size to randomly generate.
rate = 25 // Framerate (fps)
interval = float64(1) / rate // Time interval between frames.
ptsFreq = 90000 // Standard PTS frequency base.
)
frames := genFrames(numOfFrames, minFrameSize, maxFrameSize)
var (
clip bytes.Buffer // This will hold the MPEG-TS data.
want Clip // This is the Clip that we should get.
err error
)
// Now write frames.
var curTime float64
for i, frame := range frames {
// Check to see if it's time to write another lot of PSI.
if i%psiInterval == 0 && i != len(frames)-1 {
// We'll add the frame number as meta.
Meta.Add("frameNum", strconv.Itoa(i))
err = writePSIWithMeta(&clip)
if err != nil {
t.Fatalf("did not expect error writing psi: %v", err)
}
}
nextPTS := uint64(curTime * ptsFreq)
err = writeFrame(&clip, frame, uint64(nextPTS))
if err != nil {
t.Fatalf("did not expect error writing frame: %v", err)
}
curTime += interval
// Need the meta map for the new expected Frame.
metaMap, err := meta.GetAllAsMap(Meta.Encode())
if err != nil {
t.Fatalf("did not expect error getting meta map: %v", err)
}
// Create an equivalent Frame and append to our Clip want.
want.frames = append(want.frames, Frame{
Media: frame,
PTS: nextPTS,
ID: H264ID,
Meta: metaMap,
})
}
// Now use Extract to get frames from clip.
got, err := Extract(clip.Bytes())
if err != nil {
t.Fatalf("did not expect error using Extract. Err: %v", err)
}
// Check length of got and want.
if len(want.frames) != len(got.frames) {
t.Fatalf("did not get expected length for got.\nGot: %v\n, Want: %v\n", len(got.frames), len(want.frames))
}
// Check frames individually.
for i, frame := range want.frames {
// Check media data.
wantMedia := frame.Media
gotMedia := got.frames[i].Media
if !bytes.Equal(wantMedia, gotMedia) {
t.Fatalf("did not get expected data for frame: %v\nGot: %v\nWant: %v\n", i, gotMedia, wantMedia)
}
// Check stream ID.
wantID := frame.ID
gotID := got.frames[i].ID
if wantID != gotID {
t.Fatalf("did not get expected ID for frame: %v\nGot: %v\nWant: %v\n", i, gotID, wantID)
}
// Check meta.
wantMeta := frame.Meta
gotMeta := got.frames[i].Meta
if !reflect.DeepEqual(wantMeta, gotMeta) {
t.Fatalf("did not get expected meta for frame: %v\nGot: %v\nwant: %v\n", i, gotMeta, wantMeta)
}
}
}
// writePSIWithMeta writes PSI to b with updated metadata.
func writePSIWithMeta(b *bytes.Buffer) error {
// Write PAT.
pat := Packet{
PUSI: true,
PID: PatPid,
CC: 0,
AFC: HasPayload,
Payload: psi.AddPadding(patTable),
}
_, err := b.Write(pat.Bytes(nil))
if err != nil {
return err
}
// Update the meta in the pmt table.
pmtTable, err = updateMeta(pmtTable)
if err != nil {
return err
}
// Write PMT.
pmt := Packet{
PUSI: true,
PID: PmtPid,
CC: 0,
AFC: HasPayload,
Payload: psi.AddPadding(pmtTable),
}
_, err = b.Write(pmt.Bytes(nil))
if err != nil {
return err
}
return nil
}
// TestClipBytes checks that Clip.Bytes correctly returns the concatendated media
// data from the Clip's frames slice.
func TestClipBytes(t *testing.T) {
Meta = meta.New()
const (
psiInterval = 5 // Write PSI at start and after every 5 frames.
numOfFrames = 30 // Total number of frames to write.
maxFrameSize = 1000 // Max frame size to randomly generate.
minFrameSize = 100 // Min frame size to randomly generate.
rate = 25 // Framerate (fps)
interval = float64(1) / rate // Time interval between frames.
ptsFreq = 90000 // Standard PTS frequency base.
)
frames := genFrames(numOfFrames, minFrameSize, maxFrameSize)
var (
clip bytes.Buffer // This will hold the MPEG-TS data.
want []byte // This is the Clip that we should get.
err error
)
// Now write frames.
var curTime float64
for i, frame := range frames {
// Check to see if it's time to write another lot of PSI.
if i%psiInterval == 0 && i != len(frames)-1 {
// We'll add the frame number as meta.
Meta.Add("frameNum", strconv.Itoa(i))
err = writePSIWithMeta(&clip)
if err != nil {
t.Fatalf("did not expect error writing psi: %v", err)
}
}
nextPTS := uint64(curTime * ptsFreq)
err = writeFrame(&clip, frame, uint64(nextPTS))
if err != nil {
t.Fatalf("did not expect error writing frame: %v", err)
}
curTime += interval
// Append the frame straight to the expected pure media slice.
want = append(want, frame...)
}
// Now use Extract to get Clip and then use Bytes to get the slice of straight media.
gotClip, err := Extract(clip.Bytes())
if err != nil {
t.Fatalf("did not expect error using Extract. Err: %v", err)
}
got := gotClip.Bytes()
// Check length and equality of got and want.
if len(want) != len(got) {
t.Fatalf("did not get expected length for got.\nGot: %v\n, Want: %v\n", len(got), len(want))
}
if !bytes.Equal(want, got) {
t.Error("did not get expected result")
}
}
// genFrames is a helper function to generate a series of dummy media frames
// with randomized size. n is the number of frames to generate, min is the min
// size is min size of random frame and max is max size of random frames.
func genFrames(n, min, max int) [][]byte {
// Generate randomly sized data for each frame and fill.
rand.Seed(time.Now().UnixNano())
frames := make([][]byte, n)
for i := range frames {
frames[i] = make([]byte, rand.Intn(max-min)+min)
for j := 0; j < len(frames[i]); j++ {
frames[i][j] = byte(j)
}
}
return frames
}
// TestTrimToPTSRange checks that Clip.TrimToPTSRange will correctly return a
// sub Clip of the given PTS range.
func TestTrimToPTSRange(t *testing.T) {
const (
numOfTestFrames = 10
ptsInterval = 4
frameSize = 3
)
clip := &Clip{}
// Generate test frames.
for i := 0; i < numOfTestFrames; i++ {
clip.backing = append(clip.backing, []byte{byte(i), byte(i), byte(i)}...)
clip.frames = append(
clip.frames,
Frame{
Media: clip.backing[i*frameSize : (i+1)*frameSize],
PTS: uint64(i * ptsInterval),
idx: i * frameSize,
},
)
}
// We test each of these scenarios.
tests := []struct {
from uint64
to uint64
expect []byte
err error
}{
{
from: 6,
to: 15,
expect: []byte{
0x01, 0x01, 0x01,
0x02, 0x02, 0x02,
0x03, 0x03, 0x03,
},
err: nil,
},
{
from: 4,
to: 16,
expect: []byte{
0x01, 0x01, 0x01,
0x02, 0x02, 0x02,
0x03, 0x03, 0x03,
},
err: nil,
},
{
from: 10,
to: 5,
expect: nil,
err: errPTSRange,
},
{
from: 50,
to: 70,
expect: nil,
err: errPTSLowerBound,
},
{
from: 5,
to: 70,
expect: nil,
err: errPTSUpperBound,
},
}
// Run tests.
for i, test := range tests {
got, err := clip.TrimToPTSRange(test.from, test.to)
// First check the error.
if err != nil && err != test.err {
t.Errorf("unexpected error: %v for test: %v", err, i)
continue
} else if err != test.err {
t.Errorf("expected to get error: %v for test: %v", test.err, i)
continue
}
// Now check data.
if test.err == nil && !bytes.Equal(test.expect, got.Bytes()) {
t.Errorf("did not get expected data for test: %v\n Got: %v\n, Want: %v\n", i, got, test.expect)
}
}
}
// TestTrimToMetaRange checks that Clip.TrimToMetaRange correctly provides a
// sub Clip for a given meta range.
func TestClipTrimToMetaRange(t *testing.T) {
const (
numOfTestFrames = 10
ptsInterval = 4
frameSize = 3
key = "n"
)
clip := &Clip{}
// Generate test frames.
for i := 0; i < numOfTestFrames; i++ {
clip.backing = append(clip.backing, []byte{byte(i), byte(i), byte(i)}...)
clip.frames = append(
clip.frames,
Frame{
Media: clip.backing[i*frameSize : (i+1)*frameSize],
idx: i * frameSize,
Meta: map[string]string{
key: strconv.Itoa(i),
},
},
)
}
// We test each of these scenarios.
tests := []struct {
from string
to string
expect []byte
err error
}{
{
from: "1",
to: "3",
expect: []byte{
0x01, 0x01, 0x01,
0x02, 0x02, 0x02,
0x03, 0x03, 0x03,
},
err: nil,
},
{
from: "1",
to: "1",
expect: nil,
err: errMetaRange,
},
{
from: "20",
to: "1",
expect: nil,
err: errMetaLowerBound,
},
{
from: "1",
to: "20",
expect: nil,
err: errMetaUpperBound,
},
}
// Run tests.
for i, test := range tests {
got, err := clip.TrimToMetaRange(key, test.from, test.to)
// First check the error.
if err != nil && err != test.err {
t.Errorf("unexpected error: %v for test: %v", err, i)
continue
} else if err != test.err {
t.Errorf("expected to get error: %v for test: %v", test.err, i)
continue
}
// Now check data.
if test.err == nil && !bytes.Equal(test.expect, got.Bytes()) {
t.Errorf("did not get expected data for test: %v\n Got: %v\n, Want: %v\n", i, got, test.expect)
}
}
}
// TestClipSegmentForMeta checks that Clip.SegmentForMeta correctly returns
// segments from a clip with consistent meta defined by a key and value.
func TestClipSegmentForMeta(t *testing.T) {
const (
nFrames = 10 // The number of test frames we want to create.
fSize = 3 // The size of the frame media.
key = "n" // Meta key we will use.
val = "*" // The meta val of interest.
)
tests := []struct {
metaVals []string // These will be the meta vals each frame has.
fIndices []rng // These are the indices of the segments of interest.
}{
{
metaVals: []string{"1", "2", "*", "*", "*", "3", "*", "*", "4", "5"},
fIndices: []rng{{2, 5}, {6, 8}},
},
{
metaVals: []string{"1", "2", "*", "*", "*", "", "*", "*", "4", "5"},
fIndices: []rng{{2, 5}, {6, 8}},
},
{
metaVals: []string{"1", "2", "*", "*", "*", "3", "4", "5", "*", "*"},
fIndices: []rng{{2, 5}, {8, nFrames}},
},
{
metaVals: []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"},
fIndices: nil,
},
}
// Run the tests.
for testn, test := range tests {
clip := &Clip{}
// Generate test frames.
for i := 0; i < nFrames; i++ {
clip.backing = append(clip.backing, []byte{byte(i), byte(i), byte(i)}...)
clip.frames = append(
clip.frames,
Frame{
Media: clip.backing[i*fSize : (i+1)*fSize],
idx: i * fSize,
Meta: map[string]string{
key: test.metaVals[i],
},
},
)
}
// Use function we're testing to get segments.
got := clip.SegmentForMeta(key, val)
// Now get expected segments using indices defined in the test.
var want []Clip
for _, indices := range test.fIndices {
// Calculate the indices for the backing array from the original clip.
backStart := clip.frames[indices.start].idx
backEnd := clip.frames[indices.end-1].idx + len(clip.frames[indices.end-1].Media)
// Use calculated indices to create Clip for current expected segment.
want = append(want, Clip{
frames: clip.frames[indices.start:indices.end],
backing: clip.backing[backStart:backEnd],
})
}
if !reflect.DeepEqual(got, want) {
t.Errorf("did not get expected result for test %v\nGot: %v\nWant: %v\n", testn, got, want)
}
}
}