From e1852b6ae086a7f488c4feef8063d8a5b4767892 Mon Sep 17 00:00:00 2001
From: Muaz Ahmad <mahmad2000@protonmail.com>
Date: Thu, 21 Sep 2023 13:36:11 +0500
Subject: [PATCH] stream key checking

---
 srt/protocol.go   | 23 +++++++++++++++++++----
 srt/stream_ids.go | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 4 deletions(-)
 create mode 100644 srt/stream_ids.go

diff --git a/srt/protocol.go b/srt/protocol.go
index 1636c4b..54e7c17 100644
--- a/srt/protocol.go
+++ b/srt/protocol.go
@@ -137,15 +137,31 @@ func (agent *SRTManager) create_conclusion_resp() (*Packet) {
 }
 
 func (agent *SRTManager) process_conclusion(packet *Packet) (*Packet) {
+	resp_packet := agent.create_conclusion_resp()
 	if packet.packet_type == HANDSHAKE {
 		hs_cif := packet.cif.(*HandshakeCIF)
 		if hs_cif.hs_type == 0xffffffff && hs_cif.syn_cookie == agent.syn_cookie {
+			for _, v := range hs_cif.hs_extensions {
+				switch v.ext_type {
+				case 5:
+					writer, ok := CheckStreamID(v.ext_contents.([]byte))
+					if !ok {
+						resp_packet.cif.(*HandshakeCIF).hs_type = 1003
+						return resp_packet
+					} else {
+						agent.output = writer
+					}
+				}
+			}
 			agent.pings[0][1] = time.Now()
-			agent.state = DATA_LOOP
-			return agent.create_conclusion_resp()
+			if agent.output != nil {
+				agent.state = DATA_LOOP
+				return resp_packet
+			}
 		}
 	}
-	return nil
+	resp_packet.cif.(*HandshakeCIF).hs_type = 1000
+	return resp_packet
 }
 
 func (agent *SRTManager) create_ack_report() (*Packet) {
@@ -174,7 +190,6 @@ func (agent *SRTManager) create_ack_report() (*Packet) {
 			rtt_2_sum += uint32(math.Pow(float64(rtt), 2))
 		}
 	}
-	fmt.Println(rtt, rtt_sum, rtt_2_sum, rtt_n)
 	cif.rtt = rtt
 	cif.var_rtt = uint32(rtt_2_sum / rtt_n) - uint32(math.Pow(float64(rtt_sum / rtt_n), 2))
 
diff --git a/srt/stream_ids.go b/srt/stream_ids.go
new file mode 100644
index 0000000..88bde45
--- /dev/null
+++ b/srt/stream_ids.go
@@ -0,0 +1,43 @@
+package srt
+
+import (
+	"os"
+	"io"
+	"strings"
+)
+
+func NewWriter(stream_key string) (io.WriteCloser, error) {
+	base_dir, _ := os.UserHomeDir()
+	file, err := os.Create(base_dir + "/live/" + stream_key + "/rawdump.ts")
+	if err != nil {
+		return nil, err
+	}
+	return file, nil
+}
+
+func CheckStreamID(stream_id []byte) (io.WriteCloser, bool) {
+	stream_key := make([]byte, 0)
+	for i := len(stream_id) - 1; i >= 0; i-- {
+		if stream_id[i] == 0 {
+			break
+		}
+		stream_key = append(stream_key, stream_id[i])
+	}
+	stream_key_string := strings.TrimSpace(string(stream_key))
+	if check_stream_key(stream_key_string) {
+		writer, err := NewWriter(stream_key_string)
+		if err != nil {
+			return nil, false
+		}
+		return writer, true
+	}
+	return nil, false
+}
+
+func check_stream_key(stream_key string) (bool) {
+	base_dir, _ := os.UserHomeDir()
+	if fileinfo, err := os.Stat(base_dir + "/live/" + stream_key); err == nil && fileinfo.IsDir() {
+		return true
+	}
+	return false
+}