Store the clients that are connected to a telnet shell in the connected viewers stats

stream/telnet/telnet.go

@@ -14,6 +14,7 @@ var (
 	// TODO Config should not be exported
 	Cfg            *Options
 	currentMessage map[string]*string
+	clientCount    map[string]int
 )
 
 // Options holds telnet package configuration
@@ -34,6 +35,7 @@ func Serve(config *Options) {
 	}
 
 	currentMessage = make(map[string]*string)
+	clientCount = make(map[string]int)
 
 	listener, err := net.Listen("tcp", Cfg.ListenAddress)
 	if err != nil {
@@ -94,15 +96,19 @@ func Serve(config *Options) {
 					}
 				}
 
+				clientCount[streamID]++
+
 				for {
 					n, err := s.Write([]byte(*currentMessage[streamID]))
 					if err != nil {
 						log.Printf("Error while sending TCP data: %s", err)
 						_ = s.Close()
+						clientCount[streamID]--
 						break
 					}
 					if n == 0 {
 						_ = s.Close()
+						clientCount[streamID]--
 						break
 					}
 					time.Sleep(time.Duration(Cfg.Delay) * time.Millisecond)
@@ -114,6 +120,14 @@ func Serve(config *Options) {
 	log.Println("Telnet server initialized")
 }
 
+// GetNumberConnectedSessions returns the numbers of clients that are viewing the stream through a telnet shell
+func GetNumberConnectedSessions(streamID string) int {
+	if !Cfg.Enabled {
+		return 0
+	}
+	return clientCount[streamID]
+}
+
 func asciiChar(pixel byte) string {
 	asciiChars := []string{"@", "#", "$", "%", "?", "*", "+", ";", ":", ",", ".", " "}
 	return asciiChars[(255-pixel)/22]

web/handler.go

@@ -4,6 +4,7 @@ package web
 import (
 	"bytes"
 	"encoding/json"
+	"gitlab.crans.org/nounous/ghostream/stream/telnet"
 	"html/template"
 	"log"
 	"net"
@@ -144,7 +145,9 @@ func statisticsHandler(w http.ResponseWriter, r *http.Request) {
 	enc := json.NewEncoder(w)
 	err := enc.Encode(struct {
 		ConnectedViewers int
-	}{webrtc.GetNumberConnectedSessions(streamID) + srt.GetNumberConnectedSessions(streamID)})
+	}{webrtc.GetNumberConnectedSessions(streamID) +
+		srt.GetNumberConnectedSessions(streamID) +
+		telnet.GetNumberConnectedSessions(streamID)})
 	if err != nil {
 		http.Error(w, "Failed to generate JSON.", http.StatusInternalServerError)
 		log.Printf("Failed to generate JSON: %s", err)