From 2c8fdda8e30a39f82ddb72f2b8a76129aa9f7a9e Mon Sep 17 00:00:00 2001
From: dexter <178529795@qq.com>
Date: Thu, 7 Apr 2022 20:13:11 +0800
Subject: [PATCH] fix:WriteAnnexB

---
 track/h264.go  |  4 +++-
 track/h265.go  |  4 +++-
 track/video.go | 13 +++++++------
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/track/h264.go b/track/h264.go
index b7c9f97..035d4ae 100644
--- a/track/h264.go
+++ b/track/h264.go
@@ -33,7 +33,9 @@ func NewH264(stream IStream) (vt *H264) {
 	return
 }
 func (vt *H264) WriteAnnexB(pts uint32, dts uint32, frame AnnexBFrame) {
-	vt.Video.WriteAnnexB(pts, dts, frame)
+	for _, slice := range vt.Video.WriteAnnexB(pts, dts, frame) {
+		vt.WriteSlice(slice)
+	}
 	vt.Flush()
 }
 func (vt *H264) WriteSlice(slice NALUSlice) {
diff --git a/track/h265.go b/track/h265.go
index 469c0b7..ceac0b0 100644
--- a/track/h265.go
+++ b/track/h265.go
@@ -31,7 +31,9 @@ func NewH265(stream IStream) (vt *H265) {
 	return
 }
 func (vt *H265) WriteAnnexB(pts uint32, dts uint32, frame AnnexBFrame) {
-	vt.Video.WriteAnnexB(pts, dts, frame)
+	for _, slice := range vt.Video.WriteAnnexB(pts, dts, frame) {
+		vt.WriteSlice(slice)
+	}
 	vt.Flush()
 }
 func (vt *H265) WriteSlice(slice NALUSlice) {
diff --git a/track/video.go b/track/video.go
index 53df72c..2b6e969 100644
--- a/track/video.go
+++ b/track/video.go
@@ -54,26 +54,26 @@ func (t *Video) ComputeGOP() {
 	t.IDRing = t.Ring
 }
 
-func (vt *Video) writeAnnexBSlice(annexb AnnexBFrame) {
+func (vt *Video) writeAnnexBSlice(annexb AnnexBFrame, s *[]NALUSlice) {
 	for len(annexb) > 0 {
 		before, after, found := bytes.Cut(annexb, codec.NALU_Delimiter1)
 		if !found {
-			vt.WriteSlice(NALUSlice{annexb})
+			*s = append(*s, NALUSlice{annexb})
 			return
 		}
 		if len(before) > 0 {
-			vt.WriteSlice(NALUSlice{before})
+			*s = append(*s, NALUSlice{before})
 		}
 		annexb = after
 	}
 }
 
-func (vt *Video) WriteAnnexB(pts uint32, dts uint32, frame AnnexBFrame) {
+func (vt *Video) WriteAnnexB(pts uint32, dts uint32, frame AnnexBFrame) (s []NALUSlice) {
 	// vt.Stream.Tracef("WriteAnnexB:pts %d,dts %d,len %d", pts, dts, len(frame))
 	for len(frame) > 0 {
 		before, after, found := bytes.Cut(frame, codec.NALU_Delimiter2)
 		if !found {
-			vt.writeAnnexBSlice(frame)
+			vt.writeAnnexBSlice(frame, &s)
 			if len(vt.Value.Raw) > 0 {
 				vt.Value.PTS = pts
 				vt.Value.DTS = dts
@@ -81,10 +81,11 @@ func (vt *Video) WriteAnnexB(pts uint32, dts uint32, frame AnnexBFrame) {
 			return
 		}
 		if len(before) > 0 {
-			vt.writeAnnexBSlice(AnnexBFrame(before))
+			vt.writeAnnexBSlice(AnnexBFrame(before), &s)
 		}
 		frame = after
 	}
+	return
 }
 func (vt *Video) WriteAVCC(ts uint32, frame AVCCFrame) {
 	vt.Media.WriteAVCC(ts, frame)