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) } } }