From 7532484041d58669552015c16982e10f853491b1 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Sat, 12 Aug 2023 19:34:21 +0200
Subject: [PATCH] add option MaxPacketSize to Client and Server

---
 client.go               | 14 +++++++++++---
 client_media.go         |  2 +-
 server.go               | 12 ++++++++++--
 server_record_test.go   |  2 +-
 server_session.go       |  4 +++-
 server_session_media.go |  2 +-
 server_stream.go        |  2 +-
 7 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/client.go b/client.go
index f572ab5a..c20e797b 100644
--- a/client.go
+++ b/client.go
@@ -243,6 +243,10 @@ type Client struct {
 	// It allows to queue packets before sending them.
 	// It defaults to 256.
 	WriteBufferCount int
+	// maximum size of outgoing RTP / RTCP packets.
+	// This must be less than the UDP MTU (1472 bytes).
+	// It defaults to 1472.
+	MaxPacketSize int
 	// user agent header.
 	// It defaults to "gortsplib"
 	UserAgent string
@@ -343,10 +347,14 @@ func (c *Client) Start(scheme string, host string) error {
 	}
 	if c.WriteBufferCount == 0 {
 		c.WriteBufferCount = 256
-	}
-	if (c.WriteBufferCount & (c.WriteBufferCount - 1)) != 0 {
+	} else if (c.WriteBufferCount & (c.WriteBufferCount - 1)) != 0 {
 		return fmt.Errorf("WriteBufferCount must be a power of two")
 	}
+	if c.MaxPacketSize == 0 {
+		c.MaxPacketSize = udpMaxPayloadSize
+	} else if c.MaxPacketSize > udpMaxPayloadSize {
+		return fmt.Errorf("MaxPacketSize must be less than %d", udpMaxPayloadSize)
+	}
 	if c.UserAgent == "" {
 		c.UserAgent = "gortsplib"
 	}
@@ -1650,7 +1658,7 @@ func (c *Client) WritePacketRTP(medi *media.Media, pkt *rtp.Packet) error {
 
 // WritePacketRTPWithNTP writes a RTP packet to the media stream.
 func (c *Client) WritePacketRTPWithNTP(medi *media.Media, pkt *rtp.Packet, ntp time.Time) error {
-	byts := make([]byte, udpMaxPayloadSize)
+	byts := make([]byte, c.MaxPacketSize)
 	n, err := pkt.MarshalTo(byts)
 	if err != nil {
 		return err
diff --git a/client_media.go b/client_media.go
index 9f386189..7e89f1fa 100644
--- a/client_media.go
+++ b/client_media.go
@@ -118,7 +118,7 @@ func (cm *clientMedia) start() {
 
 		cm.tcpRTPFrame = &base.InterleavedFrame{Channel: cm.tcpChannel}
 		cm.tcpRTCPFrame = &base.InterleavedFrame{Channel: cm.tcpChannel + 1}
-		cm.tcpBuffer = make([]byte, udpMaxPayloadSize+4)
+		cm.tcpBuffer = make([]byte, cm.c.MaxPacketSize+4)
 	}
 
 	for _, ct := range cm.formats {
diff --git a/server.go b/server.go
index 0fef65f2..9983219a 100644
--- a/server.go
+++ b/server.go
@@ -90,6 +90,10 @@ type Server struct {
 	// It allows to queue packets before sending them.
 	// It defaults to 256.
 	WriteBufferCount int
+	// maximum size of outgoing RTP / RTCP packets.
+	// This must be less than the UDP MTU (1472 bytes).
+	// It defaults to 1472.
+	MaxPacketSize int
 	// disable automatic RTCP sender reports.
 	DisableRTCPSenderReports bool
 
@@ -154,10 +158,14 @@ func (s *Server) Start() error {
 	}
 	if s.WriteBufferCount == 0 {
 		s.WriteBufferCount = 256
-	}
-	if (s.WriteBufferCount & (s.WriteBufferCount - 1)) != 0 {
+	} else if (s.WriteBufferCount & (s.WriteBufferCount - 1)) != 0 {
 		return fmt.Errorf("WriteBufferCount must be a power of two")
 	}
+	if s.MaxPacketSize == 0 {
+		s.MaxPacketSize = udpMaxPayloadSize
+	} else if s.MaxPacketSize > udpMaxPayloadSize {
+		return fmt.Errorf("MaxPacketSize must be less than %d", udpMaxPayloadSize)
+	}
 
 	// system functions
 	if s.Listen == nil {
diff --git a/server_record_test.go b/server_record_test.go
index 184be30a..35163343 100644
--- a/server_record_test.go
+++ b/server_record_test.go
@@ -520,7 +520,7 @@ func TestServerRecord(t *testing.T) {
 						}, nil, nil
 					},
 					onRecord: func(ctx *ServerHandlerOnRecordCtx) (*base.Response, error) {
-						// send RTCP packets directly to the session.
+						// queue sending of RTCP packets.
 						// these are sent after the response, only if onRecord returns StatusOK.
 						err := ctx.Session.WritePacketRTCP(ctx.Session.AnnouncedMedias()[0], &testRTCPPacket)
 						require.NoError(t, err)
diff --git a/server_session.go b/server_session.go
index 567d2bfb..e109d8e6 100644
--- a/server_session.go
+++ b/server_session.go
@@ -1161,10 +1161,12 @@ func (ss *ServerSession) writePacketRTP(medi *media.Media, byts []byte) {
 
 // WritePacketRTP writes a RTP packet to the session.
 func (ss *ServerSession) WritePacketRTP(medi *media.Media, pkt *rtp.Packet) error {
-	byts, err := pkt.Marshal()
+	byts := make([]byte, ss.s.MaxPacketSize)
+	n, err := pkt.MarshalTo(byts)
 	if err != nil {
 		return err
 	}
+	byts = byts[:n]
 
 	ss.writePacketRTP(medi, byts)
 	return nil
diff --git a/server_session_media.go b/server_session_media.go
index c472bf8c..d911621d 100644
--- a/server_session_media.go
+++ b/server_session_media.go
@@ -93,7 +93,7 @@ func (sm *serverSessionMedia) start() {
 
 		sm.tcpRTPFrame = &base.InterleavedFrame{Channel: sm.tcpChannel}
 		sm.tcpRTCPFrame = &base.InterleavedFrame{Channel: sm.tcpChannel + 1}
-		sm.tcpBuffer = make([]byte, udpMaxPayloadSize+4)
+		sm.tcpBuffer = make([]byte, sm.ss.s.MaxPacketSize+4)
 	}
 }
 
diff --git a/server_stream.go b/server_stream.go
index 8438afbc..3a488374 100644
--- a/server_stream.go
+++ b/server_stream.go
@@ -251,7 +251,7 @@ func (st *ServerStream) WritePacketRTP(medi *media.Media, pkt *rtp.Packet) error
 // 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) error {
-	byts := make([]byte, udpMaxPayloadSize)
+	byts := make([]byte, st.s.MaxPacketSize)
 	n, err := pkt.MarshalTo(byts)
 	if err != nil {
 		return err