Use messaging system in SRT package

stream/srt/handler.go

@@ -5,23 +5,30 @@ import (
 	"log"
 
 	"github.com/haivision/srtgo"
+	"gitlab.crans.org/nounous/ghostream/stream"
 )
 
-func handleStreamer(s *srtgo.SrtSocket, name string, clientDataChannels map[string][]chan Packet, forwardingChannel, webrtcChannel chan Packet) {
+func handleStreamer(socket *srtgo.SrtSocket, streams map[string]stream.Stream, name string) {
+	// Check stream does not exist
+	if _, ok := streams[name]; ok {
+		log.Print("Stream already exists, refusing new streamer")
+		socket.Close()
+		return
+	}
+
+	// Create stream
 	log.Printf("New SRT streamer for stream %s", name)
+	st := *stream.New()
+	streams[name] = st
 
 	// Create a new buffer
 	// UDP packet cannot be larger than MTU (1500)
 	buff := make([]byte, 1500)
 
-	// Setup stream forwarding
-	forwardingChannel <- Packet{StreamName: name, PacketType: "register", Data: nil}
-	webrtcChannel <- Packet{StreamName: name, PacketType: "register", Data: nil}
-
 	// Read RTP packets forever and send them to the WebRTC Client
 	for {
 		// 5s timeout
-		n, err := s.Read(buff, 5000)
+		n, err := socket.Read(buff, 5000)
 		if err != nil {
 			log.Println("Error occurred while reading SRT socket:", err)
 			break
@@ -33,40 +40,50 @@ func handleStreamer(s *srtgo.SrtSocket, name string, clientDataChannels map[stri
 			break
 		}
 
-		// Send raw packet to other streams
+		// Send raw data to other streams
 		// Copy data in another buffer to ensure that the data would not be overwritten
+		// FIXME: might be unnecessary
 		data := make([]byte, n)
 		copy(data, buff[:n])
-		forwardingChannel <- Packet{StreamName: name, PacketType: "sendData", Data: data}
-		webrtcChannel <- Packet{StreamName: name, PacketType: "sendData", Data: data}
-		for _, dataChannel := range clientDataChannels[name] {
-			dataChannel <- Packet{StreamName: name, PacketType: "sendData", Data: data}
-		}
+		st.Broadcast <- data
 	}
 
-	forwardingChannel <- Packet{StreamName: name, PacketType: "close", Data: nil}
-	webrtcChannel <- Packet{StreamName: name, PacketType: "close", Data: nil}
+	// Close stream
+	st.Close()
+	socket.Close()
+	delete(streams, name)
 }
 
-func handleViewer(s *srtgo.SrtSocket, name string, dataChannel chan Packet, dataChannels map[string][]chan Packet) {
-	// FIXME Should not pass all dataChannels to one viewer
-
+func handleViewer(s *srtgo.SrtSocket, streams map[string]stream.Stream, name string) {
 	log.Printf("New SRT viewer for stream %s", name)
 
-	// Receive packets from channel and send them
+	// Get requested stream
+	st, ok := streams[name]
+	if !ok {
+		log.Println("Stream does not exist, refusing new viewer")
+		return
+	}
+
+	// Register new output
+	c := make(chan []byte, 128)
+	st.Register(c)
+
+	// Receive data and send them
 	for {
-		packet := <-dataChannel
-		if packet.PacketType == "sendData" {
-			_, err := s.Write(packet.Data, 10000)
-			if err != nil {
-				s.Close()
-				for i, channel := range dataChannels[name] {
-					if channel == dataChannel {
-						dataChannels[name] = append(dataChannels[name][:i], dataChannels[name][i+1:]...)
-					}
-				}
-				return
-			}
+		data := <-c
+		if len(data) < 1 {
+			log.Print("Remove SRT viewer because of end of stream")
+			break
+		}
+
+		_, err := s.Write(data, 1000)
+		if err != nil {
+			log.Printf("Remove SRT viewer because of sending error, %s", err)
+			break
 		}
 	}
+
+	// Close output
+	st.Unregister(c)
+	s.Close()
 }