fix: record mp4

plugin.go

@@ -442,7 +442,7 @@ func (p *Plugin) SubscribeWithConfig(ctx context.Context, streamPath string, con
 	err = p.Server.Streams.AddTask(subscriber, ctx).WaitStarted()
 	if err == nil {
 		select {
-		case <-subscriber.waitPublishDone:
+		case <-subscriber.waitPublishDone.Done():
 			err = subscriber.Publisher.WaitTrack()
 		case <-subscriber.Done():
 			err = subscriber.Err()

plugin/mp4/index.go

@@ -183,24 +183,12 @@ func (p *MP4Plugin) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			durVideo = video.Timestamp - lastVTime
 		}
 		bs := video.Memory.ToBytes()
-		b0 := bs[0]
-		idr := b0&0b0111_0000>>4 == 1
-		if b0&0b1000_0000 == 0 {
-			offsetVideo = 5
-			if bs[1] == 0 {
-				return nil
-			}
+		if ctx, ok := sub.VideoReader.Track.ICodecCtx.(*rtmp.H265Ctx); ok && ctx.Enhanced && bs[1]&0b1111 == rtmp.PacketTypeCodedFrames {
+			offsetVideo = 8
 		} else {
-			switch packetType := b0 & 0b1111; packetType {
-			case rtmp.PacketTypeSequenceStart:
-				return nil
-			case rtmp.PacketTypeCodedFrames:
-				offsetVideo = 8
-			case rtmp.PacketTypeCodedFramesX:
-				offsetVideo = 5
-			}
+			offsetVideo = 5
 		}
-		ctx.video.Push(&ctx, video.Timestamp, durVideo, bs[offsetVideo:], util.Conditoinal(idr, mp4.SyncSampleFlags, mp4.NonSyncSampleFlags))
+		ctx.video.Push(&ctx, video.Timestamp, durVideo, bs[offsetVideo:], util.Conditoinal(sub.VideoReader.Value.IDR, mp4.SyncSampleFlags, mp4.NonSyncSampleFlags))
 		lastVTime = video.Timestamp
 		return nil
 	})

plugin/mp4/pkg/box/mp4track.go

@@ -1,7 +1,6 @@
 package box
 
 import (
-	"encoding/binary"
 	"github.com/yapingcat/gomedia/go-codec"
 	"io"
 )
@@ -274,98 +273,25 @@ func (track *mp4track) makeEmptyStblTable() {
 	track.stbltable.stss = &movstss{}
 }
 
-func (track *mp4track) writeH264(nalus [][]byte, pts, dts uint64) (err error) {
-	//h264extra, ok := track.extra.(*h264ExtraData)
-	//if !ok {
-	//	panic("must init h264ExtraData first")
-	//}
-	for _, nalu := range nalus {
-		nalu_type := codec.H264_NAL_TYPE(nalu[0] & 0x1F)
-		//aud/sps/pps/sei 为帧间隔
-		//通过first_slice_in_mb来判断，改nalu是否为一帧的开头
-		if track.lastSample.hasVcl && isH264NewAccessUnit(nalu) {
-			var currentOffset int64
-			if currentOffset, err = track.writer.Seek(0, io.SeekCurrent); err != nil {
-				return
-			}
-			entry := sampleEntry{
-				pts:                    track.lastSample.pts,
-				dts:                    track.lastSample.dts,
-				size:                   0,
-				isKeyFrame:             track.lastSample.isKey,
-				SampleDescriptionIndex: 1,
-				offset:                 uint64(currentOffset),
-			}
-			n := 0
-			if n, err = track.writer.Write(track.lastSample.cache); err != nil {
-				return
-			}
-			entry.size = uint64(n)
-			track.addSampleEntry(entry)
-			track.lastSample.cache = track.lastSample.cache[:0]
-			track.lastSample.hasVcl = false
-		}
-		if codec.IsH264VCLNaluType(nalu_type) {
-			track.lastSample.pts = pts
-			track.lastSample.dts = dts
-			track.lastSample.hasVcl = true
-			track.lastSample.isKey = false
-			if nalu_type == codec.H264_NAL_I_SLICE {
-				track.lastSample.isKey = true
-			}
-		}
-		naluLen := uint32(len(nalu))
-		var lenBytes [4]byte
-		binary.BigEndian.PutUint32(lenBytes[:], naluLen)
-		track.lastSample.cache = append(append(track.lastSample.cache, lenBytes[:]...), nalu...)
+func (track *mp4track) write(sample Sample) (err error) {
+	var currentOffset int64
+	if currentOffset, err = track.writer.Seek(0, io.SeekCurrent); err != nil {
+		return
 	}
-	return
-}
-
-func (track *mp4track) writeH265(nalus [][]byte, pts, dts uint64) (err error) {
-	//h265extra, ok := track.extra.(*h265ExtraData)
-	//if !ok {
-	//	panic("must init h265ExtraData first")
-	//}
-	for _, nalu := range nalus {
-		nalu_type := codec.H265_NAL_TYPE((nalu[0] >> 1) & 0x3F)
-		if track.lastSample.hasVcl && isH265NewAccessUnit(nalu) {
-			var currentOffset int64
-			if currentOffset, err = track.writer.Seek(0, io.SeekCurrent); err != nil {
-				return
-			}
-			entry := sampleEntry{
-				pts:                    track.lastSample.pts,
-				dts:                    track.lastSample.dts,
-				size:                   0,
-				isKeyFrame:             track.lastSample.isKey,
-				SampleDescriptionIndex: 1,
-				offset:                 uint64(currentOffset),
-			}
-			n := 0
-			if n, err = track.writer.Write(track.lastSample.cache); err != nil {
-				return
-			}
-			entry.size = uint64(n)
-			track.addSampleEntry(entry)
-			track.lastSample.cache = track.lastSample.cache[:0]
-			track.lastSample.hasVcl = false
-		}
-		if codec.IsH265VCLNaluType(nalu_type) {
-			track.lastSample.pts = pts
-			track.lastSample.dts = dts
-			track.lastSample.hasVcl = true
-			track.lastSample.isKey = false
-			if nalu_type >= codec.H265_NAL_SLICE_BLA_W_LP && nalu_type <= codec.H265_NAL_SLICE_CRA {
-				track.lastSample.isKey = true
-			}
-		}
-		naluLen := uint32(len(nalu))
-		var lenBytes [4]byte
-		binary.BigEndian.PutUint32(lenBytes[:], naluLen)
-		track.lastSample.cache = append(append(track.lastSample.cache, lenBytes[:]...), nalu...)
+	entry := sampleEntry{
+		pts:                    uint64(sample.PTS),
+		dts:                    uint64(sample.DTS),
+		size:                   0,
+		isKeyFrame:             sample.KeyFrame,
+		SampleDescriptionIndex: 1,
+		offset:                 uint64(currentOffset),
 	}
-
+	n := 0
+	if n, err = track.writer.Write(sample.Data); err != nil {
+		return
+	}
+	entry.size = uint64(n)
+	track.addSampleEntry(entry)
 	return
 }
 

plugin/mp4/pkg/box/muxer.go

@@ -176,36 +176,30 @@ func (muxer *Movmuxer) WriteAudio(track uint32, sample []byte, dts uint64) (err
 	return err
 }
 
-func (muxer *Movmuxer) WriteVideo(track uint32, nalus [][]byte, pts uint64, dts uint64) (err error) {
-	mp4track := muxer.tracks[track]
-	switch mp4track.cid {
-	case MP4_CODEC_H264:
-		err = mp4track.writeH264(nalus, pts, dts)
-	case MP4_CODEC_H265:
-		err = mp4track.writeH265(nalus, pts, dts)
-	}
+func (muxer *Movmuxer) WriteSample(trackId uint32, data Sample) (err error) {
+	track := muxer.tracks[trackId]
+	err = track.write(data)
 	if err != nil {
-		return err
+		return
 	}
-
 	if !muxer.movFlag.isFragment() && !muxer.movFlag.isDash() {
 		return err
 	}
 	// isCustion := muxer.movFlag.has(MP4_FLAG_CUSTOM)
 	isKeyFrag := muxer.movFlag.has(MP4_FLAG_KEYFRAME)
 	if isKeyFrag {
-		if mp4track.lastSample.isKey && mp4track.duration > 0 {
+		if data.KeyFrame && track.duration > 0 {
 			err = muxer.flushFragment()
 			if err != nil {
-				return err
+				return
 			}
 			if muxer.onNewFragment != nil {
-				muxer.onNewFragment(mp4track.duration, mp4track.startPts, mp4track.startDts)
+				muxer.onNewFragment(track.duration, track.startPts, track.startDts)
 			}
 		}
 	}
 
-	return nil
+	return
 }
 
 func (muxer *Movmuxer) WriteTrailer() (err error) {

plugin/mp4/pkg/box/video.go

@@ -0,0 +1,8 @@
+package box
+
+type Sample struct {
+	KeyFrame bool
+	Data     []byte
+	PTS      uint32
+	DTS      uint32
+}

plugin/mp4/pkg/record.go

@@ -1,14 +1,14 @@
 package mp4
 
 import (
-	"os"
-	"time"
-
 	"m7s.live/m7s/v5"
 	"m7s.live/m7s/v5/pkg"
 	"m7s.live/m7s/v5/pkg/codec"
 	"m7s.live/m7s/v5/pkg/task"
 	"m7s.live/m7s/v5/plugin/mp4/pkg/box"
+	rtmp "m7s.live/m7s/v5/plugin/rtmp/pkg"
+	"os"
+	"time"
 )
 
 type WriteTrailerQueueTask struct {
@@ -21,7 +21,7 @@ func init() {
 	m7s.Servers.AddTaskLazy(&writeTrailerQueueTask)
 }
 
-func NewRecorder() *Recorder {
+func NewRecorder() m7s.IRecorder {
 	return &Recorder{}
 }
 
@@ -46,49 +46,77 @@ func (task *writeTrailerTask) Start() (err error) {
 }
 
 func (r *Recorder) Run() (err error) {
-	ctx := &r.RecordJob
+	recordJob := &r.RecordJob
+	sub := recordJob.Subscriber
 	var file *os.File
 	var muxer *box.Movmuxer
 	var audioId, videoId uint32
 	// TODO: fragment
-	if file, err = os.OpenFile(ctx.FilePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666); err != nil {
+	if file, err = os.OpenFile(recordJob.FilePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666); err != nil {
 		return
 	}
 	muxer, err = box.CreateMp4Muxer(file)
-	task := &writeTrailerTask{
+	if err != nil {
+		return
+	}
+	defer writeTrailerQueueTask.AddTask(&writeTrailerTask{
 		file:  file,
 		muxer: muxer,
-	}
-	task.Logger = r.Logger
-	defer writeTrailerQueueTask.AddTask(task)
-	ar, vr := ctx.Subscriber.AudioReader, ctx.Subscriber.VideoReader
-	if ar != nil {
-		audioTrack := ar.Track
-		switch ctx := audioTrack.ICodecCtx.GetBase().(type) {
-		case *codec.AACCtx:
-			audioId = muxer.AddAudioTrack(box.MP4_CODEC_AAC, box.WithExtraData(ctx.ConfigBytes))
-		case *codec.PCMACtx:
-			audioId = muxer.AddAudioTrack(box.MP4_CODEC_G711A, box.WithAudioSampleRate(uint32(ctx.SampleRate)), box.WithAudioChannelCount(uint8(ctx.Channels)), box.WithAudioSampleBits(uint8(ctx.SampleSize)))
-		case *codec.PCMUCtx:
-			audioId = muxer.AddAudioTrack(box.MP4_CODEC_G711U, box.WithAudioSampleRate(uint32(ctx.SampleRate)), box.WithAudioChannelCount(uint8(ctx.Channels)), box.WithAudioSampleBits(uint8(ctx.SampleSize)))
+	}, r.Logger)
+	var at, vt *pkg.AVTrack
+	//err = muxer.WriteInitSegment(file)
+	return m7s.PlayBlock(sub, func(audio *pkg.RawAudio) error {
+		if at == nil {
+			at = sub.AudioReader.Track
+			switch ctx := at.ICodecCtx.GetBase().(type) {
+			case *codec.AACCtx:
+				audioId = muxer.AddAudioTrack(box.MP4_CODEC_AAC, box.WithExtraData(ctx.ConfigBytes))
+			case *codec.PCMACtx:
+				audioId = muxer.AddAudioTrack(box.MP4_CODEC_G711A, box.WithAudioSampleRate(uint32(ctx.SampleRate)), box.WithAudioChannelCount(uint8(ctx.Channels)), box.WithAudioSampleBits(uint8(ctx.SampleSize)))
+			case *codec.PCMUCtx:
+				audioId = muxer.AddAudioTrack(box.MP4_CODEC_G711U, box.WithAudioSampleRate(uint32(ctx.SampleRate)), box.WithAudioChannelCount(uint8(ctx.Channels)), box.WithAudioSampleBits(uint8(ctx.SampleSize)))
+			}
 		}
-	}
-	if vr != nil {
-		videoTrack := vr.Track
-		switch ctx := videoTrack.ICodecCtx.GetBase().(type) {
+		return muxer.WriteSample(audioId, box.Sample{
+			Data: audio.ToBytes(),
+			PTS:  uint32(audio.Timestamp / time.Millisecond),
+			DTS:  uint32(audio.Timestamp / time.Millisecond),
+		})
+	}, func(video *rtmp.RTMPVideo) error {
+		offset := 5
+		bytes := video.ToBytes()
+		if vt == nil {
+			vt = sub.VideoReader.Track
+			switch ctx := vt.ICodecCtx.GetBase().(type) {
+			case *codec.H264Ctx:
+				videoId = muxer.AddVideoTrack(box.MP4_CODEC_H264, box.WithExtraData(ctx.Record), box.WithVideoWidth(uint32(ctx.Width())), box.WithVideoHeight(uint32(ctx.Height())))
+			case *codec.H265Ctx:
+				videoId = muxer.AddVideoTrack(box.MP4_CODEC_H265, box.WithExtraData(ctx.Record), box.WithVideoWidth(uint32(ctx.Width())), box.WithVideoHeight(uint32(ctx.Height())))
+			}
+		}
+		switch ctx := vt.ICodecCtx.(type) {
 		case *codec.H264Ctx:
-			videoId = muxer.AddVideoTrack(box.MP4_CODEC_H264, box.WithExtraData(ctx.Record))
-		case *codec.H265Ctx:
-			videoId = muxer.AddVideoTrack(box.MP4_CODEC_H265, box.WithExtraData(ctx.Record))
+			if bytes[1] == 0 {
+				return nil
+			}
+		case *rtmp.H265Ctx:
+			if ctx.Enhanced {
+				switch bytes[1] & 0b1111 {
+				case rtmp.PacketTypeCodedFrames:
+					offset += 3
+				case rtmp.PacketTypeSequenceStart:
+					return nil
+				}
+			} else if bytes[1] == 0 {
+				return nil
+			}
 		}
-	}
-	return m7s.PlayBlock(ctx.Subscriber, func(audio *pkg.RawAudio) error {
-		return muxer.WriteAudio(audioId, audio.ToBytes(), uint64(audio.Timestamp/time.Millisecond))
-	}, func(video *pkg.H26xFrame) error {
-		var nalus [][]byte
-		for _, nalu := range video.Nalus {
-			nalus = append(nalus, nalu.ToBytes())
-		}
-		return muxer.WriteVideo(videoId, nalus, uint64(video.Timestamp/time.Millisecond), uint64(video.CTS/time.Millisecond))
+
+		return muxer.WriteSample(videoId, box.Sample{
+			KeyFrame: sub.VideoReader.Value.IDR,
+			Data:     bytes[offset:],
+			PTS:      video.Timestamp + video.CTS,
+			DTS:      video.Timestamp,
+		})
 	})
 }

plugin/rtmp/pkg/net-connection.go

@@ -305,14 +305,18 @@ func (nc *NetConnection) RecvMessage() (msg *Chunk, err error) {
 					err = r.WriteAudio(msg.AVData.WrapAudio())
 				} else {
 					msg.AVData.Recycle()
-					nc.Warn("ReceiveAudio", "MessageStreamID", msg.MessageStreamID)
+					if r.PubAudio {
+						nc.Warn("ReceiveAudio", "MessageStreamID", msg.MessageStreamID)
+					}
 				}
 			case RTMP_MSG_VIDEO:
 				if r, ok := nc.Receivers[msg.MessageStreamID]; ok && r.PubVideo {
 					err = r.WriteVideo(msg.AVData.WrapVideo())
 				} else {
 					msg.AVData.Recycle()
-					nc.Warn("ReceiveVideo", "MessageStreamID", msg.MessageStreamID)
+					if r.PubVideo {
+						nc.Warn("ReceiveVideo", "MessageStreamID", msg.MessageStreamID)
+					}
 				}
 			}
 		}

plugin/rtmp/pkg/video.go

@@ -2,6 +2,7 @@ package rtmp
 
 import (
 	"encoding/binary"
+	"github.com/deepch/vdk/codec/h264parser"
 	"github.com/deepch/vdk/codec/h265parser"
 	"io"
 	"time"
@@ -44,16 +45,21 @@ func (avcc *RTMPVideo) Parse(t *AVTrack) (err error) {
 		switch fourCC {
 		case codec.FourCC_H264:
 			var ctx codec.H264Ctx
-			if _, err = ctx.RecordInfo.Unmarshal(cloneFrame.Buffers[0][reader.Offset():]); err == nil {
+			ctx.Record = cloneFrame.Buffers[0][reader.Offset():]
+			if _, err = ctx.RecordInfo.Unmarshal(ctx.Record); err == nil {
 				t.SequenceFrame = &cloneFrame
 				t.ICodecCtx = &ctx
+				ctx.SPSInfo, err = h264parser.ParseSPS(ctx.SPS())
 			}
 		case codec.FourCC_H265:
 			var ctx H265Ctx
-			if _, err = ctx.RecordInfo.Unmarshal(cloneFrame.Buffers[0][reader.Offset():]); err == nil {
+			ctx.Enhanced = enhanced
+			ctx.Record = cloneFrame.Buffers[0][reader.Offset():]
+			if _, err = ctx.RecordInfo.Unmarshal(ctx.Record); err == nil {
 				ctx.RecordInfo.LengthSizeMinusOne = 3 // Unmarshal wrong LengthSizeMinusOne
 				t.SequenceFrame = &cloneFrame
 				t.ICodecCtx = &ctx
+				ctx.SPSInfo, err = h265parser.ParseSPS(ctx.SPS())
 			}
 		case codec.FourCC_AV1:
 			var ctx AV1Ctx

publisher.go

@@ -107,7 +107,9 @@ func (t *AVTracks) Dispose() {
 	t.Lock()
 	defer t.Unlock()
 	for track := range t.Range {
-		track.Dispose()
+		if track == t.AVTrack || track.RingWriter != t.AVTrack.RingWriter {
+			track.Dispose()
+		}
 	}
 	t.AVTrack = nil
 	t.Clear()
@@ -261,7 +263,7 @@ func (p *Publisher) RemoveSubscriber(subscriber *Subscriber) {
 
 func (p *Publisher) AddSubscriber(subscriber *Subscriber) {
 	subscriber.Publisher = p
-	close(subscriber.waitPublishDone)
+	subscriber.waitPublishDone.Resolve()
 	if p.Subscribers.AddUnique(subscriber) {
 		p.Info("subscriber +1", "count", p.Subscribers.Length)
 		if subscriber.BufferTime > p.BufferTime {

subscriber.go

@@ -60,13 +60,13 @@ type Subscriber struct {
 	PubSubBase
 	config.Subscribe
 	Publisher                  *Publisher
-	waitPublishDone            chan struct{}
+	waitPublishDone            *util.Promise
 	AudioReader, VideoReader   *AVRingReader
 	StartAudioTS, StartVideoTS time.Duration
 }
 
 func createSubscriber(p *Plugin, streamPath string, conf config.Subscribe) *Subscriber {
-	subscriber := &Subscriber{Subscribe: conf, waitPublishDone: make(chan struct{})}
+	subscriber := &Subscriber{Subscribe: conf, waitPublishDone: util.NewPromise(p)}
 	subscriber.ID = task.GetNextTaskID()
 	subscriber.Plugin = p
 	subscriber.TimeoutTimer = time.NewTimer(subscriber.WaitTimeout)