From b3de3cf80e51fdc25d7bdab925843adff08a83c8 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Mon, 19 Dec 2022 13:06:08 +0100
Subject: [PATCH] move media-related logics into serverStreamMedia

---
 serverstream.go      | 27 ++-------------------------
 serverstreammedia.go | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 25 deletions(-)

diff --git a/serverstream.go b/serverstream.go
index 0aeac1d7..bc351d56 100644
--- a/serverstream.go
+++ b/serverstream.go
@@ -40,7 +40,7 @@ func NewServerStream(medias media.Medias) *ServerStream {
 
 	st.streamMedias = make(map[*media.Media]*serverStreamMedia, len(medias))
 	for _, media := range medias {
-		ssm := &serverStreamMedia{}
+		ssm := newServerStreamMedia(media)
 
 		ssm.formats = make(map[uint8]*serverStreamFormat)
 		for _, forma := range media.Formats {
@@ -279,13 +279,6 @@ func (st *ServerStream) WritePacketRTP(medi *media.Media, pkt *rtp.Packet) {
 // ntp is the absolute time of the packet, and is needed to generate RTCP sender reports
 // that allows the receiver to reconstruct the absolute time of the packet.
 func (st *ServerStream) WritePacketRTPWithNTP(medi *media.Media, pkt *rtp.Packet, ntp time.Time) {
-	byts := make([]byte, maxPacketSize)
-	n, err := pkt.MarshalTo(byts)
-	if err != nil {
-		return
-	}
-	byts = byts[:n]
-
 	st.mutex.RLock()
 	defer st.mutex.RUnlock()
 
@@ -294,23 +287,7 @@ func (st *ServerStream) WritePacketRTPWithNTP(medi *media.Media, pkt *rtp.Packet
 	}
 
 	sm := st.streamMedias[medi]
-
-	forma := sm.formats[pkt.PayloadType]
-
-	forma.rtcpSender.ProcessPacket(pkt, ntp, forma.format.PTSEqualsDTS(pkt))
-
-	// send unicast
-	for r := range st.activeUnicastReaders {
-		sm, ok := r.setuppedMedias[medi]
-		if ok {
-			sm.writePacketRTP(byts)
-		}
-	}
-
-	// send multicast
-	if sm.multicastHandler != nil {
-		sm.multicastHandler.writePacketRTP(byts)
-	}
+	sm.WritePacketRTPWithNTP(st, pkt, ntp)
 }
 
 // WritePacketRTCP writes a RTCP packet to all the readers of the stream.
diff --git a/serverstreammedia.go b/serverstreammedia.go
index 19deec6c..8832b5ea 100644
--- a/serverstreammedia.go
+++ b/serverstreammedia.go
@@ -1,10 +1,25 @@
 package gortsplib
 
+import (
+	"time"
+
+	"github.com/pion/rtp"
+
+	"github.com/aler9/gortsplib/v2/pkg/media"
+)
+
 type serverStreamMedia struct {
+	media            *media.Media
 	formats          map[uint8]*serverStreamFormat
 	multicastHandler *serverMulticastHandler
 }
 
+func newServerStreamMedia(medi *media.Media) *serverStreamMedia {
+	return &serverStreamMedia{
+		media: medi,
+	}
+}
+
 func (sm *serverStreamMedia) close() {
 	for _, tr := range sm.formats {
 		if tr.rtcpSender != nil {
@@ -28,3 +43,29 @@ func (sm *serverStreamMedia) allocateMulticastHandler(s *Server) error {
 	}
 	return nil
 }
+
+func (sm *serverStreamMedia) WritePacketRTPWithNTP(ss *ServerStream, pkt *rtp.Packet, ntp time.Time) {
+	byts := make([]byte, maxPacketSize)
+	n, err := pkt.MarshalTo(byts)
+	if err != nil {
+		return
+	}
+	byts = byts[:n]
+
+	forma := sm.formats[pkt.PayloadType]
+
+	forma.rtcpSender.ProcessPacket(pkt, ntp, forma.format.PTSEqualsDTS(pkt))
+
+	// send unicast
+	for r := range ss.activeUnicastReaders {
+		sm, ok := r.setuppedMedias[sm.media]
+		if ok {
+			sm.writePacketRTP(byts)
+		}
+	}
+
+	// send multicast
+	if sm.multicastHandler != nil {
+		sm.multicastHandler.writePacketRTP(byts)
+	}
+}