diff --git a/container/mts/encoder_test.go b/container/mts/encoder_test.go index 24fb823d..2d4501d9 100644 --- a/container/mts/encoder_test.go +++ b/container/mts/encoder_test.go @@ -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) + } +} diff --git a/container/mts/meta/meta.go b/container/mts/meta/meta.go index 188c2d4e..0e67aa96 100644 --- a/container/mts/meta/meta.go +++ b/container/mts/meta/meta.go @@ -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:])) { diff --git a/container/mts/meta/meta_test.go b/container/mts/meta/meta_test.go index 38e4dbb6..9316ce55 100644 --- a/container/mts/meta/meta_test.go +++ b/container/mts/meta/meta_test.go @@ -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) { diff --git a/container/mts/metaEncode_test.go b/container/mts/metaEncode_test.go deleted file mode 100644 index 83660777..00000000 --- a/container/mts/metaEncode_test.go +++ /dev/null @@ -1,100 +0,0 @@ -/* -NAME - metaEncode_test.go - -DESCRIPTION - See Readme.md - -AUTHOR - Saxon Nelson-Milton - -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) - } -} diff --git a/container/mts/mpegts.go b/container/mts/mpegts.go index 60b1c484..00028478 100644 --- a/container/mts/mpegts.go +++ b/container/mts/mpegts.go @@ -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 +} diff --git a/container/mts/mpegts_test.go b/container/mts/mpegts_test.go index 1a6409c7..3a603a20 100644 --- a/container/mts/mpegts_test.go +++ b/container/mts/mpegts_test.go @@ -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, + } +} diff --git a/container/mts/payload.go b/container/mts/payload.go new file mode 100644 index 00000000..089eba76 --- /dev/null +++ b/container/mts/payload.go @@ -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 + +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)], + }, + ) +} diff --git a/container/mts/payload_test.go b/container/mts/payload_test.go new file mode 100644 index 00000000..cc8f13fa --- /dev/null +++ b/container/mts/payload_test.go @@ -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 + +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) + } + } +}