From f165efcdc23541bf1aac0fa1c8d8b3725e135841 Mon Sep 17 00:00:00 2001
From: Muaz Ahmad <m.ahmed_2000@hotmail.com>
Date: Fri, 18 Aug 2023 14:43:52 +0500
Subject: [PATCH] Audio+video handling, break data loop on command message
 since it should only ever be an unpublish if working correctly, or exit
 anyway for non-proper protocol

---
 rtmp/data_loop.go  | 14 ++++++++++----
 rtmp/flv/writer.go | 24 ++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/rtmp/data_loop.go b/rtmp/data_loop.go
index e889944..51df3d0 100644
--- a/rtmp/data_loop.go
+++ b/rtmp/data_loop.go
@@ -1,7 +1,6 @@
 package rtmp
 
 import (
-	"fmt"
 	"stream_server/rtmp/flv"
 )
 
@@ -15,12 +14,19 @@ func HandleDataLoop(chnk_wrp_ptr *ChunkWrapper) {
 		return 
 	}
 	for {
-		p, err := chnk_wrp_ptr.ReadChunk()
+		msg_ptr, err := chnk_wrp_ptr.ReadChunk()
 		if err != nil {
 			return
+		} else if msg_ptr == nil {
+			continue
 		}
-		if p != nil {
-			fmt.Println(p.msg_type, p.msg_len)
+		switch msg_ptr.msg_type {
+		case 20:
+			return
+		case 8, 9:
+			if err = file_writer.WriteMediaTag(&(msg_ptr.data), msg_ptr.timestamp, msg_ptr.msg_type); err != nil {
+				return
+			}
 		}
 	}
 }
diff --git a/rtmp/flv/writer.go b/rtmp/flv/writer.go
index 492f601..8bf7d9f 100644
--- a/rtmp/flv/writer.go
+++ b/rtmp/flv/writer.go
@@ -49,6 +49,30 @@ func (writer *FLVWriter) WriteMetadataTag(data *[]byte) (err error) {
 	return	
 }
 
+func (writer *FLVWriter) WriteMediaTag(data *[]byte, timestamp uint32, media_type uint8) (err error) {
+	uint24_buf := make([]byte, 4)
+	tag_header := make([]byte, 11)
+	tag_header[0] = media_type
+	data_len := uint32(len(*data))
+	binary.BigEndian.PutUint32(uint24_buf, data_len)
+	copy(tag_header[1:4], uint24_buf[1:])
+
+	binary.BigEndian.PutUint32(uint24_buf, timestamp)
+	copy(tag_header[4:7], uint24_buf[1:])
+	tag_header[7] = uint24_buf[0]
+
+	if _, err = writer.w.Write(tag_header); err != nil {
+		return
+	}
+	if _, err = writer.w.Write(*data); err != nil {
+		return
+	}
+
+	tag_len_buf := make([]byte, 4)
+	binary.BigEndian.PutUint32(tag_len_buf, data_len + 11)
+	_, err = writer.w.Write(tag_len_buf)
+	return 
+}
 
 func (writer *FLVWriter) write_flv_header() (err error) {
 	header := make([]byte, 13)