1 files changed, 67 insertions, 4 deletions
diff --git a/internal/api/websocket.go b/internal/api/websocket.go
index 6bd8c88..b5bf728 100644
--- a/internal/api/websocket.go
+++ b/internal/api/websocket.go
@@ -1,13 +1,27 @@
 package api
 
 import (
+	"errors"
 	"log/slog"
 	"net/http"
+	"strings"
 	"sync"
+	"time"
 
 	"golang.org/x/net/websocket"
 )
 
+// wsPingInterval and wsPingDeadline control heartbeat timing.
+// Exposed as vars so tests can override them without rebuilding.
+var (
+	wsPingInterval = 30 * time.Second
+	wsPingDeadline = 10 * time.Second
+
+	// maxWsClients caps the number of concurrent WebSocket connections.
+	// Exposed as a var so tests can override it.
+	maxWsClients = 1000
+)
+
 // Hub manages WebSocket connections and broadcasts messages.
 type Hub struct {
 	mu      sync.RWMutex
@@ -25,10 +39,14 @@ func NewHub() *Hub {
 // Run is a no-op loop kept for future cleanup/heartbeat logic.
 func (h *Hub) Run() {}
 
-func (h *Hub) Register(ws *websocket.Conn) {
+func (h *Hub) Register(ws *websocket.Conn) error {
 	h.mu.Lock()
+	defer h.mu.Unlock()
+	if len(h.clients) >= maxWsClients {
+		return errors.New("max WebSocket clients reached")
+	}
 	h.clients[ws] = true
-	h.mu.Unlock()
+	return nil
 }
 
 func (h *Hub) Unregister(ws *websocket.Conn) {
@@ -56,11 +74,56 @@ func (h *Hub) ClientCount() int {
 }
 
 func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
+	if s.hub.ClientCount() >= maxWsClients {
+		http.Error(w, "too many connections", http.StatusServiceUnavailable)
+		return
+	}
+
+	if s.apiToken != "" {
+		token := r.URL.Query().Get("token")
+		if token == "" {
+			auth := r.Header.Get("Authorization")
+			if strings.HasPrefix(auth, "Bearer ") {
+				token = strings.TrimPrefix(auth, "Bearer ")
+			}
+		}
+		if token != s.apiToken {
+			http.Error(w, "unauthorized", http.StatusUnauthorized)
+			return
+		}
+	}
+
 	handler := websocket.Handler(func(ws *websocket.Conn) {
-		s.hub.Register(ws)
+		if err := s.hub.Register(ws); err != nil {
+			return
+		}
 		defer s.hub.Unregister(ws)
 
-		// Keep connection alive until client disconnects.
+		// Ping goroutine: detect dead connections by sending periodic pings.
+		// A write failure (including write deadline exceeded) closes the conn,
+		// causing the read loop below to exit and unregister the client.
+		done := make(chan struct{})
+		defer close(done)
+		go func() {
+			ticker := time.NewTicker(wsPingInterval)
+			defer ticker.Stop()
+			for {
+				select {
+				case <-done:
+					return
+				case <-ticker.C:
+					ws.SetWriteDeadline(time.Now().Add(wsPingDeadline))
+					err := websocket.Message.Send(ws, "ping")
+					ws.SetWriteDeadline(time.Time{})
+					if err != nil {
+						ws.Close()
+						return
+					}
+				}
+			}
+		}()
+
+		// Keep connection alive until client disconnects or ping fails.
 		buf := make([]byte, 1024)
 		for {
 			if _, err := ws.Read(buf); err != nil {