From 8f2ff955e9ba1c2f477edb971ac6ad65e60ee51a Mon Sep 17 00:00:00 2001
From: Kioubit <kioubit@localhost.localdomain>
Date: Tue, 18 Jun 2024 15:48:53 +0300
Subject: [PATCH] Various optimizations

---
 go.mod  |  4 ++--
 go.sum  |  4 ++--
 http.go | 23 +++++++++++++++++------
 main.go | 10 +++++++---
 4 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/go.mod b/go.mod
index 4f2185a..2a4cc00 100644
--- a/go.mod
+++ b/go.mod
@@ -1,7 +1,7 @@
 module ColorPing
 
-go 1.20
+go 1.22.4
 
 require github.com/songgao/water v0.0.0-20200317203138-2b4b6d7c09d8
 
-require golang.org/x/sys v0.6.0 // indirect
+require golang.org/x/sys v0.21.0 // indirect
diff --git a/go.sum b/go.sum
index d97ca71..13d68dd 100644
--- a/go.sum
+++ b/go.sum
@@ -1,4 +1,4 @@
 github.com/songgao/water v0.0.0-20200317203138-2b4b6d7c09d8 h1:TG/diQgUe0pntT/2D9tmUCz4VNwm9MfrtPr0SU2qSX8=
 github.com/songgao/water v0.0.0-20200317203138-2b4b6d7c09d8/go.mod h1:P5HUIBuIWKbyjl083/loAegFkfbFNx5i2qEP4CNbm7E=
-golang.org/x/sys v0.6.0 h1:MVltZSvRTcU2ljQOhs94SXPftV6DCNnZViHeQps87pQ=
-golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws=
+golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
diff --git a/http.go b/http.go
index 889ef43..e4fe926 100644
--- a/http.go
+++ b/http.go
@@ -50,6 +50,9 @@ func getClientID() uint32 {
 func clearClients() {
 	clientMutex.Lock()
 	defer clientMutex.Unlock()
+	for _, c := range clients {
+		close(c.channel)
+	}
 	clients = make(map[uint32]*client)
 }
 
@@ -59,7 +62,7 @@ func httpServer() {
 	htmlTemplate = template.Must(template.ParseFS(embedFS, "template.html"))
 	http.HandleFunc("/stream", stream)
 	http.HandleFunc("/", serveRoot)
-	err = http.ListenAndServe("0.0.0.0:9090", nil)
+	err = http.ListenAndServe(":9090", nil)
 	if err != nil {
 		log.Fatalln(err)
 	}
@@ -157,7 +160,7 @@ func streamServer() {
 
 		dataInitial, dataUpdate := getPicture(requiresInitial, requiresUpdate)
 
-		for _, v := range clients {
+		for clientID, v := range clients {
 			if v.state == INITIAL {
 				v.state = ACTIVE
 				select {
@@ -170,6 +173,9 @@ func streamServer() {
 					select {
 					case v.channel <- dataUpdate:
 					default:
+						// Client cannot keep up
+						close(v.channel)
+						delete(clients, clientID)
 						continue
 					}
 				}
@@ -189,7 +195,6 @@ func stream(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "text/event-stream")
 	w.Header().Set("Cache-Control", "no-cache")
 	w.Header().Set("Connection", "keep-alive")
-	w.Header().Set("Access-Control-Allow-Origin", "*")
 	messageChan := make(chan string, 40)
 	id := getClientID()
 	newClient := client{
@@ -199,18 +204,24 @@ func stream(w http.ResponseWriter, r *http.Request) {
 	clientMutex.Lock()
 	clients[id] = &newClient
 	clientMutex.Unlock()
+
+	// For when clients are removed prior to connection closed, to avoid a call to delete(clients, id)
+	var channelClosedFirst = false
 	go func() {
 		// Listen for connection close
 		<-r.Context().Done()
 		clientMutex.Lock()
+		if !channelClosedFirst {
+			delete(clients, id)
+		}
 		close(messageChan)
-		delete(clients, id)
 		clientMutex.Unlock()
 	}()
 
 	for {
-		data := <-messageChan
-		if data == "" {
+		data, ok := <-messageChan
+		if !ok {
+			channelClosedFirst = true
 			return
 		}
 		_, _ = w.Write([]byte(data))
diff --git a/main.go b/main.go
index 165e606..0788bc7 100644
--- a/main.go
+++ b/main.go
@@ -25,7 +25,7 @@ func main() {
 	go startInterface(packetChan)
 	go streamServer()
 	fmt.Println("Kioubit ColorPing started")
-	fmt.Println("Interface name:", interfaceName)
+	fmt.Println("Interface name:", interfaceName, "HTTP server port: 9090")
 	httpServer()
 }
 
@@ -41,6 +41,10 @@ func prePopulatePixelArray() {
 	}
 }
 
+var pktPool = sync.Pool{
+	New: func() interface{} { return make([]byte, 2000) },
+}
+
 func startInterface(packetChan chan *[]byte) {
 	config := water.Config{
 		DeviceType: water.TUN,
@@ -52,7 +56,7 @@ func startInterface(packetChan chan *[]byte) {
 	}
 
 	for {
-		packet := make([]byte, 2000)
+		packet := pktPool.Get().([]byte)
 		n, err := iFace.Read(packet)
 		if err != nil {
 			log.Fatal(err)
@@ -99,7 +103,7 @@ func packetHandler(packetChan chan *[]byte) {
 			obj.changed = true
 		}
 		obj.Unlock()
-
+		pktPool.Put(*packet)
 	}
 
 }