container/mts: wrote function TrimToMtsRange and added related testing

This function will return a sub slice of MPEG-TS corresponding to an interval of metadata.
Also wrote testing for this function.

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
+		}
+	}
+}

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)
+		}
+	}
+}

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