blumenjs 0.2.0 → 0.2.2

package/dist/templates/go-server/websocket.go

@@ -0,0 +1,430 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+// ─── WebSocket Hub ─────────────────────────────────────────────
+// Manages all WebSocket connections with room/channel support.
+// Each connection runs in its own goroutine — Go handles 100K+
+// concurrent connections with minimal memory overhead.
+
+// WSMessage represents a message sent over WebSocket
+type WSMessage struct {
+	Type    string          `json:"type"`              // Message type (e.g. "chat", "notification", "ping")
+	Room    string          `json:"room,omitempty"`    // Target room (empty = broadcast to all)
+	Payload json.RawMessage `json:"payload,omitempty"` // Arbitrary JSON payload
+	Sender  string          `json:"sender,omitempty"`  // Connection ID of the sender
+}
+
+// WSClient represents a single WebSocket connection
+type WSClient struct {
+	hub  *WSHub
+	conn *websocket.Conn
+	send chan []byte
+	id   string
+	rooms map[string]bool
+	mu    sync.RWMutex
+}
+
+// WSHub manages all active WebSocket clients and rooms
+type WSHub struct {
+	clients    map[*WSClient]bool
+	rooms      map[string]map[*WSClient]bool
+	broadcast  chan []byte
+	register   chan *WSClient
+	unregister chan *WSClient
+	mu         sync.RWMutex
+
+	// OnConnect is called when a client connects. Return false to reject.
+	OnConnect func(client *WSClient, r *http.Request) bool
+
+	// OnMessage is called when a message is received from a client.
+	// Return a response message or nil to not respond.
+	OnMessage func(client *WSClient, msg WSMessage) *WSMessage
+
+	// OnDisconnect is called when a client disconnects.
+	OnDisconnect func(client *WSClient)
+}
+
+// NewWSHub creates a new WebSocket hub
+func NewWSHub() *WSHub {
+	return &WSHub{
+		clients:    make(map[*WSClient]bool),
+		rooms:      make(map[string]map[*WSClient]bool),
+		broadcast:  make(chan []byte, 256),
+		register:   make(chan *WSClient),
+		unregister: make(chan *WSClient),
+	}
+}
+
+// Run starts the hub's event loop. Call as a goroutine.
+func (h *WSHub) Run() {
+	for {
+		select {
+		case client := <-h.register:
+			h.mu.Lock()
+			h.clients[client] = true
+			h.mu.Unlock()
+			log.Printf("🔌 WebSocket client connected: %s (total: %d)", client.id, len(h.clients))
+
+		case client := <-h.unregister:
+			h.mu.Lock()
+			if _, ok := h.clients[client]; ok {
+				delete(h.clients, client)
+				close(client.send)
+
+				// Remove from all rooms
+				client.mu.RLock()
+				for room := range client.rooms {
+					if clients, ok := h.rooms[room]; ok {
+						delete(clients, client)
+						if len(clients) == 0 {
+							delete(h.rooms, room)
+						}
+					}
+				}
+				client.mu.RUnlock()
+			}
+			h.mu.Unlock()
+			if h.OnDisconnect != nil {
+				h.OnDisconnect(client)
+			}
+			log.Printf("🔌 WebSocket client disconnected: %s (total: %d)", client.id, len(h.clients))
+
+		case message := <-h.broadcast:
+			h.mu.RLock()
+			for client := range h.clients {
+				select {
+				case client.send <- message:
+				default:
+					// Client buffer full — disconnect
+					close(client.send)
+					delete(h.clients, client)
+				}
+			}
+			h.mu.RUnlock()
+		}
+	}
+}
+
+// BroadcastJSON sends a WSMessage to all connected clients
+func (h *WSHub) BroadcastJSON(msg WSMessage) {
+	data, err := json.Marshal(msg)
+	if err != nil {
+		log.Printf("WebSocket broadcast marshal error: %v", err)
+		return
+	}
+	h.broadcast <- data
+}
+
+// SendToRoom sends a message to all clients in a specific room
+func (h *WSHub) SendToRoom(room string, msg WSMessage) {
+	data, err := json.Marshal(msg)
+	if err != nil {
+		return
+	}
+
+	h.mu.RLock()
+	clients, ok := h.rooms[room]
+	h.mu.RUnlock()
+
+	if !ok {
+		return
+	}
+
+	for client := range clients {
+		select {
+		case client.send <- data:
+		default:
+			close(client.send)
+			h.mu.Lock()
+			delete(h.clients, client)
+			delete(clients, client)
+			h.mu.Unlock()
+		}
+	}
+}
+
+// SendToClient sends a message to a specific client by ID
+func (h *WSHub) SendToClient(clientID string, msg WSMessage) {
+	data, err := json.Marshal(msg)
+	if err != nil {
+		return
+	}
+
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	for client := range h.clients {
+		if client.id == clientID {
+			select {
+			case client.send <- data:
+			default:
+			}
+			return
+		}
+	}
+}
+
+// ClientCount returns the number of connected clients
+func (h *WSHub) ClientCount() int {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	return len(h.clients)
+}
+
+// RoomCount returns the number of clients in a specific room
+func (h *WSHub) RoomCount(room string) int {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	if clients, ok := h.rooms[room]; ok {
+		return len(clients)
+	}
+	return 0
+}
+
+// ─── Client Methods ────────────────────────────────────────────
+
+// JoinRoom adds the client to a room
+func (c *WSClient) JoinRoom(room string) {
+	c.mu.Lock()
+	c.rooms[room] = true
+	c.mu.Unlock()
+
+	c.hub.mu.Lock()
+	if _, ok := c.hub.rooms[room]; !ok {
+		c.hub.rooms[room] = make(map[*WSClient]bool)
+	}
+	c.hub.rooms[room][c] = true
+	c.hub.mu.Unlock()
+
+	log.Printf("🔌 Client %s joined room: %s", c.id, room)
+}
+
+// LeaveRoom removes the client from a room
+func (c *WSClient) LeaveRoom(room string) {
+	c.mu.Lock()
+	delete(c.rooms, room)
+	c.mu.Unlock()
+
+	c.hub.mu.Lock()
+	if clients, ok := c.hub.rooms[room]; ok {
+		delete(clients, c)
+		if len(clients) == 0 {
+			delete(c.hub.rooms, room)
+		}
+	}
+	c.hub.mu.Unlock()
+
+	log.Printf("🔌 Client %s left room: %s", c.id, room)
+}
+
+// Send sends a message directly to this client
+func (c *WSClient) Send(msg WSMessage) {
+	data, err := json.Marshal(msg)
+	if err != nil {
+		return
+	}
+	select {
+	case c.send <- data:
+	default:
+	}
+}
+
+// ID returns the client's unique identifier
+func (c *WSClient) ID() string {
+	return c.id
+}
+
+// ─── WebSocket Upgrader ────────────────────────────────────────
+
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+	// Allow all origins in development. In production, configure this.
+	CheckOrigin: func(r *http.Request) bool {
+		return true
+	},
+}
+
+// WebSocket read/write configuration
+const (
+	writeWait      = 10 * time.Second
+	pongWait       = 60 * time.Second
+	pingPeriod     = (pongWait * 9) / 10
+	maxMessageSize = 65536 // 64KB
+)
+
+// readPump pumps messages from the WebSocket connection to the hub.
+func (c *WSClient) 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 {
+		_, rawMsg, err := c.conn.ReadMessage()
+		if err != nil {
+			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
+				log.Printf("WebSocket read error: %v", err)
+			}
+			break
+		}
+
+		// Parse the incoming message
+		var msg WSMessage
+		if err := json.Unmarshal(rawMsg, &msg); err != nil {
+			log.Printf("WebSocket message parse error: %v", err)
+			continue
+		}
+
+		msg.Sender = c.id
+
+		// Handle built-in message types
+		switch msg.Type {
+		case "join":
+			var room string
+			json.Unmarshal(msg.Payload, &room)
+			if room != "" {
+				c.JoinRoom(room)
+			}
+			continue
+
+		case "leave":
+			var room string
+			json.Unmarshal(msg.Payload, &room)
+			if room != "" {
+				c.LeaveRoom(room)
+			}
+			continue
+
+		case "ping":
+			c.Send(WSMessage{Type: "pong"})
+			continue
+		}
+
+		// Call the user-defined message handler
+		if c.hub.OnMessage != nil {
+			if response := c.hub.OnMessage(c, msg); response != nil {
+				c.Send(*response)
+			}
+		}
+
+		// If the message targets a room, broadcast to that room
+		if msg.Room != "" {
+			c.hub.SendToRoom(msg.Room, msg)
+		}
+	}
+}
+
+// writePump pumps messages from the hub to the WebSocket connection.
+func (c *WSClient) 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 {
+				// Hub closed the channel
+				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
+				return
+			}
+
+			w, err := c.conn.NextWriter(websocket.TextMessage)
+			if err != nil {
+				return
+			}
+			w.Write(message)
+
+			// Drain queued messages into the current write
+			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
+			}
+		}
+	}
+}
+
+// ─── HTTP Handler ──────────────────────────────────────────────
+
+var clientIDCounter uint64
+var clientIDMu sync.Mutex
+
+func nextClientID() string {
+	clientIDMu.Lock()
+	clientIDCounter++
+	id := clientIDCounter
+	clientIDMu.Unlock()
+	return log.Prefix() + "ws-" + time.Now().Format("150405") + "-" + fmt.Sprintf("%d", id)
+}
+
+// WSHandler returns an HTTP handler that upgrades to WebSocket
+func WSHandler(hub *WSHub) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			log.Printf("WebSocket upgrade error: %v", err)
+			return
+		}
+
+		client := &WSClient{
+			hub:   hub,
+			conn:  conn,
+			send:  make(chan []byte, 256),
+			id:    nextClientID(),
+			rooms: make(map[string]bool),
+		}
+
+		// Check if the connection should be accepted
+		if hub.OnConnect != nil && !hub.OnConnect(client, r) {
+			conn.WriteMessage(websocket.CloseMessage,
+				websocket.FormatCloseMessage(websocket.ClosePolicyViolation, "connection rejected"))
+			conn.Close()
+			return
+		}
+
+		hub.register <- client
+
+		// Send the client its ID
+		welcome := WSMessage{
+			Type:    "connected",
+			Payload: json.RawMessage(`"` + client.id + `"`),
+		}
+		client.Send(welcome)
+
+		// Start read/write pumps in goroutines
+		go client.writePump()
+		go client.readPump()
+	}
+}