From 848a9b2a397eab871474b45ae7fe98a3e3010e29 Mon Sep 17 00:00:00 2001
From: nareix <hi.nareix@gmail.com>
Date: Sun, 6 Dec 2015 19:50:46 +0800
Subject: [PATCH] writePMT() bugfix

---
 example/test.go | 26 ++++++++++++++++++++++++--
 reader.go       | 41 +++++++++++++++++++++++------------------
 writer.go       | 28 ++++++++++++++--------------
 3 files changed, 61 insertions(+), 34 deletions(-)

diff --git a/example/test.go b/example/test.go
index 6ab86a1..f305a34 100644
--- a/example/test.go
+++ b/example/test.go
@@ -159,9 +159,9 @@ func main() {
 			W: file,
 			PID: 0,
 		}
-		pat := ts.PAT {
+		pat := ts.PAT{
 			Entries: []ts.PATEntry{
-				{ProgramNumber: 1, ProgramMapPID: 4096},
+				{ProgramNumber: 1, ProgramMapPID: 0x1000},
 			},
 		}
 		bw := &bytes.Buffer{}
@@ -174,8 +174,30 @@ func main() {
 		return
 	}
 
+	writePMT := func() (err error) {
+		w := &ts.TSWriter{
+			W: file,
+			PID: 0x1000,
+		}
+		pmt := ts.PMT{
+			PCRPID: 0x100,
+			ElementaryStreamInfos: []ts.ElementaryStreamInfo{
+				{StreamType: ts.ElementaryStreamTypeH264, ElementaryPID: 0x100},
+			},
+		}
+		bw := &bytes.Buffer{}
+		if err = ts.WritePMT(bw, pmt); err != nil {
+			return
+		}
+		if err = w.Write(bw.Bytes(), false); err != nil {
+			return
+		}
+		return
+	}
+
 	if file != nil {
 		writePAT()
+		writePMT()
 		file.Close()
 	}
 
diff --git a/reader.go b/reader.go
index dc49f9a..23972de 100644
--- a/reader.go
+++ b/reader.go
@@ -275,6 +275,25 @@ func ReadPSI(r io.Reader) (self PSI, lr *io.LimitedReader, cr *Crc32Reader, err
 }
 
 func ReadPMT(r io.Reader) (self PMT, err error) {
+	readDescs := func(lr *io.LimitedReader) (res []Descriptor, err error) {
+		var desc Descriptor
+		for lr.N > 0 {
+			if desc.Tag, err = ReadUInt(lr, 1); err != nil {
+				return
+			}
+			var length uint
+			if length, err = ReadUInt(lr, 1); err != nil {
+				return
+			}
+			desc.Data = make([]byte, length)
+			if _, err = lr.Read(desc.Data); err != nil {
+				return
+			}
+			res = append(res, desc)
+		}
+		return
+	}
+
 	var lr *io.LimitedReader
 	var cr *Crc32Reader
 	//var psi PSI
@@ -306,30 +325,16 @@ func ReadPMT(r io.Reader) (self PMT, err error) {
 		})
 	}
 
-	// Reserved(6)
+	// Reserved(4)=0xf
+	// Reserved(2)=0x0
 	// Program info length(10)
 	if flags, err = ReadUInt(lr, 2); err != nil {
 		return
 	}
 	length = flags & 0x3ff
 
-	readDescs := func(lr *io.LimitedReader) (res []Descriptor, err error) {
-		var desc Descriptor
-		for lr.N > 0 {
-			if desc.Tag, err = ReadUInt(lr, 1); err != nil {
-				return
-			}
-			var length uint
-			if length, err = ReadUInt(lr, 1); err != nil {
-				return
-			}
-			desc.Data = make([]byte, length)
-			if _, err = lr.Read(desc.Data); err != nil {
-				return
-			}
-			res = append(res, desc)
-		}
-		return
+	if DebugReader {
+		fmt.Printf("pmt: ProgramDescriptorsLen=%d\n", length)
 	}
 
 	if length > 0 {
diff --git a/writer.go b/writer.go
index 73b62ea..d3945b5 100644
--- a/writer.go
+++ b/writer.go
@@ -367,7 +367,19 @@ func WritePMT(w io.Writer, self PMT) (err error) {
 	}
 
 	writeBody := func(w io.Writer) (err error) {
-		if err = writeDescs(w, self.ProgramDescriptors); err != nil {
+		if err = WriteUInt(w, self.PCRPID|7<<13, 2); err != nil {
+			return
+		}
+
+		bw := &bytes.Buffer{}
+		if err = writeDescs(bw, self.ProgramDescriptors); err != nil {
+			return
+		}
+
+		if err = WriteUInt(w, 0xf<<12|uint(bw.Len()), 2); err != nil {
+			return
+		}
+		if _, err = w.Write(bw.Bytes()); err != nil {
 			return
 		}
 
@@ -401,20 +413,8 @@ func WritePMT(w io.Writer, self PMT) (err error) {
 		return
 	}
 
-	writeAll := func(w io.Writer) (err error) {
-		if err = WriteUInt(w, self.PCRPID|7<<13, 2); err != nil {
-			return
-		}
-
-		bw := &bytes.Buffer{}
-		if err = writeBody(bw); err != nil {
-			return
-		}
-		return
-	}
-
 	bw := &bytes.Buffer{}
-	if err = writeAll(bw); err != nil {
+	if err = writeBody(bw); err != nil {
 		return
 	}