Changeset 0.28.8 (#194)

server/events/ws.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"log"
 	"net/http"
+	"strings"
 	"sync"
 	"time"
 
@@ -11,12 +12,6 @@ import (
 	"github.com/gorilla/websocket"
 )
 
-var upgrader = websocket.Upgrader{
-	ReadBufferSize:  1024,
-	WriteBufferSize: 1024,
-	CheckOrigin:     func(r *http.Request) bool { return true },
-}
-
 const (
 	writeWait  = 10 * time.Second
 	pongWait   = 60 * time.Second
@@ -26,9 +21,10 @@ const (
 
 // Hub manages WebSocket connections and bridges them to the Bus.
 type Hub struct {
-	bus   *Bus
-	mu    sync.RWMutex
-	conns map[string]*Conn // connID → Conn
+	bus      *Bus
+	mu       sync.RWMutex
+	conns    map[string]*Conn // connID → Conn
+	upgrader websocket.Upgrader
 }
 
 // Conn represents a single WebSocket connection.
@@ -43,11 +39,55 @@ type Conn struct {
 }
 
 // NewHub creates a Hub bound to an event bus.
-func NewHub(bus *Bus) *Hub {
-	return &Hub{
+// allowedOrigins controls the WebSocket upgrader's CheckOrigin behavior.
+// Empty or "*" allows all origins (dev/test). Comma-separated list
+// restricts to those origins (production).
+func NewHub(bus *Bus, allowedOrigins string) *Hub {
+	h := &Hub{
 		bus:   bus,
 		conns: make(map[string]*Conn),
 	}
+	h.upgrader = websocket.Upgrader{
+		ReadBufferSize:  1024,
+		WriteBufferSize: 1024,
+		CheckOrigin:     h.makeCheckOrigin(allowedOrigins),
+	}
+	return h
+}
+
+// makeCheckOrigin returns a CheckOrigin function that validates the
+// request Origin header against the configured allowed origins.
+func (h *Hub) makeCheckOrigin(allowedOrigins string) func(r *http.Request) bool {
+	origins := strings.TrimSpace(allowedOrigins)
+
+	// No restriction: dev/test or explicit "*"
+	if origins == "" || origins == "*" {
+		return func(r *http.Request) bool { return true }
+	}
+
+	// Parse comma-separated origin list
+	var allowed []string
+	for _, o := range strings.Split(origins, ",") {
+		o = strings.TrimSpace(o)
+		if o != "" {
+			allowed = append(allowed, o)
+		}
+	}
+
+	return func(r *http.Request) bool {
+		origin := r.Header.Get("Origin")
+		if origin == "" {
+			// No Origin header — same-origin request (non-browser client).
+			return true
+		}
+		for _, a := range allowed {
+			if a == origin {
+				return true
+			}
+		}
+		log.Printf("[ws] rejected WebSocket from origin %q (allowed: %s)", origin, origins)
+		return false
+	}
 }
 
 // HandleWebSocket is a Gin handler for WebSocket upgrade.
@@ -60,7 +100,7 @@ func (h *Hub) HandleWebSocket(c *gin.Context) {
 		return
 	}
 
-	ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	ws, err := h.upgrader.Upgrade(c.Writer, c.Request, nil)
 	if err != nil {
 		log.Printf("[ws] upgrade failed: %v", err)
 		return