From ab8f020f63658d624da2e3e5bfb9d6578d537235 Mon Sep 17 00:00:00 2001
From: nareix <hi.nareix@gmail.com>
Date: Thu, 14 Jul 2016 23:31:27 +0800
Subject: [PATCH] ts: add packet size check

---
 format/ts/demuxer.go | 17 ++++++++++-------
 format/ts/stream.go  |  1 +
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/format/ts/demuxer.go b/format/ts/demuxer.go
index df26b86..866fc18 100644
--- a/format/ts/demuxer.go
+++ b/format/ts/demuxer.go
@@ -2,6 +2,7 @@ package ts
 
 import (
 	"bufio"
+	"fmt"
 	"time"
 	"github.com/nareix/pio"
 	"github.com/nareix/joy4/av"
@@ -207,6 +208,11 @@ func (self *Stream) payloadEnd() (n int, err error) {
 	if payload == nil {
 		return
 	}
+	if self.datalen != 0 && len(payload) != self.datalen {
+		err = fmt.Errorf("ts: packet size mismatch size=%d correct=%d", len(payload), self.datalen)
+		return
+	}
+	self.data = nil
 
 	switch self.streamType {
 	case ElementaryStreamTypeAdtsAAC:
@@ -258,8 +264,6 @@ func (self *Stream) payloadEnd() (n int, err error) {
 		}
 	}
 
-	self.data = nil
-
 	return
 }
 
@@ -269,15 +273,14 @@ func (self *Stream) handleTSPacket(start bool, iskeyframe bool, payload []byte)
 			return
 		}
 		var hdrlen int
-		var datalen int
-		if hdrlen, _, datalen, self.pts, self.dts, err = tsio.ParsePESHeader(payload); err != nil {
+		if hdrlen, _, self.datalen, self.pts, self.dts, err = tsio.ParsePESHeader(payload); err != nil {
 			return
 		}
 		self.iskeyframe = iskeyframe
-		if datalen == 0 {
-			self.data = make([]byte, 0, 16000)
+		if self.datalen == 0 {
+			self.data = make([]byte, 0, 4096)
 		} else {
-			self.data = make([]byte, 0, datalen)
+			self.data = make([]byte, 0, self.datalen)
 		}
 		self.data = append(self.data, payload[hdrlen:]...)
 	} else {
diff --git a/format/ts/stream.go b/format/ts/stream.go
index f301fb1..e649481 100644
--- a/format/ts/stream.go
+++ b/format/ts/stream.go
@@ -24,6 +24,7 @@ type Stream struct {
 	iskeyframe bool
 	pts, dts time.Duration
 	data []byte
+	datalen int
 }
 
 func timeToPesTs(tm time.Duration) uint64 {