path: root/api/ws/ws.go

package ws

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/gorilla/websocket"

	"penguins.lan/api/types"
	"penguins.lan/database"
	"penguins.lan/utils"
)

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

	historyLimit = 50
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	// It's a LAN intranet; we don't police origins.
	CheckOrigin: func(r *http.Request) bool { return true },
}

// Message is the wire format shared by client and server. `Type` discriminates
// the payload so we can grow beyond chat without breaking old clients.
type Message struct {
	Type string `json:"type"` // "chat" | "system" | "history"
	Nick string `json:"nick,omitempty"`
	Body string `json:"body,omitempty"`
	At   int64  `json:"at,omitempty"` // unix millis
	// History carries a batch of past messages on join.
	History []Message `json:"history,omitempty"`
}

// Hub fans messages out to every connected client.
type Hub struct {
	db         *sql.DB
	clients    map[*Client]bool
	broadcast  chan []byte
	register   chan *Client
	unregister chan *Client
	online     chan chan map[string]bool
}

func NewHub(db *sql.DB) *Hub {
	return &Hub{
		db:         db,
		clients:    make(map[*Client]bool),
		broadcast:  make(chan []byte),
		register:   make(chan *Client),
		unregister: make(chan *Client),
		online:     make(chan chan map[string]bool),
	}
}

// Run owns all mutations to the client set, so no locks are needed elsewhere.
func (h *Hub) Run() {
	for {
		select {
		case client := <-h.register:
			h.clients[client] = true
			h.systemf("%s waddled in 🐧", h.displayName(client.userID))
		case client := <-h.unregister:
			if _, ok := h.clients[client]; ok {
				delete(h.clients, client)
				close(client.send)
				h.systemf("%s slid away", h.displayName(client.userID))
			}
		case payload := <-h.broadcast:
			h.deliver(payload)
		case reply := <-h.online:
			set := make(map[string]bool, len(h.clients))
			for client := range h.clients {
				set[client.userID] = true
			}
			reply <- set
		}
	}
}

// Online returns the set of userIDs currently connected to the chat.
func (h *Hub) Online() map[string]bool {
	reply := make(chan map[string]bool)
	h.online <- reply
	return <-reply
}

// displayName resolves a userID to its current username. The chat hub never
// caches names, so a rename is reflected everywhere immediately.
func (h *Hub) displayName(userID string) string {
	user, err := database.GetUserByID(h.db, userID)
	if err != nil {
		log.Println("ws: could not resolve user", userID, err)
		return "a mystery penguin"
	}
	return user.Username
}

// deliver fans a payload out to every client, dropping any that can't keep up.
// Only safe to call from the Run goroutine, which owns the clients map.
func (h *Hub) deliver(payload []byte) {
	for client := range h.clients {
		select {
		case client.send <- payload:
		default:
			// slow consumer: drop them rather than block the hub
			delete(h.clients, client)
			close(client.send)
		}
	}
}

// sendJSON queues a message for broadcast from outside the hub goroutine
// (e.g. a chat message arriving on a client's readPump).
func (h *Hub) sendJSON(msg Message) {
	payload, err := json.Marshal(msg)
	if err != nil {
		log.Println("ws: marshal error", err)
		return
	}
	h.broadcast <- payload
}

// systemf emits a system notice. It's only ever called from within Run, so it
// fans out directly — re-entering the broadcast channel here would deadlock the
// hub (the loop can't both send and receive on it at once).
func (h *Hub) systemf(format string, args ...interface{}) {
	payload, err := json.Marshal(Message{
		Type: "system",
		Body: fmt.Sprintf(format, args...),
		At:   time.Now().UnixMilli(),
	})
	if err != nil {
		log.Println("ws: marshal error", err)
		return
	}
	h.deliver(payload)
}

// Client is one websocket connection. Reads happen on readPump, writes on
// writePump, communicating through the buffered `send` channel. Identity is the
// userID alone — the display name is always resolved from it, never cached.
type Client struct {
	hub    *Hub
	conn   *websocket.Conn
	send   chan []byte
	userID string
}

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.CloseAbnormalClosure) {
				log.Println("ws: read error", err)
			}
			break
		}

		var incoming Message
		if err := json.Unmarshal(raw, &incoming); err != nil {
			continue
		}

		// The server is authoritative about who you are and when it happened.
		body := utils.Sanitize(incoming.Body)
		if incoming.Type != "chat" || body == "" {
			continue
		}

		stored := &database.ChatMessage{
			ID:     utils.RandomId(),
			UserID: c.userID,
			Body:   utils.Truncate(incoming.Body, 500),
		}
		if _, err := database.SaveChatMessage(c.hub.db, stored); err != nil {
			log.Println("ws: failed to persist chat", err)
		}

		c.hub.sendJSON(Message{
			Type: "chat",
			Nick: c.hub.displayName(c.userID),
			Body: stored.Body,
			At:   time.Now().UnixMilli(),
		})
	}
}

func (c *Client) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()

	for {
		select {
		case payload, ok := <-c.send:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if !ok {
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			if err := c.conn.WriteMessage(websocket.TextMessage, payload); err != nil {
				return
			}
		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

// sendHistory replays recent messages to a freshly-connected client only.
func (c *Client) sendHistory() {
	rows, err := database.GetRecentChat(c.hub.db, historyLimit)
	if err != nil {
		log.Println("ws: failed to load history", err)
		return
	}

	history := make([]Message, 0, len(rows))
	for _, row := range rows {
		history = append(history, Message{
			Type: "chat",
			Nick: row.Nick,
			Body: row.Body,
			At:   row.CreatedAt.UnixMilli(),
		})
	}

	payload, err := json.Marshal(Message{Type: "history", History: history})
	if err != nil {
		log.Println("ws: marshal history error", err)
		return
	}
	c.send <- payload
}

// ServeWSContinuation upgrades the request to a websocket and wires the client
// into the hub. It's terminal in a route chain — after upgrade there's no HTTP
// response left to render.
func ServeWSContinuation(hub *Hub) types.Continuation {
	return func(context *types.RequestContext, req *http.Request, resp http.ResponseWriter) types.ContinuationChain {
		return func(success types.Continuation, failure types.Continuation) types.ContinuationChain {
			conn, err := upgrader.Upgrade(resp, req, nil)
			if err != nil {
				log.Println("ws: upgrade failed", err)
				return failure(context, req, resp)
			}

			client := &Client{
				hub:    hub,
				conn:   conn,
				send:   make(chan []byte, 256),
				userID: context.User.ID,
			}

			// Queue history first (send is buffered), then announce the join.
			client.sendHistory()
			client.hub.register <- client

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

			return success(context, req, resp)
		}
	}
}