package pusher import ( "context" "encoding/json" "log/slog" "sync" "nhooyr.io/websocket" ) const sendBufferSize = 256 // Conn wraps a WebSocket connection with identity and a send buffer. type Conn struct { id string humanID string ws *websocket.Conn send chan []byte once sync.Once // ensures close logic runs once } func newConn(id, humanID string, ws *websocket.Conn) *Conn { return &Conn{ id: id, humanID: humanID, ws: ws, send: make(chan []byte, sendBufferSize), } } // ReadPump reads messages from the WebSocket and forwards them to the hub. // It blocks until the connection is closed or the context is cancelled. func (c *Conn) ReadPump(ctx context.Context, hub *Hub) { defer hub.disconnect(c) for { _, data, err := c.ws.Read(ctx) if err != nil { if ctx.Err() != nil { slog.Info("websocket context cancelled", "connId", c.id, "humanId", c.humanID, "error", ctx.Err()) return } slog.Warn("websocket read error", "connId", c.id, "humanId", c.humanID, "error", err) return } // Respond to keep-alive pings if string(data) == "ping" { if err := c.ws.Write(ctx, websocket.MessageText, []byte("pong")); err != nil { slog.Warn("websocket pong write error", "connId", c.id, "error", err) } continue } var msg ClientMessage if err := json.Unmarshal(data, &msg); err != nil { c.sendError("invalid message format") continue } switch msg.Type { case TypeSubscribe: if msg.Channel == "" { c.sendError("channel is required") continue } hub.subscribeCh <- &subscribeRequest{conn: c, channelID: msg.Channel} case TypeUnsubscribe: if msg.Channel == "" { c.sendError("channel is required") continue } hub.unsubscribeCh <- &unsubscribeRequest{conn: c, channelID: msg.Channel} case TypeMessage: if msg.Channel == "" { c.sendError("channel is required") continue } hub.broadcastCh <- &broadcastRequest{conn: c, channelID: msg.Channel, payload: msg.Payload} default: c.sendError("unknown message type: " + msg.Type) } } } // WritePump drains the send buffer and writes messages to the WebSocket. func (c *Conn) WritePump(ctx context.Context) { for { select { case <-ctx.Done(): return case data, ok := <-c.send: if !ok { return } if err := c.ws.Write(ctx, websocket.MessageText, data); err != nil { slog.Debug("websocket write error", "connId", c.id, "error", err) return } } } } // Send enqueues a ServerMessage to be written to the WebSocket. // If the send buffer is full, the connection is closed (slow client). func (c *Conn) Send(msg ServerMessage) { data, err := json.Marshal(msg) if err != nil { slog.Error("failed to marshal server message", "error", err) return } select { case c.send <- data: default: slog.Warn("slow client, closing connection", "connId", c.id, "humanId", c.humanID) c.Close() } } // Close closes the WebSocket connection and the send channel. func (c *Conn) Close() { c.once.Do(func() { c.ws.Close(websocket.StatusNormalClosure, "closing") close(c.send) }) } func (c *Conn) sendError(msg string) { c.Send(ServerMessage{Type: TypeError, Message: msg}) }