From a7b233681c1793d803fa3f1a8bd05f9327237c07 Mon Sep 17 00:00:00 2001
From: ydajiang <diyidacheng@gmail.com>
Date: Sun, 8 Jun 2025 09:42:29 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20rtmp=E6=8E=A8=E6=8B=89=E6=B5=81=E9=A2=91?=
 =?UTF-8?q?=E7=B9=81=E5=88=9B=E5=BB=BA=E5=86=85=E5=AD=98=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 rtmp/rtmp_server.go  |  6 +-----
 rtmp/rtmp_session.go | 18 +++++++++++++-----
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/rtmp/rtmp_server.go b/rtmp/rtmp_server.go
index 6d602db..c2e9bf4 100644
--- a/rtmp/rtmp_server.go
+++ b/rtmp/rtmp_server.go
@@ -65,11 +65,7 @@ func (s *server) OnPacket(conn net.Conn, data []byte) []byte {
 		_ = conn.Close()
 	}
 
-	if session.isPublisher {
-		return stream.TCPReceiveBufferPool.Get().([]byte)
-	}
-
-	return nil
+	return session.receiveBuffer
 }
 
 func NewServer() Server {
diff --git a/rtmp/rtmp_session.go b/rtmp/rtmp_session.go
index 42ecfd1..cd2885c 100644
--- a/rtmp/rtmp_session.go
+++ b/rtmp/rtmp_session.go
@@ -10,10 +10,11 @@ import (
 
 // Session RTMP会话, 解析处理Message
 type Session struct {
-	conn        net.Conn
-	stack       *rtmp.ServerStack // rtmp协议栈, 解析message
-	handle      interface{}       // 持有具体会话句柄(推流端/拉流端)， 在@see OnPublish @see OnPlay回调中赋值
-	isPublisher bool              // 是否是推流会话
+	conn          net.Conn
+	stack         *rtmp.ServerStack // rtmp协议栈, 解析message
+	handle        interface{}       // 持有具体会话句柄(推流端/拉流端)， 在@see OnPublish @see OnPlay回调中赋值
+	isPublisher   bool              // 是否是推流会话
+	receiveBuffer []byte
 }
 
 func (s *Session) generateSourceID(app, stream string) string {
@@ -95,6 +96,11 @@ func (s *Session) Close() {
 			s.stack.Close()
 			s.stack = nil
 		}
+
+		if s.receiveBuffer != nil {
+			stream.TCPReceiveBufferPool.Put(s.receiveBuffer[:cap(s.receiveBuffer)])
+			s.receiveBuffer = nil
+		}
 	}()
 
 	// 还未确定会话类型, 无需处理
@@ -118,7 +124,9 @@ func (s *Session) Close() {
 }
 
 func NewSession(conn net.Conn) *Session {
-	session := &Session{}
+	session := &Session{
+		receiveBuffer: stream.TCPReceiveBufferPool.Get().([]byte),
+	}
 
 	stackServer := rtmp.NewStackServer(false)
 	stackServer.SetOnStreamHandler(session)