From e2fb4a644f80292b5ec96dd54b51548c660ab9de Mon Sep 17 00:00:00 2001
From: nareix <hi.nareix@gmail.com>
Date: Sun, 17 Jul 2016 15:35:08 +0800
Subject: [PATCH] ts: fix FillPSI BUG and ParsePSI datalen bug

---
 format/ts/muxer.go     |  2 ++
 format/ts/tsio/tsio.go | 12 ++++++++----
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/format/ts/muxer.go b/format/ts/muxer.go
index 6fb64b5..7b880b4 100644
--- a/format/ts/muxer.go
+++ b/format/ts/muxer.go
@@ -102,6 +102,7 @@ func (self *Muxer) WritePATPMT() (err error) {
 		},
 	}
 	patlen := pat.Marshal(self.psidata[tsio.PSIHeaderLength:])
+	tsio.FillPSI(self.psidata, tsio.TableIdPAT, tsio.TableExtPAT, patlen)
 	self.datav[0] = self.psidata[:tsio.PSIHeaderLength+patlen]
 	if err = self.tswpat.WritePackets(self.w, self.datav[:1], 0, false, true); err != nil {
 		return
@@ -133,6 +134,7 @@ func (self *Muxer) WritePATPMT() (err error) {
 		return
 	}
 	pmt.Marshal(self.psidata[tsio.PSIHeaderLength:])
+	tsio.FillPSI(self.psidata, tsio.TableIdPMT, tsio.TableExtPMT, pmtlen)
 	self.datav[0] = self.psidata[:tsio.PSIHeaderLength+pmtlen]
 	if err = self.tswpmt.WritePackets(self.w, self.datav[:1], 0, false, true); err != nil {
 		return
diff --git a/format/ts/tsio/tsio.go b/format/ts/tsio/tsio.go
index fb2ba6a..9a7c37b 100644
--- a/format/ts/tsio/tsio.go
+++ b/format/ts/tsio/tsio.go
@@ -306,6 +306,11 @@ func ParsePSI(h []byte) (tableid uint8, tableext uint16, hdrlen int, datalen int
 	datalen = int(pio.U16BE(h[hdrlen:]))&0x3ff - 9
 	hdrlen += 2
 
+	if datalen < 0 {
+		err = ErrPSIHeader
+		return
+	}
+
 	// Table ID extension(16)
 	tableext = pio.U16BE(h[hdrlen:])
 	hdrlen += 2
@@ -330,7 +335,7 @@ func ParsePSI(h []byte) (tableid uint8, tableext uint16, hdrlen int, datalen int
 
 const PSIHeaderLength = 9
 
-func FillPSI(h []byte, tableid uint8, tableext uint16, data []byte) (n int) {
+func FillPSI(h []byte, tableid uint8, tableext uint16, datalen int) (n int) {
 	// pointer(8)
 	h[n] = 0
 	n++
@@ -340,7 +345,7 @@ func FillPSI(h []byte, tableid uint8, tableext uint16, data []byte) (n int) {
 	n++
 
 	// section_syntax_indicator(1)=1,private_bit(1)=0,reserved(2)=3,unused(2)=0,section_length(10)
-	pio.PutU16BE(h[n:], uint16(0xa<<12 | 2+3+4+len(data)))
+	pio.PutU16BE(h[n:], uint16(0xa<<12 | 2+3+4+datalen))
 	n += 2
 
 	// Table ID extension(16)
@@ -359,8 +364,7 @@ func FillPSI(h []byte, tableid uint8, tableext uint16, data []byte) (n int) {
 	h[n] = 0
 	n++
 
-	copy(h[n:], data)
-	n += len(data)
+	n += datalen
 
 	crc := calcCRC32(0xffffffff, h[1:n])
 	pio.PutU32LE(h[n:], crc)