From 3eab25e18be49352cb77d5852aade9aefeab5d78 Mon Sep 17 00:00:00 2001
From: Saxon <saxon@ausocean.org>
Date: Mon, 17 Jun 2019 04:30:52 +0930
Subject: [PATCH] container/mts: added SegmentForMeta func and testing

This func will further segment an MTS segment to each series of packets that
correspond to a given key and val of meta. Testing for this func has also been
written.
---
 container/mts/meta/meta.go   |  6 +--
 container/mts/mpegts.go      | 41 ++++++++++++++++++
 container/mts/mpegts_test.go | 83 ++++++++++++++++++++++++++++++++++++
 container/mts/payload.go     |  2 +-
 4 files changed, 128 insertions(+), 4 deletions(-)

diff --git a/container/mts/meta/meta.go b/container/mts/meta/meta.go
index 2baa33b7..b45c5796 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,7 +209,7 @@ 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)
 	}
@@ -237,7 +237,7 @@ func GetAllAsMap(d []byte) (map[string]string, error) {
 		// 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
+			return nil, ErrUnexpectedMetaFormat
 		}
 		all[kv[0]] = kv[1]
 	}
diff --git a/container/mts/mpegts.go b/container/mts/mpegts.go
index 4e6cbd51..622a19aa 100644
--- a/container/mts/mpegts.go
+++ b/container/mts/mpegts.go
@@ -442,3 +442,44 @@ func TrimToMetaRange(d []byte, key, from, to string) ([]byte, error) {
 		}
 	}
 }
+
+func SegmentForMeta(d []byte, key, val string) ([][]byte, error) {
+	var (
+		pkt        packet.Packet
+		segmenting bool
+		res        [][]byte
+		start      int
+	)
+
+	for i := 0; i < len(d); i += PacketSize {
+		copy(pkt[:], d[i:i+PacketSize])
+		if pkt.PID() == PmtPid {
+			_meta, err := ExtractMeta(pkt[:])
+			switch err {
+			case errNoMeta, meta.ErrUnexpectedMetaFormat:
+				if segmenting {
+					res = append(res, d[start:i])
+					segmenting = false
+				}
+				continue
+			case nil: // do nothing.
+			default:
+				return nil, err
+			}
+
+			if _meta[key] == val && !segmenting {
+				start = i
+				segmenting = true
+			} else if _meta[key] != val && segmenting {
+				res = append(res, d[start:i])
+				segmenting = false
+			}
+		}
+	}
+
+	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 5ac633f7..5d74c7ce 100644
--- a/container/mts/mpegts_test.go
+++ b/container/mts/mpegts_test.go
@@ -30,6 +30,7 @@ package mts
 import (
 	"bytes"
 	"math/rand"
+	"reflect"
 	"strconv"
 	"testing"
 	"time"
@@ -334,3 +335,85 @@ func TestTrimToMetaRange(t *testing.T) {
 		}
 	}
 }
+
+func TestSegmentForMeta(t *testing.T) {
+	Meta = meta.New()
+
+	const (
+		nPSI = 10
+		key  = "n"
+		val  = "*"
+	)
+
+	tests := []struct {
+		metaVals   [nPSI]string
+		expectIdxs []rng
+	}{
+		{
+			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(2, 5),
+			},
+		},
+		{
+			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 {
+		clip.Reset()
+		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)
+			}
+		}
+		var want [][]byte
+		for _, idxs := range test.expectIdxs {
+			want = append(want, clip.Bytes()[idxs.start:idxs.end])
+		}
+		got, err := SegmentForMeta(clip.Bytes(), key, val)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+
+		if !reflect.DeepEqual(want, got) {
+			t.Errorf("did not get expected result for test %v\nGot: %v\nWant: %v\n", testn, got, want)
+		}
+	}
+}
+
+type rng struct {
+	start int
+	end   int
+}
+
+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
index 37394ce6..b7ccba4e 100644
--- a/container/mts/payload.go
+++ b/container/mts/payload.go
@@ -8,7 +8,7 @@ import (
 	"github.com/Comcast/gots/pes"
 )
 
-// TODO: write function for getting subslice of mpegts for meta interval
+// TODO(saxon): write function to get sub clips that posess a particular meta value.
 
 // 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.