From a20de73ab2dc6c7f7fc0a149053f6a7480b0bb8e Mon Sep 17 00:00:00 2001
From: Alexey Khit <alexey.khit@gmail.com>
Date: Mon, 24 Apr 2023 06:40:11 +0300
Subject: [PATCH] Add pkt_size option fort RTSP server

---
 cmd/rtsp/rtsp.go     |  9 ++++++++-
 pkg/h264/rtp.go      |  4 ++++
 pkg/h265/rtp.go      |  4 ++++
 pkg/rtsp/conn.go     |  1 +
 pkg/rtsp/consumer.go | 13 +++++++++++--
 pkg/webrtc/api.go    |  1 +
 6 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/cmd/rtsp/rtsp.go b/cmd/rtsp/rtsp.go
index 31ab985b..93306d88 100644
--- a/cmd/rtsp/rtsp.go
+++ b/cmd/rtsp/rtsp.go
@@ -22,6 +22,7 @@ func Init() {
 			Username     string `yaml:"username" json:"-"`
 			Password     string `yaml:"password" json:"-"`
 			DefaultQuery string `yaml:"default_query" json:"default_query"`
+			PacketSize   uint16 `yaml:"pkt_size"`
 		} `yaml:"rtsp"`
 	}
 
@@ -67,6 +68,7 @@ func Init() {
 			}
 
 			c := rtsp.NewServer(conn)
+			c.PacketSize = conf.Mod.PacketSize
 			// skip check auth for localhost
 			if conf.Mod.Username != "" && !conn.RemoteAddr().(*net.TCPAddr).IP.IsLoopback() {
 				c.Auth(conf.Mod.Username, conf.Mod.Password)
@@ -174,13 +176,18 @@ func tcpHandler(conn *rtsp.Conn) {
 
 			conn.SessionName = app.UserAgent
 
-			conn.Medias = mp4.ParseQuery(conn.URL.Query())
+			query := conn.URL.Query()
+			conn.Medias = mp4.ParseQuery(query)
 			if conn.Medias == nil {
 				for _, media := range defaultMedias {
 					conn.Medias = append(conn.Medias, media.Clone())
 				}
 			}
 
+			if s := query.Get("pkt_size"); s != "" {
+				conn.PacketSize = uint16(core.Atoi(s))
+			}
+
 			if err := stream.AddConsumer(conn); err != nil {
 				log.Warn().Err(err).Str("stream", name).Msg("[rtsp]")
 				return
diff --git a/pkg/h264/rtp.go b/pkg/h264/rtp.go
index 8cccb637..a1bc93ea 100644
--- a/pkg/h264/rtp.go
+++ b/pkg/h264/rtp.go
@@ -94,6 +94,10 @@ func RTPDepay(codec *core.Codec, handler core.HandlerFunc) core.HandlerFunc {
 }
 
 func RTPPay(mtu uint16, handler core.HandlerFunc) core.HandlerFunc {
+	if mtu == 0 {
+		mtu = 1472
+	}
+
 	payloader := &Payloader{IsAVC: true}
 	sequencer := rtp.NewRandomSequencer()
 	mtu -= 12 // rtp.Header size
diff --git a/pkg/h265/rtp.go b/pkg/h265/rtp.go
index 333ca6d4..3e027b83 100644
--- a/pkg/h265/rtp.go
+++ b/pkg/h265/rtp.go
@@ -76,6 +76,10 @@ func RTPDepay(codec *core.Codec, handler core.HandlerFunc) core.HandlerFunc {
 }
 
 func RTPPay(mtu uint16, handler core.HandlerFunc) core.HandlerFunc {
+	if mtu == 0 {
+		mtu = 1472
+	}
+
 	payloader := &Payloader{}
 	sequencer := rtp.NewRandomSequencer()
 	mtu -= 12 // rtp.Header size
diff --git a/pkg/rtsp/conn.go b/pkg/rtsp/conn.go
index 1a00ac8b..2bdb91bb 100644
--- a/pkg/rtsp/conn.go
+++ b/pkg/rtsp/conn.go
@@ -22,6 +22,7 @@ type Conn struct {
 	// public
 
 	Backchannel bool
+	PacketSize  uint16
 	SessionName string
 
 	Medias    []*core.Media
diff --git a/pkg/rtsp/consumer.go b/pkg/rtsp/consumer.go
index 7f5f1d1a..38125147 100644
--- a/pkg/rtsp/consumer.go
+++ b/pkg/rtsp/consumer.go
@@ -104,14 +104,23 @@ func (c *Conn) packetWriter(codec *core.Codec, channel, payloadType uint8) core.
 	if !codec.IsRTP() {
 		switch codec.Name {
 		case core.CodecH264:
-			handlerFunc = h264.RTPPay(1500, handlerFunc)
+			handlerFunc = h264.RTPPay(c.PacketSize, handlerFunc)
 		case core.CodecH265:
-			handlerFunc = h265.RTPPay(1500, handlerFunc)
+			handlerFunc = h265.RTPPay(c.PacketSize, handlerFunc)
 		case core.CodecAAC:
 			handlerFunc = aac.RTPPay(handlerFunc)
 		case core.CodecJPEG:
 			handlerFunc = mjpeg.RTPPay(handlerFunc)
 		}
+	} else if c.PacketSize != 0 {
+		switch codec.Name {
+		case core.CodecH264:
+			handlerFunc = h264.RTPPay(c.PacketSize, handlerFunc)
+			handlerFunc = h264.RTPDepay(codec, handlerFunc)
+		case core.CodecH265:
+			handlerFunc = h265.RTPPay(c.PacketSize, handlerFunc)
+			handlerFunc = h265.RTPDepay(codec, handlerFunc)
+		}
 	}
 
 	return handlerFunc
diff --git a/pkg/webrtc/api.go b/pkg/webrtc/api.go
index 8a5c7668..5086109d 100644
--- a/pkg/webrtc/api.go
+++ b/pkg/webrtc/api.go
@@ -9,6 +9,7 @@ import (
 )
 
 // ReceiveMTU = Ethernet MTU (1500) - IP Header (20) - UDP Header (8)
+// https://ffmpeg.org/ffmpeg-all.html#Muxer
 const ReceiveMTU = 1472
 
 func NewAPI(address string) (*webrtc.API, error) {