mirror of https://bitbucket.org/ausocean/av.git
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
commit
aeb10e0ab7
|
|
@ -29,12 +29,14 @@ import (
|
||||||
"bytes"
|
"bytes"
|
||||||
"io"
|
"io"
|
||||||
"io/ioutil"
|
"io/ioutil"
|
||||||
|
"reflect"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
"github.com/Comcast/gots/packet"
|
"github.com/Comcast/gots/packet"
|
||||||
"github.com/Comcast/gots/pes"
|
"github.com/Comcast/gots/pes"
|
||||||
|
|
||||||
"bitbucket.org/ausocean/av/container/mts/meta"
|
"bitbucket.org/ausocean/av/container/mts/meta"
|
||||||
|
"bitbucket.org/ausocean/av/container/mts/psi"
|
||||||
)
|
)
|
||||||
|
|
||||||
type nopCloser struct{ io.Writer }
|
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")
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -51,7 +51,7 @@ const (
|
||||||
var (
|
var (
|
||||||
errKeyAbsent = errors.New("Key does not exist in map")
|
errKeyAbsent = errors.New("Key does not exist in map")
|
||||||
errInvalidMeta = errors.New("Invalid metadata given")
|
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
|
// 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 {
|
for i, entry := range entries {
|
||||||
kv := strings.Split(entry, "=")
|
kv := strings.Split(entry, "=")
|
||||||
if len(kv) != 2 {
|
if len(kv) != 2 {
|
||||||
return nil, errUnexpectedMetaFormat
|
return nil, ErrUnexpectedMetaFormat
|
||||||
}
|
}
|
||||||
copy(all[i][:], kv)
|
copy(all[i][:], kv)
|
||||||
}
|
}
|
||||||
return all, nil
|
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.
|
// checkHeader checks that a valid metadata header exists in the given data.
|
||||||
func checkMeta(d []byte) error {
|
func checkMeta(d []byte) error {
|
||||||
if len(d) == 0 || d[0] != 0 || binary.BigEndian.Uint16(d[2:headSize]) != uint16(len(d[headSize:])) {
|
if len(d) == 0 || d[0] != 0 || binary.BigEndian.Uint16(d[2:headSize]) != uint16(len(d[headSize:])) {
|
||||||
|
|
|
||||||
|
|
@ -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
|
// TestKeys checks that we can successfully get keys from some metadata using
|
||||||
// the meta.Keys method.
|
// the meta.Keys method.
|
||||||
func TestKeys(t *testing.T) {
|
func TestKeys(t *testing.T) {
|
||||||
|
|
|
||||||
|
|
@ -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)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
@ -35,6 +35,9 @@ import (
|
||||||
"github.com/Comcast/gots/packet"
|
"github.com/Comcast/gots/packet"
|
||||||
"github.com/Comcast/gots/pes"
|
"github.com/Comcast/gots/pes"
|
||||||
"github.com/pkg/errors"
|
"github.com/pkg/errors"
|
||||||
|
|
||||||
|
"bitbucket.org/ausocean/av/container/mts/meta"
|
||||||
|
"bitbucket.org/ausocean/av/container/mts/psi"
|
||||||
)
|
)
|
||||||
|
|
||||||
const PacketSize = 188
|
const PacketSize = 188
|
||||||
|
|
@ -173,11 +176,17 @@ func FindPat(d []byte) ([]byte, int, error) {
|
||||||
return FindPid(d, PatPid)
|
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
|
// 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.
|
// 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) {
|
func FindPid(d []byte, pid uint16) (pkt []byte, i int, err error) {
|
||||||
if len(d) < PacketSize {
|
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 {
|
for i = 0; i < len(d); i += PacketSize {
|
||||||
p := (uint16(d[i+1]&0x1f) << 8) | uint16(d[i+2])
|
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
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
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
|
// 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.
|
// GetPTSRange retreives the first and last PTS of an MPEGTS clip.
|
||||||
func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) {
|
func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) {
|
||||||
var _pkt packet.Packet
|
var _pkt packet.Packet
|
||||||
|
|
@ -370,4 +382,130 @@ func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) {
|
||||||
return
|
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
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -30,9 +30,12 @@ package mts
|
||||||
import (
|
import (
|
||||||
"bytes"
|
"bytes"
|
||||||
"math/rand"
|
"math/rand"
|
||||||
|
"reflect"
|
||||||
|
"strconv"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"bitbucket.org/ausocean/av/container/mts/meta"
|
||||||
"bitbucket.org/ausocean/av/container/mts/pes"
|
"bitbucket.org/ausocean/av/container/mts/pes"
|
||||||
"bitbucket.org/ausocean/av/container/mts/psi"
|
"bitbucket.org/ausocean/av/container/mts/psi"
|
||||||
"github.com/Comcast/gots/packet"
|
"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)
|
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,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -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)],
|
||||||
|
},
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue