From d08dd0b0a1ba73f996c7b74478f570444edd1c8b Mon Sep 17 00:00:00 2001
From: nareix <hi.nareix@gmail.com>
Date: Fri, 1 Jul 2016 18:37:16 +0800
Subject: [PATCH] add h264parser.CheckNALUsType for demuxer and muxer

---
 flv.go | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/flv.go b/flv.go
index ab14d11..bb03bd7 100644
--- a/flv.go
+++ b/flv.go
@@ -77,10 +77,16 @@ func (self *Muxer) WritePacket(pkt av.Packet) (err error) {
 
 	switch stream.Type() {
 	case av.H264:
+		if typ := h264parser.CheckNALUsType(pkt.Data); typ != h264parser.NALU_RAW {
+			err = fmt.Errorf("flv: h264 nalu format=%d invalid", typ)
+			return
+		}
+		var b [4]byte
+		pio.PutU32BE(b[:], uint32(len(pkt.Data)))
 		tag := &flvio.Videodata{
 			AVCPacketType: flvio.AVC_NALU,
 			CodecID: flvio.VIDEO_H264,
-			Data: pkt.Data,
+			Datav: [][]byte{b[:], pkt.Data},
 			CompositionTime: timeToTs(pkt.CompositionTime),
 		}
 		if pkt.IsKeyFrame {
@@ -245,7 +251,11 @@ func (self *Demuxer) ReadPacket() (pkt av.Packet, err error) {
 				pkt.Idx = int8(self.videostreamidx)
 				pkt.CompositionTime = tsToTime(tag.CompositionTime)
 				pkt.IsKeyFrame = tag.FrameType == flvio.FRAME_KEY
-				pkt.Data = tag.Data
+				if typ := h264parser.CheckNALUsType(tag.Data); typ != h264parser.NALU_AVCC {
+					err = fmt.Errorf("flv: h264 nalu format=%d invalid", typ)
+					return
+				}
+				pkt.Data = tag.Data[4:]
 				break loop
 			}