diff --git a/container/mts/mpegts.go b/container/mts/mpegts.go index 0aabd5a1..4e6cbd51 100644 --- a/container/mts/mpegts.go +++ b/container/mts/mpegts.go @@ -31,7 +31,6 @@ package mts import ( "errors" - "fmt" "bitbucket.org/ausocean/av/container/mts/meta" "bitbucket.org/ausocean/av/container/mts/psi" @@ -175,11 +174,16 @@ func FindPat(d []byte) ([]byte, int, error) { return FindPid(d, PatPid) } +var ( + errInvalidLen = errors.New("MPEG-TS data not of valid length") + errCouldNotFind = errors.New("could not find packet with given PID") +) + // FindPid will take a clip of MPEG-TS and try to find a packet with given PID - if one // is found, then it is returned along with its index, otherwise nil, -1 and an error is returned. func FindPid(d []byte, pid uint16) (pkt []byte, i int, err error) { if len(d) < PacketSize { - return nil, -1, 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]) @@ -188,7 +192,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 @@ -363,6 +367,8 @@ func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) { return [2]uint64{}, errors.New("could only find one access unit in mpegts clip") } +var errNoMeta = errors.New("PMT does not contain meta") + // ExtractMeta returns a map of metadata from the first PMT's metaData // descriptor, that is found in the MPEG-TS clip d. d must contain a series of // complete MPEG-TS packets. @@ -377,10 +383,62 @@ func ExtractMeta(d []byte) (map[string]string, error) { _, metaDescriptor := pmt.HasDescriptor(psi.MetadataTag) if metaDescriptor == nil { - return nil, errors.New("PMT does not have descriptor") + return nil, errNoMeta } // Skip the descriptor head. m := metaDescriptor[2:] return meta.GetAllAsMap(m) } + +// TrimToMetaRange trims a slice of MPEG-TS to a segment between two points of +// meta data described by key, from and to. +func TrimToMetaRange(d []byte, key, from, to string) ([]byte, error) { + if len(d)%PacketSize != 0 { + return nil, errors.New("MTS clip is not of valid size") + } + + if from == to { + return nil, errors.New("'from' and 'to' cannot be equivalent") + } + + var ( + start = -1 // Index of the start of the segment in d. + end = -1 // Index of the end of segment in d. + off int // Index of remaining slice of d to check after each PMT found. + ) + + for { + // Find the next PMT. + pmt, idx, err := FindPid(d[off:], PmtPid) + if err != nil { + switch -1 { + case start: + return nil, errMetaLowerBound + case end: + return nil, errMetaUpperBound + default: + panic("should not have got error from FindPid") + } + } + off += idx + PacketSize + + meta, err := ExtractMeta(pmt) + switch err { + case nil: // do nothing + case errNoMeta: + continue + default: + return nil, err + } + + if start == -1 { + if meta[key] == from { + start = off - PacketSize + } + } else if meta[key] == to { + end = off + return d[start:end], nil + } + } +} diff --git a/container/mts/mpegts_test.go b/container/mts/mpegts_test.go index 4c90cc0e..5ac633f7 100644 --- a/container/mts/mpegts_test.go +++ b/container/mts/mpegts_test.go @@ -30,9 +30,11 @@ package mts import ( "bytes" "math/rand" + "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" @@ -264,3 +266,71 @@ 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) + } + } +} diff --git a/container/mts/payload_test.go b/container/mts/payload_test.go index b39e3ebe..4ca12f42 100644 --- a/container/mts/payload_test.go +++ b/container/mts/payload_test.go @@ -318,7 +318,7 @@ func TestTrimToPTSRange(t *testing.T) { // TestTrimToMetaRange checks that Clip.TrimToMetaRange correctly provides a // sub Clip for a given meta range. -func TestTrimToMetaRange(t *testing.T) { +func TestClipTrimToMetaRange(t *testing.T) { const ( numOfTestFrames = 10 ptsInterval = 4