From 5df64adcbd7dab7e3077c5246342da7eee7f7f65 Mon Sep 17 00:00:00 2001
From: Muaz Ahmad <m.ahmed_2000@hotmail.com>
Date: Thu, 10 Aug 2023 20:34:10 +0500
Subject: [PATCH] Refactor, wrapper around ReadChunk to manage structs

---
 rtmp/chunk_wrap.go | 47 ++++++++++++++++++++++++++++++++++++++++++++++
 rtmp/connect.go    |  5 ++---
 rtmp/server.go     |  5 ++---
 3 files changed, 51 insertions(+), 6 deletions(-)
 create mode 100644 rtmp/chunk_wrap.go

diff --git a/rtmp/chunk_wrap.go b/rtmp/chunk_wrap.go
new file mode 100644
index 0000000..3d1e688
--- /dev/null
+++ b/rtmp/chunk_wrap.go
@@ -0,0 +1,47 @@
+package rtmp
+
+import (
+	"net"
+)
+
+type ChunkWrapper struct {
+	conn net.Conn
+	params *ProtocolParams
+	open_chnkstrms map[uint32]*ChunkStream
+	open_msgs map[uint32]*Message
+	chunk_buffs *ChunkBuffers
+}
+
+func NewChunkWrapper(conn net.Conn) (chnk_wrp_ptr *ChunkWrapper) {
+	chnk_wrp_ptr = new(ChunkWrapper)
+
+	chnk_wrp_ptr.conn = conn
+	chnk_wrp_ptr.params = &ProtocolParams{1024, 512, 0, 0}
+	chnk_wrp_ptr.open_chnkstrms = make(map[uint32]*ChunkStream)
+	chnk_wrp_ptr.open_msgs = make(map[uint32]*Message)
+	buffers := ChunkBuffers{
+		make([]byte, 4),
+		make([]byte, 4),
+		make([]byte, 1),
+		make([]byte, 4),
+		make([]byte, 1),
+		make([]byte, 2),
+	}
+	chnk_wrp_ptr.chunk_buffs = &buffers
+	return
+}
+
+func (chnk_wrp_ptr *ChunkWrapper) ReadChunk() (*Message, error) {
+	full_msg_ptr, err := ReadChunk(
+		chnk_wrp_ptr.conn,
+		chnk_wrp_ptr.open_chnkstrms,
+		chnk_wrp_ptr.open_msgs,
+		chnk_wrp_ptr.params.peer_chunk_size,
+		chnk_wrp_ptr.chunk_buffs,
+		&(chnk_wrp_ptr.params.curr_read),
+	)
+	if err != nil {
+		return nil, err
+	}
+	return full_msg_ptr, nil
+}
diff --git a/rtmp/connect.go b/rtmp/connect.go
index 880d6fb..b39d09b 100644
--- a/rtmp/connect.go
+++ b/rtmp/connect.go
@@ -1,13 +1,12 @@
 package rtmp
 
 import (
-	"net"
 	"fmt"
 )
 
-func NegotiateConnect(conn net.Conn, params *ProtocolParams, open_chnkstrms map[uint32]*ChunkStream, open_msgs map[uint32]*Message, chunk_bufs *ChunkBuffers) (bool) {
+func NegotiateConnect(chnk_wrp_ptr *ChunkWrapper) (bool) {
 	for {
-		full_msg_ptr, err := ReadChunk(conn, open_chnkstrms, open_msgs, params.peer_chunk_size, chunk_bufs, &(params.curr_read))
+		full_msg_ptr, err := chnk_wrp_ptr.ReadChunk()
 		if err != nil {
 			return false
 		}
diff --git a/rtmp/server.go b/rtmp/server.go
index ed6b474..fb26812 100644
--- a/rtmp/server.go
+++ b/rtmp/server.go
@@ -40,9 +40,8 @@ func handle_conn(conn net.Conn, stream_live *bool) {
 	if !DoHandshake(conn) {
 		return
 	}
-	params := ProtocolParams{1024, 512, 0, 0}
-	open_chnkstrms, open_msgs, chunk_buffers := OpenStreamsMapInit()
-	if !NegotiateConnect(conn, &params, open_chnkstrms, open_msgs, chunk_buffers) {
+	chunk_wrapper := NewChunkWrapper(conn)
+	if !NegotiateConnect(chunk_wrapper) {
 		return
 	}