From 4bb97ee43a35d68ea395745767a472078d082017 Mon Sep 17 00:00:00 2001
From: Yifu Yu <root@jackyyf.com>
Date: Fri, 1 Mar 2024 17:34:41 +0800
Subject: [PATCH] Fix: HEAD requests may keep a stream active forever.

---
 main.go | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/main.go b/main.go
index 53733e5..0a13ed4 100644
--- a/main.go
+++ b/main.go
@@ -93,10 +93,9 @@ func NewClientChannel(name string) *ClientChannel {
 	ctx.newChanLock.Lock()
 	ctx.newClientsChanList.PushBack(ch)
 	ctx.newChanLock.Unlock()
-	if atomic.AddInt64(&ctx.viewer, 1) == 1 {
-		ctx.viewerCond.Broadcast()
-	}
-	fmt.Printf("Current viewers for stream %s: %d\n", name, ctx.viewer)
+	cntViewer := atomic.AddInt64(&ctx.viewer, 1)
+	ctx.viewerCond.Broadcast()
+	fmt.Printf("Current viewers for stream %s: %d\n", name, cntViewer)
 	return ch
 }
 
@@ -166,6 +165,7 @@ func broadcastFiber(name string) {
 		chunk := <-ctx.upstreamChan
 		// Send to all existing clients
 		e := ctx.clientsChanList.Front()
+	enumerate_client:
 		for e != nil {
 			ch := e.Value.(*ClientChannel)
 			// Note: select in Golang is not ordered!
@@ -176,7 +176,7 @@ func broadcastFiber(name string) {
 				next := e.Next()
 				ctx.clientsChanList.Remove(e)
 				e = next
-				break
+				continue enumerate_client
 			default:
 			}
 			select {
@@ -200,13 +200,17 @@ func clientHandler(w http.ResponseWriter, r *http.Request) {
 		remote = r.Header.Get("X-Remote-Addr")
 	}
 	fmt.Printf("Client connection from %s accepted.\n", remote)
-	w.Header().Add("Content-Type", "video/MP2T")
-	w.WriteHeader(200)
 	name := r.URL.Path
 	if _, ok := upstreams[name]; !ok {
 		http.NotFound(w, r)
 		return
 	}
+	w.Header().Add("Content-Type", "video/MP2T")
+	w.WriteHeader(200)
+	if r.Method == http.MethodHead {
+		// Fuck you golang, you should never eat write but fail
+		return
+	}
 	ch := NewClientChannel(name)
 	defer ch.Close()
 	for {