From 7f56571004467b94c50cff0c1eb106cbc4c4ecc7 Mon Sep 17 00:00:00 2001
From: hobokenchicken <dustin@dustin.coffee>
Date: Mon, 29 Jun 2026 14:33:39 -0400
Subject: [PATCH] fix: WS origin check now uses hostname matching instead of
 exact strings

---
 internal/gateway/client.go | 63 ++++++++++++++++++++++++++------------
 1 file changed, 44 insertions(+), 19 deletions(-)

diff --git a/internal/gateway/client.go b/internal/gateway/client.go
index 27392fb..cff3a96 100644
--- a/internal/gateway/client.go
+++ b/internal/gateway/client.go
@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"log/slog"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -21,15 +22,30 @@ const (
 var upgrader = websocket.Upgrader{
 	ReadBufferSize:  1024,
 	WriteBufferSize: 1024,
-	CheckOrigin:     func(r *http.Request) bool { return true }, // default: allow all (overridden by SetAllowedOrigins)
+	CheckOrigin:     func(r *http.Request) bool { return true },
 }
 
 // SetAllowedOrigins replaces the default upgrader with one that validates
 // the Origin header against a trusted set. Call once at startup.
+// Exact matches are checked first, then hostname-based matching is used.
 func SetAllowedOrigins(origins []string) {
 	originSet := make(map[string]struct{}, len(origins))
+	hostSet := make(map[string]struct{}, len(origins))
 	for _, o := range origins {
 		originSet[o] = struct{}{}
+		h := o
+		if idx := strings.Index(h, "://"); idx >= 0 {
+			h = h[idx+3:]
+		}
+		if idx := strings.Index(h, "/"); idx >= 0 {
+			h = h[:idx]
+		}
+		if idx := strings.LastIndex(h, ":"); idx >= 0 {
+			h = h[:idx]
+		}
+		if h != "" {
+			hostSet[h] = struct{}{}
+		}
 	}
 	upgrader = websocket.Upgrader{
 		ReadBufferSize:  1024,
@@ -37,10 +53,30 @@ func SetAllowedOrigins(origins []string) {
 		CheckOrigin: func(r *http.Request) bool {
 			origin := r.Header.Get("Origin")
 			if origin == "" {
-				return false
+				return true // non-browser clients
 			}
-			_, ok := originSet[origin]
-			return ok
+			if _, ok := originSet[origin]; ok {
+				return true
+			}
+			oh := origin
+			if idx := strings.Index(oh, "://"); idx >= 0 {
+				oh = oh[idx+3:]
+			}
+			if idx := strings.Index(oh, "/"); idx >= 0 {
+				oh = oh[:idx]
+			}
+			if idx := strings.LastIndex(oh, ":"); idx >= 0 {
+				oh = oh[:idx]
+			}
+			oh = strings.TrimPrefix(oh, "[")
+			oh = strings.TrimSuffix(oh, "]")
+			if _, ok := hostSet[oh]; ok {
+				return true
+			}
+			if oh == "localhost" || oh == "127.0.0.1" || oh == "::1" {
+				return true
+			}
+			return false
 		},
 	}
 }
@@ -82,11 +118,8 @@ func (c *Client) readPump() {
 			continue
 		}
 
-		// Record activity for idle detection.
 		c.Hub.RecordActivity(c.UserID)
 
-		// Handle client-sent events (e.g. TYPING_START, PRESENCE_UPDATE)
-		// For now, broadcast them to all clients
 		switch event.Type {
 		case EventTypingStart, EventPresenceUpdate:
 			c.Hub.BroadcastEvent(event)
@@ -119,7 +152,6 @@ func (c *Client) writePump() {
 			}
 			w.Write(message)
 
-			// Drain queued messages into the current write
 			n := len(c.send)
 			for i := 0; i < n; i++ {
 				w.Write([]byte{'\n'})
@@ -139,15 +171,13 @@ func (c *Client) writePump() {
 	}
 }
 
-// authMessage is the first frame the client must send after connecting.
+// authMessage is the first frame the client must send after connecting (for TUI/bot clients).
 type authMessage struct {
 	Token string `json:"token"`
 }
 
-// ServeWS handles websocket requests from the peer. It upgrades the
-// connection first, then waits for an auth message frame containing the
-// session token. This avoids leaking the token in the URL (which would
-// appear in server logs, browser history, and Referer headers).
+// ServeWS handles websocket requests from the peer.
+// It authenticates via cookie (browser) or message-frame (TUI/bots).
 func ServeWS(db *sql.DB, hub *Hub, logger *slog.Logger, w http.ResponseWriter, r *http.Request, cookieName string) {
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
@@ -155,9 +185,7 @@ func ServeWS(db *sql.DB, hub *Hub, logger *slog.Logger, w http.ResponseWriter, r
 		return
 	}
 
-	// Short deadline for the auth frame; we don't want unauthenticated
-	// connections hanging around consuming resources.
-	// Try cookie-based auth first (browser clients send cookies automatically).
+	// Try cookie-based auth first (browser clients).
 	var token string
 	if cookie, cookieErr := r.Cookie(cookieName); cookieErr == nil && cookie.Value != "" {
 		token = cookie.Value
@@ -195,10 +223,7 @@ func ServeWS(db *sql.DB, hub *Hub, logger *slog.Logger, w http.ResponseWriter, r
 		return
 	}
 
-	// Auth succeeded. Clear the auth deadline (readPump sets its own).
 	conn.SetReadDeadline(time.Time{})
-
-	// Send ready signal so the client knows auth passed.
 	conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"ready"}`))
 
 	client := &Client{