From aa671349692bfdbdfdcb8aafed9b69cbcca4dfbf Mon Sep 17 00:00:00 2001
From: scruzin <anoble@gmail.com>
Date: Fri, 12 Jul 2019 14:01:49 +0930
Subject: [PATCH] Added GetPTS() function.

---
 container/mts/mpegts.go | 89 +++++++++++++++++++++++++----------------
 1 file changed, 55 insertions(+), 34 deletions(-)

diff --git a/container/mts/mpegts.go b/container/mts/mpegts.go
index 28ad802b..ee898177 100644
--- a/container/mts/mpegts.go
+++ b/container/mts/mpegts.go
@@ -33,7 +33,6 @@ import (
 	"fmt"
 
 	"github.com/Comcast/gots/packet"
-	"github.com/Comcast/gots/pes"
 	"github.com/pkg/errors"
 
 	"bitbucket.org/ausocean/av/container/mts/meta"
@@ -379,8 +378,8 @@ var errNoPTS = errors.New("could not find PTS")
 // GetPTSRange retreives the first and last PTS of an MPEGTS clip.
 // If there is only one PTS, it is included twice in the pts return value.
 func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) {
-	var _pkt packet.Packet
-	// Find the first packet with PID pidType and PUSI.
+	var _pts int64
+	// Get the first PTS for the given PID.
 	var i int
 	for {
 		if i >= len(clip) {
@@ -390,25 +389,14 @@ func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) {
 		if err != nil {
 			return pts, errors.Wrap(err, fmt.Sprintf("could not find packet of PID: %d", pid))
 		}
-		copy(_pkt[:], pkt)
-		if _pkt.PayloadUnitStartIndicator() {
+		_pts, err = GetPTS(pkt)
+		if err == nil {
 			break
 		}
 		i += _i + PacketSize
 	}
 
-	// Get the payload of the packet, which will be the start of the PES packet.
-	payload, err := packet.Payload(&_pkt)
-	if err != nil {
-		return pts, err
-	}
-
-	// Get the the first PTS from the PES header.
-	_pes, err := pes.NewPESHeader(payload)
-	if err != nil {
-		return pts, err
-	}
-	pts[0] = _pes.PTS()
+	pts[0] = uint64(_pts)
 	pts[1] = pts[0] // Until we have find a second PTS.
 
 	// Get the last PTS searching in reverse from end of the clip.
@@ -419,31 +407,64 @@ func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) {
 		if err != nil || i <= first {
 			return pts, nil
 		}
-
-		copy(_pkt[:], pkt)
-		if packet.PayloadUnitStartIndicator(&_pkt) {
+		_pts, err = GetPTS(pkt)
+		if err == nil {
 			break
 		}
-
 		i = _i
 	}
 
-	// Get the payload of the packet.
-	payload, err = packet.Payload(&_pkt)
-	if err != nil {
-		return
-	}
-
-	// Get the the last PTS from the PES header.
-	_pes, err = pes.NewPESHeader(payload)
-	if err != nil {
-		return
-	}
-	pts[1] = _pes.PTS()
+	pts[1] = uint64(_pts)
 
 	return
 }
 
+var errNoPesPayload = errors.New("no PES payload")
+var errNoPesPTS = errors.New("no PES PTS")
+var errInvalidPesHeader = errors.New("invalid PES header")
+var errInvalidPesPayload = errors.New("invalid PES payload")
+
+// GetPTS returns a PTS from an MTS packet that has PES payload, or an error otherwise.
+func GetPTS(pkt []byte) (pts int64, err error) {
+	// Bail if packet does not contain a PES.
+	if pkt[1]&0x040 == 0 {
+		err = errNoPesPayload
+		return
+	}
+
+	// Compute start of PES payload and check its length.
+	start := HeadSize
+	if pkt[3]&0x20 != 0 {
+		// Take into account length of the adaptation field.
+		start += 1 + int(pkt[4])
+	}
+	pes := pkt[start:]
+
+	if len(pes) < 9 {
+		err = errInvalidPesHeader
+		return
+	}
+
+	// Check that PES has a PTS.
+	if pes[7]&0xc0 == 0 {
+		err = errNoPesPTS
+		return
+	}
+
+	// Extract the PTS.
+	if len(pes) < 14 {
+		err = errInvalidPesHeader
+		return
+	}
+	pts = extractPTS(pes[9:14])
+	return
+}
+
+// extractTime extracts a PTS from the given data.
+func extractPTS(d []byte) int64 {
+	return (int64((d[0]>>1)&0x07) << 30) | (int64(d[1]) << 22) | (int64((d[2]>>1)&0x7f) << 15) | (int64(d[3]) << 7) | int64((d[4]>>1)&0x7f)
+}
+
 var errNoMeta = errors.New("PMT does not contain meta")
 
 // ExtractMeta returns a map of metadata from the first PMT's metaData
@@ -460,7 +481,7 @@ func ExtractMeta(d []byte) (map[string]string, error) {
 // metaFromPMT returns metadata, if any, from a PMT.
 func metaFromPMT(d []byte) (m map[string]string, err error) {
 	// Get as PSI type, skipping the MTS header.
-	pmt := psi.PSIBytes(d[4:])
+	pmt := psi.PSIBytes(d[HeadSize:])
 
 	// Get the metadata descriptor.
 	_, desc := pmt.HasDescriptor(psi.MetadataTag)