From e4bce3bcb33d2bed529f40937badc769da396a92 Mon Sep 17 00:00:00 2001
From: Saxon <saxon@ausocean.org>
Date: Wed, 12 Jun 2019 01:01:51 +0930
Subject: [PATCH] container/mts: added ExtractMeta function and testing

Added function to extract meta from first PMT found in a mpegts clip. This simply
wraps the logic we've used in vidrecord, i.e. Finding the PMT, converting to a
comcast gots psi.PSIBytesm, using HasDescriptor to get the meta descriptor, and
then mapping metadata. Also added testing.
---
 container/mts/metaEncode_test.go | 26 ++++++++++++++++++++++++++
 container/mts/mpegts.go          | 25 ++++++++++++++++++++++++-
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/container/mts/metaEncode_test.go b/container/mts/metaEncode_test.go
index 83660777..3a0dbd62 100644
--- a/container/mts/metaEncode_test.go
+++ b/container/mts/metaEncode_test.go
@@ -29,6 +29,7 @@ package mts
 
 import (
 	"bytes"
+	"reflect"
 	"testing"
 
 	"bitbucket.org/ausocean/av/container/mts/meta"
@@ -98,3 +99,28 @@ func TestMetaEncode2(t *testing.T) {
 		t.Errorf(errNotExpectedOut, 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(errUnexpectedErr, err.Error())
+	}
+	out := buf.Bytes()
+	got, err := ExtractMeta(out)
+	if err != nil {
+		t.Errorf(errUnexpectedErr, err.Error())
+	}
+	want := map[string]string{
+		"ts":  "12345678",
+		"loc": "1234,4321,1234",
+	}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf(errNotExpectedOut, got, want)
+	}
+}
diff --git a/container/mts/mpegts.go b/container/mts/mpegts.go
index eb4bee5d..0aabd5a1 100644
--- a/container/mts/mpegts.go
+++ b/container/mts/mpegts.go
@@ -33,6 +33,8 @@ import (
 	"errors"
 	"fmt"
 
+	"bitbucket.org/ausocean/av/container/mts/meta"
+	"bitbucket.org/ausocean/av/container/mts/psi"
 	"github.com/Comcast/gots/packet"
 	"github.com/Comcast/gots/pes"
 )
@@ -331,7 +333,6 @@ func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) {
 	copy(_pkt[:], pkt)
 	payload, err := packet.Payload(&_pkt)
 	if err != nil {
-		fmt.Printf("_pkt: %v\n", _pkt)
 		return [2]uint64{}, err
 	}
 
@@ -361,3 +362,25 @@ 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")
 }
+
+// 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, errors.New("PMT does not have descriptor")
+	}
+	// Skip the descriptor head.
+	m := metaDescriptor[2:]
+
+	return meta.GetAllAsMap(m)
+}