dumpsterChat/internal/gateway/client.go

package gateway

import (
	"context"
	"database/sql"
	"encoding/json"
	"log/slog"
	"net/http"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

const (
	writeWait      = 10 * time.Second
	pongWait       = 60 * time.Second
	pingPeriod     = (pongWait * 9) / 10
	maxMessageSize = 4096
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

// SetAllowedOrigins replaces the default upgrader with one that validates
// the Origin header. Accepts if:
// 1. Exact match in origins list, OR
// 2. Hostname matches a trusted hostname, OR
// 3. Hostname matches the request's own Host header (same-origin), OR
// 4. Localhost variant
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,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			origin := r.Header.Get("Origin")
			if origin == "" {
				return true
			}
			if _, ok := originSet[origin]; ok {
				return true
			}
			oh := extractHostname(origin)

			// Match trusted hostnames
			if _, ok := hostSet[oh]; ok {
				return true
			}

			// Match request's own Host header (same-origin)
			reqHost := extractHostname(r.Host)
			if reqHost != "" && oh == reqHost {
				return true
			}

			// Localhost variants
			if oh == "localhost" || oh == "127.0.0.1" || oh == "::1" {
				return true
			}

			return false
		},
	}
}

func extractHostname(s string) string {
	h := s
	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]
	}
	h = strings.TrimPrefix(h, "[")
	h = strings.TrimSuffix(h, "]")
	return h
}

// Client is a middleman between the websocket connection and the hub.
type Client struct {
	Hub    *Hub
	Conn   *websocket.Conn
	UserID string
	send   chan []byte
}

// readPump pumps messages from the websocket connection to the hub.
func (c *Client) readPump() {
	defer func() {
		c.Hub.Unregister(c)
		c.Conn.Close()
	}()

	c.Conn.SetReadLimit(maxMessageSize)
	c.Conn.SetReadDeadline(time.Now().Add(pongWait))
	c.Conn.SetPongHandler(func(string) error {
		c.Conn.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})

	for {
		_, raw, err := c.Conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				c.Hub.logger.Warn("websocket read error", "error", err, "user_id", c.UserID)
			}
			break
		}

		var event Event
		if err := json.Unmarshal(raw, &event); err != nil {
			c.Hub.logger.Warn("invalid event JSON", "error", err, "user_id", c.UserID)
			continue
		}

		c.Hub.RecordActivity(c.UserID)

		switch event.Type {
		case EventTypingStart, EventPresenceUpdate:
			c.Hub.BroadcastEvent(event)
		default:
			c.Hub.logger.Info("received event from client", "type", event.Type, "user_id", c.UserID)
		}
	}
}

// writePump pumps messages from the hub to the websocket connection.
func (c *Client) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.Conn.Close()
	}()

	for {
		select {
		case message, ok := <-c.send:
			c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
			if !ok {
				c.Conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			w, err := c.Conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			w.Write(message)

			n := len(c.send)
			for i := 0; i < n; i++ {
				w.Write([]byte{'\n'})
				w.Write(<-c.send)
			}

			if err := w.Close(); err != nil {
				return
			}

		case <-ticker.C:
			c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

// 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 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 {
		logger.Error("websocket upgrade failed", "error", err)
		return
	}

	// Try cookie-based auth first (browser clients).
	var token string
	if cookie, cookieErr := r.Cookie(cookieName); cookieErr == nil && cookie.Value != "" {
		token = cookie.Value
	}

	// Fall back to message-frame auth (TUI, bots).
	if token == "" {
		conn.SetReadDeadline(time.Now().Add(10 * time.Second))
		_, raw, readErr := conn.ReadMessage()
		if readErr != nil {
			logger.Warn("ws auth: failed to read auth message", "error", readErr)
			conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"auth timeout"}`))
			conn.Close()
			return
		}
		var auth authMessage
		if jsonErr := json.Unmarshal(raw, &auth); jsonErr != nil || auth.Token == "" {
			logger.Warn("ws auth: invalid auth message")
			conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"missing token"}`))
			conn.Close()
			return
		}
		token = auth.Token
	}

	var userID string
	err = db.QueryRowContext(context.Background(),
		`SELECT user_id FROM sessions WHERE token = $1 AND expires_at > NOW()`,
		token,
	).Scan(&userID)
	if err != nil {
		logger.Warn("ws auth: invalid session", "error", err)
		conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"invalid session"}`))
		conn.Close()
		return
	}

	conn.SetReadDeadline(time.Time{})
	conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"ready"}`))

	client := &Client{
		Hub:    hub,
		Conn:   conn,
		UserID: userID,
		send:   make(chan []byte, 256),
	}

	hub.Register(client)

	go client.writePump()
	go client.readPump()
}