Messenger/server/internal/realtime/client.go

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
}