From 320a2b99af6a0f4580c786af6a02d2d9035cc73b Mon Sep 17 00:00:00 2001
From: Alessandro Ros <aler9.dev@gmail.com>
Date: Fri, 28 Jul 2023 23:40:31 +0200
Subject: [PATCH] fix race condition in WritePacketRTP() (#334)

---
 client_format.go       | 25 ++++++++++++++++++++++++-
 server_stream_media.go |  2 +-
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/client_format.go b/client_format.go
index 2c5e8127..e77a67d5 100644
--- a/client_format.go
+++ b/client_format.go
@@ -2,6 +2,7 @@ package gortsplib
 
 import (
 	"fmt"
+	"io"
 	"time"
 
 	"github.com/pion/rtcp"
@@ -14,6 +15,28 @@ import (
 	"github.com/bluenviron/gortsplib/v3/pkg/rtpreorderer"
 )
 
+// workaround until this gets tagged:
+// https://github.com/pion/rtp/pull/234
+func rtpPacketMarshalToSafe(p *rtp.Packet, buf []byte) (n int, err error) {
+	n, err = p.Header.MarshalTo(buf)
+	if err != nil {
+		return 0, err
+	}
+
+	// Make sure the buffer is large enough to hold the packet.
+	if n+len(p.Payload)+int(p.PaddingSize) > len(buf) {
+		return 0, io.ErrShortBuffer
+	}
+
+	m := copy(buf[n:], p.Payload)
+
+	if p.Header.Padding {
+		buf[n+m+int(p.PaddingSize-1)] = p.PaddingSize
+	}
+
+	return n + m + int(p.PaddingSize), nil
+}
+
 type clientFormat struct {
 	c               *Client
 	cm              *clientMedia
@@ -76,7 +99,7 @@ func (ct *clientFormat) stop() {
 
 func (ct *clientFormat) writePacketRTPWithNTP(pkt *rtp.Packet, ntp time.Time) error {
 	byts := make([]byte, udpMaxPayloadSize)
-	n, err := pkt.MarshalTo(byts)
+	n, err := rtpPacketMarshalToSafe(pkt, byts)
 	if err != nil {
 		return err
 	}
diff --git a/server_stream_media.go b/server_stream_media.go
index f9f711b0..e814887c 100644
--- a/server_stream_media.go
+++ b/server_stream_media.go
@@ -69,7 +69,7 @@ func (sm *serverStreamMedia) allocateMulticastHandler(s *Server) error {
 
 func (sm *serverStreamMedia) WritePacketRTPWithNTP(ss *ServerStream, pkt *rtp.Packet, ntp time.Time) {
 	byts := make([]byte, udpMaxPayloadSize)
-	n, err := pkt.MarshalTo(byts)
+	n, err := rtpPacketMarshalToSafe(pkt, byts)
 	if err != nil {
 		return
 	}