diff --git a/srt/protocol.go b/srt/protocol.go index 819ab38..2b945f6 100644 --- a/srt/protocol.go +++ b/srt/protocol.go @@ -31,6 +31,7 @@ type SRTManager struct { bw uint32 mtu uint32 output io.WriteCloser + stream_key string } func NewSRTManager(l net.PacketConn) (*SRTManager) { @@ -145,7 +146,8 @@ func (agent *SRTManager) process_conclusion(packet *Packet) (*Packet) { for _, v := range hs_cif.hs_extensions { switch v.ext_type { case 5: - writer, ok := CheckStreamID(v.ext_contents.([]byte)) + writer, stream_key, ok := CheckStreamID(v.ext_contents.([]byte)) + agent.stream_key = stream_key if !ok { resp_packet.cif.(*HandshakeCIF).hs_type = 1003 return resp_packet diff --git a/srt/stream_ids.go b/srt/stream_ids.go index d6548dc..bd18e53 100644 --- a/srt/stream_ids.go +++ b/srt/stream_ids.go @@ -5,6 +5,8 @@ import ( "io" "strings" "stream_server/transcoder" + "time" + "path/filepath" ) func NewWriter(stream_key string) (io.WriteCloser, error) { @@ -15,7 +17,20 @@ func NewWriter(stream_key string) (io.WriteCloser, error) { return transcoder_in, nil } -func CheckStreamID(stream_id []byte) (io.WriteCloser, bool) { +func CleanFiles(stream_key string, delay time.Duration) { + time.Sleep(delay * time.Second) + base_dir, _ := os.UserHomeDir() + stream_dir := base_dir + "/live/" + stream_key + fileinfo, _ := os.Stat(stream_dir + "/stream.m3u8") + if time.Now().Sub(fileinfo.ModTime()) > delay * time.Second { + leftover_files, _ := filepath.Glob(stream_dir + "/*") + for _, file := range leftover_files { + os.Remove(file) + } + } +} + +func CheckStreamID(stream_id []byte) (io.WriteCloser, string, bool) { stream_key := make([]byte, 0) for i := len(stream_id) - 1; i >= 0; i-- { if stream_id[i] == 0 { @@ -27,11 +42,11 @@ func CheckStreamID(stream_id []byte) (io.WriteCloser, bool) { if check_stream_key(stream_key_string) { writer, err := NewWriter(stream_key_string) if err != nil { - return nil, false + return nil, stream_key_string, false } - return writer, true + return writer, stream_key_string, true } - return nil, false + return nil, stream_key_string, false } func check_stream_key(stream_key string) (bool) { diff --git a/srt/tunnel.go b/srt/tunnel.go index ac0acbd..4098098 100644 --- a/srt/tunnel.go +++ b/srt/tunnel.go @@ -48,6 +48,7 @@ func (tunnel *Tunnel) Shutdown() { tunnel.WritePacket(packet) if tunnel.state.output != nil { tunnel.state.output.Close() + go CleanFiles(tunnel.state.stream_key, 10) } } }