Initial commit

1 files changed, 293 insertions, 0 deletions

diff --git a/api/ws/ws.go b/api/ws/ws.go
new file mode 100644
index 0000000..63706db
--- /dev/null
+++ b/api/ws/ws.go
@@ -0,0 +1,293 @@
+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)
+		}
+	}
+}