From 5d0c51817c6d97cf5ad6a2358df684e5d303c40c Mon Sep 17 00:00:00 2001
From: yangjiechina <1534796060@qq.com>
Date: Fri, 13 Dec 2024 19:46:44 +0800
Subject: [PATCH] =?UTF-8?q?=E5=85=BC=E5=AE=B9Adoble=20FMLE=E6=8E=A8?=
 =?UTF-8?q?=E6=B5=81=E5=B7=A5=E5=85=B7?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 rtmp/rtmp_publisher.go | 2 +-
 rtmp/rtmp_server.go    | 2 +-
 rtmp/rtmp_session.go   | 9 ++++-----
 rtmp/rtmp_sink.go      | 4 ++--
 rtmp/rtmp_stream.go    | 8 ++++----
 5 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/rtmp/rtmp_publisher.go b/rtmp/rtmp_publisher.go
index 42e5461..9d9a7b1 100644
--- a/rtmp/rtmp_publisher.go
+++ b/rtmp/rtmp_publisher.go
@@ -16,7 +16,7 @@ type Publisher struct {
 }
 
 func (p *Publisher) Input(data []byte) error {
-	return p.stack.Input(nil, data)
+	return p.stack.Input(data)
 }
 
 func (p *Publisher) OnDeMuxStream(stream utils.AVStream) {
diff --git a/rtmp/rtmp_server.go b/rtmp/rtmp_server.go
index a2b712f..c9b0c34 100644
--- a/rtmp/rtmp_server.go
+++ b/rtmp/rtmp_server.go
@@ -58,7 +58,7 @@ func (s *server) OnPacket(conn net.Conn, data []byte) []byte {
 	s.StreamServer.OnPacket(conn, data)
 
 	session := conn.(*transport.Conn).Data.(*Session)
-	err := session.Input(conn, data)
+	err := session.Input(data)
 
 	if err != nil {
 		log.Sugar.Errorf("处理rtmp包失败 err:%s conn:%s", err.Error(), conn.RemoteAddr().String())
diff --git a/rtmp/rtmp_session.go b/rtmp/rtmp_session.go
index a0ed8b1..cdd3542 100644
--- a/rtmp/rtmp_session.go
+++ b/rtmp/rtmp_session.go
@@ -76,13 +76,12 @@ func (s *Session) OnPlay(app, stream_ string) utils.HookState {
 	return state
 }
 
-func (s *Session) Input(conn net.Conn, data []byte) error {
+func (s *Session) Input(data []byte) error {
 	// 推流会话, 收到的包都将交由主协程处理
 	if s.isPublisher {
-		s.handle.(*Publisher).PublishSource.Input(data)
-		return nil
+		return s.handle.(*Publisher).PublishSource.Input(data)
 	} else {
-		return s.stack.Input(conn, data)
+		return s.stack.Input(data)
 	}
 }
 
@@ -121,7 +120,7 @@ func (s *Session) Close() {
 func NewSession(conn net.Conn) *Session {
 	session := &Session{}
 
-	stack := librtmp.NewStack(session)
+	stack := librtmp.NewStack(conn, session)
 	session.stack = stack
 	session.conn = conn
 	return session
diff --git a/rtmp/rtmp_sink.go b/rtmp/rtmp_sink.go
index b139383..8575f65 100644
--- a/rtmp/rtmp_sink.go
+++ b/rtmp/rtmp_sink.go
@@ -13,11 +13,11 @@ type Sink struct {
 }
 
 func (s *Sink) StartStreaming(_ stream.TransStream) error {
-	return s.stack.SendStreamBeginChunk(s.Conn)
+	return s.stack.SendStreamBeginChunk()
 }
 
 func (s *Sink) StopStreaming(_ stream.TransStream) {
-	_ = s.stack.SendStreamEOFChunk(s.Conn)
+	_ = s.stack.SendStreamEOFChunk()
 }
 
 func (s *Sink) Close() {
diff --git a/rtmp/rtmp_stream.go b/rtmp/rtmp_stream.go
index c6c7b5f..cf798bb 100644
--- a/rtmp/rtmp_stream.go
+++ b/rtmp/rtmp_stream.go
@@ -75,7 +75,7 @@ func (t *transStream) Input(packet utils.AVPacket) ([][]byte, int64, bool, error
 	// 写chunk header
 	chunk.Length = payloadSize
 	chunk.Timestamp = uint32(dts)
-	n := chunk.ToBytes(allocate)
+	n := chunk.MarshalHeader(allocate)
 
 	// 写flv
 	if videoPkt {
@@ -84,7 +84,7 @@ func (t *transStream) Input(packet utils.AVPacket) ([][]byte, int64, bool, error
 		n += t.muxer.WriteAudioData(allocate[n:], false)
 	}
 
-	n += chunk.WriteData(allocate[n:], data, t.chunkSize, chunkPayloadOffset)
+	n += chunk.WriteBody(allocate[n:], data, t.chunkSize, chunkPayloadOffset)
 	utils.Assert(len(allocate) == n)
 
 	// 合并写满了再发
@@ -157,7 +157,7 @@ func (t *transStream) WriteHeader() error {
 		n += len(extra)
 
 		t.audioChunk.Length = n
-		t.audioChunk.ToBytes(t.header)
+		t.audioChunk.MarshalHeader(t.header)
 		n += 12
 	}
 
@@ -169,7 +169,7 @@ func (t *transStream) WriteHeader() error {
 		n += len(extra)
 
 		t.videoChunk.Length = 5 + len(extra)
-		t.videoChunk.ToBytes(t.header[tmp:])
+		t.videoChunk.MarshalHeader(t.header[tmp:])
 		n += 12
 	}