package realtime import ( "context" "encoding/json" "time" "github.com/coder/websocket" ) const sendBufferSize = 64 // client wraps a single WebSocket connection belonging to a user. type client struct { hub *Hub conn *websocket.Conn userID int64 send chan []byte done chan struct{} } func newClient(hub *Hub, conn *websocket.Conn, userID int64) *client { return &client{ hub: hub, conn: conn, userID: userID, send: make(chan []byte, sendBufferSize), done: make(chan struct{}), } } // enqueue marshals env and queues it for delivery; drops the message if the // client cannot keep up. func (c *client) enqueue(env envelope) { payload, err := json.Marshal(env) if err != nil { return } c.enqueueRaw(payload) } func (c *client) enqueueRaw(payload []byte) { select { case <-c.done: case c.send <- payload: default: c.closeSlow() } } // closeSlow terminates a connection that cannot keep up with its send queue. func (c *client) closeSlow() { c.hub.mu.Lock() defer c.hub.mu.Unlock() if conns := c.hub.conns[c.userID]; conns != nil { if _, ok := conns[c]; ok { c.conn.Close(websocket.StatusPolicyViolation, "client too slow") } } } // writeLoop drains the send queue onto the connection. func (c *client) writeLoop() { ctx := context.Background() for { select { case <-c.done: return case msg := <-c.send: writeCtx, cancel := context.WithTimeout(ctx, 10*time.Second) err := c.conn.Write(writeCtx, websocket.MessageText, msg) cancel() if err != nil { return } } } } // readLoop processes inbound client events until the connection closes. func (c *client) readLoop() error { ctx := context.Background() for { _, data, err := c.conn.Read(ctx) if err != nil { return err } var env envelope if err := json.Unmarshal(data, &env); err != nil { c.conn.Close(websocket.StatusInvalidFramePayloadData, "bad json") return err } switch env.Type { case EventTyping: var msg clientTyping if err := json.Unmarshal(env.Data, &msg); err != nil { continue } if err := c.handleTyping(ctx, msg); err != nil { return err } case EventRead: var msg clientRead if err := json.Unmarshal(env.Data, &msg); err != nil { continue } if err := c.handleRead(ctx, msg); err != nil { return err } default: return websocket.CloseError{Code: websocket.StatusUnsupportedData, Reason: "unsupported event type"} } } } func (c *client) handleTyping(ctx context.Context, msg clientTyping) error { ok, err := c.hub.members.IsMember(ctx, msg.ConversationID, c.userID) if err != nil || !ok { return nil } c.hub.PublishTyping(msg.ConversationID, c.userID, msg.IsTyping) return nil } func (c *client) handleRead(ctx context.Context, msg clientRead) error { ok, err := c.hub.members.IsMember(ctx, msg.ConversationID, c.userID) if err != nil || !ok { return nil } if err := c.hub.members.UpdateLastRead(ctx, msg.ConversationID, c.userID); err != nil { return err } c.hub.PublishRead(msg.ConversationID, c.userID, msg.MessageID) return nil }