暂时解决nalu分帧问题

Author: langhuihui

nalu.go

@@ -1,11 +1,9 @@
 package engine
 
 import (
-	"bytes"
 	"encoding/binary"
+
 	. "github.com/Monibuca/engine/v2/avformat"
-	"github.com/Monibuca/engine/v2/pool"
-	"github.com/Monibuca/engine/v2/util"
 )
 
 const (
@@ -15,22 +13,35 @@ const (
 
 	naluTypeBitmask   = 0x1F
 	naluRefIdcBitmask = 0x60
-	fuaStartBitmask   = 0x80
-	fuaEndBitmask     = 0x40
+	fuaStartBitmask   = 0x80 //1000 0000
+	fuaEndBitmask     = 0x40 //0100 0000
 )
 
 type NALU struct {
 	Publisher
 	fuBuffer []byte //用于fua的解析暂存的缓存
+	lastTs   uint32
+	buffer   []byte //用于存储分帧的Nalu数据
 }
 
+func (r *NALU) writePicture(ts uint32, head, payload []byte) {
+	if r.VideoTag == nil {
+		return
+	}
+	if payload[1]&fuaStartBitmask != 0 {
+		if len(r.buffer) > 0 {
+			r.PushVideo(r.lastTs, r.buffer)
+			r.buffer = nil
+		}
+		r.buffer = append(r.buffer, head...)
+		r.lastTs = ts
+	}
+	nl := len(payload)
+	r.buffer = append(r.buffer, byte(nl>>24), byte(nl>>16), byte(nl>>8), byte(nl))
+	r.buffer = append(r.buffer, payload...)
+}
 func (r *NALU) WriteNALU(ts uint32, payload []byte) {
 	nalType := payload[0] & naluTypeBitmask
-	buffer := r.AVRing.Buffer
-	if buffer == nil {
-		buffer = bytes.NewBuffer([]byte{})
-		r.AVRing.Buffer = buffer
-	}
 	switch nalType {
 	case NALU_STAPA:
 		for currOffset, naluSize := stapaHeaderSize, 0; currOffset < len(payload); currOffset += naluSize {
@@ -63,32 +74,11 @@ func (r *NALU) WriteNALU(ts uint32, payload []byte) {
 	case NALU_PPS:
 		r.WritePPS(payload)
 	case NALU_Access_Unit_Delimiter:
-		r.PushVideo(ts, r.AVRing.Bytes())
-		r.GetBuffer()
+
 	case NALU_IDR_Picture:
-		if r.VideoTag == nil {
-			break
-		}
-		if buffer.Len() == 0 {
-			buffer.Write(RTMP_KEYFRAME_HEAD)
-		}
-		nl := pool.GetSlice(4)
-		util.BigEndian.PutUint32(nl, uint32(len(payload)))
-		buffer.Write(nl)
-		pool.RecycleSlice(nl)
-		buffer.Write(payload)
+		r.writePicture(ts, RTMP_KEYFRAME_HEAD, payload)
 	case NALU_Non_IDR_Picture:
-		if r.VideoTag == nil {
-			break
-		}
-		if buffer.Len() == 0 {
-			buffer.Write(RTMP_NORMALFRAME_HEAD)
-		}
-		nl := pool.GetSlice(4)
-		util.BigEndian.PutUint32(nl, uint32(len(payload)))
-		buffer.Write(nl)
-		pool.RecycleSlice(nl)
-		buffer.Write(payload)
+		r.writePicture(ts, RTMP_NORMALFRAME_HEAD, payload)
 	default:
 		Printf("nalType not support yet:%d", nalType)
 	}