From 313a5e01c50ae13d5f1fb126328a2c0fa9a973d2 Mon Sep 17 00:00:00 2001
From: langhuihui <178529795@qq.com>
Date: Fri, 10 Nov 2023 09:31:28 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=AF=B9=E6=B5=81=E4=B8=AD=E6=98=AF?=
 =?UTF-8?q?=E5=90=A6=E5=8C=85=E5=90=AB=E9=9F=B3=E8=A7=86=E9=A2=91=E8=BF=9B?=
 =?UTF-8?q?=E8=A1=8C=E6=8F=90=E5=89=8D=E5=88=A4=E6=96=AD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 mpegps/demuxer_v3.go | 21 ++++++++++++---------
 mpegps/mpegps.go     |  2 +-
 publisher.go         |  5 ++++-
 3 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/mpegps/demuxer_v3.go b/mpegps/demuxer_v3.go
index 06a32ec..1e84f41 100644
--- a/mpegps/demuxer_v3.go
+++ b/mpegps/demuxer_v3.go
@@ -30,8 +30,8 @@ type DecPSPackage struct {
 	PTS     uint32
 	DTS     uint32
 	EsHandler
-	audio MpegPsEsStream
-	video MpegPsEsStream
+	audio *MpegPsEsStream
+	video *MpegPsEsStream
 }
 
 func (dec *DecPSPackage) clean() {
@@ -95,16 +95,15 @@ loop:
 			err = dec.decProgramStreamMap()
 		case StartCodeVideo:
 			payload, err = dec.ReadPayload()
-			if err == nil {
+			if err == nil && dec.video != nil {
 				if frame, err = dec.video.parsePESPacket(payload); err == nil && frame.Buffer.Len() > 0 {
 					dec.ReceiveVideo(frame)
 				}
 			}
 		case StartCodeAudio:
 			payload, err = dec.ReadPayload()
-			if err == nil {
-				frame, err = dec.audio.parsePESPacket(payload)
-				if err == nil && frame.Buffer.Len() > 0 {
+			if err == nil && dec.audio != nil {
+				if frame, err = dec.audio.parsePESPacket(payload); err == nil && frame.Buffer.Len() > 0 {
 					dec.ReceiveAudio(frame)
 				}
 			}
@@ -156,7 +155,6 @@ func (dec *DecPSPackage) decProgramStreamMap() error {
 	if err != nil {
 		return err
 	}
-	defer dec.EsHandler.ReceivePSM(psm)
 	l := len(psm)
 	index := 2
 	programStreamInfoLen := util.BigEndian.Uint16(psm[index:])
@@ -173,9 +171,13 @@ func (dec *DecPSPackage) decProgramStreamMap() error {
 		elementaryStreamID := psm[index]
 		index++
 		if elementaryStreamID >= 0xe0 && elementaryStreamID <= 0xef {
-			dec.video.Type = streamType
+			dec.video = &MpegPsEsStream{
+				Type: streamType,
+			}
 		} else if elementaryStreamID >= 0xc0 && elementaryStreamID <= 0xdf {
-			dec.audio.Type = streamType
+			dec.audio = &MpegPsEsStream{
+				Type: streamType,
+			}
 		}
 		if l <= index+1 {
 			break
@@ -185,5 +187,6 @@ func (dec *DecPSPackage) decProgramStreamMap() error {
 		index += int(elementaryStreamInfoLength)
 		programStreamMapLen -= 4 + elementaryStreamInfoLength
 	}
+	dec.EsHandler.ReceivePSM(psm, dec.audio != nil, dec.video != nil)
 	return nil
 }
diff --git a/mpegps/mpegps.go b/mpegps/mpegps.go
index 67ad0ef..0e392da 100644
--- a/mpegps/mpegps.go
+++ b/mpegps/mpegps.go
@@ -21,7 +21,7 @@ const (
 type EsHandler interface {
 	ReceiveAudio(MpegPsEsStream)
 	ReceiveVideo(MpegPsEsStream)
-	ReceivePSM(util.Buffer)
+	ReceivePSM(util.Buffer, bool, bool)
 }
 
 type MpegPsEsStream struct {
diff --git a/publisher.go b/publisher.go
index 9a7ba75..e172b4d 100644
--- a/publisher.go
+++ b/publisher.go
@@ -310,10 +310,13 @@ func (p *PSPublisher) Replay(f *os.File) (err error) {
 	}
 	return
 }
-func (p *PSPublisher) ReceivePSM(buf util.Buffer) {
+
+func (p *PSPublisher) ReceivePSM(buf util.Buffer, hasAudio bool, hasVideo bool) {
 	if p.relayTrack != nil {
 		p.relayTrack.PSM = buf.Clone()
 	}
+	p.Config.PubAudio = hasAudio
+	p.Config.PubVideo = hasVideo
 }
 
 func (p *PSPublisher) Receive(streamPath, dump, port string, ssrc uint32, reuse bool) (err error) {