job-search/internal/websocket/client.go

package websocket

import (
	"context"
	"crypto/rand"
	"encoding/json"
	"errors"
	"log/slog"
	"sync"
	"sync/atomic"
	"time"

	"github.com/coder/websocket"
	"golang.org/x/time/rate"
)

// clientSendBuffer is how many messages are queued for a client before it
// counts as a slow consumer and is closed.
const clientSendBuffer = 256

// logMessageLimit is how much of a message is logged; longer messages are
// cut short and logged as truncated.
const logMessageLimit = 1 << 10

// msgConn is the connection a client reads and writes messages on: a
// *websocket.Conn, or a tcpConn for plain TCP clients (see tcp.go).
type msgConn interface {
	Read(ctx context.Context) (websocket.MessageType, []byte, error)
	Write(ctx context.Context, messageType websocket.MessageType, payload []byte) error
	Ping(ctx context.Context) error
	Close(code websocket.StatusCode, reason string) error
	CloseNow() error
}

// client is one connection registered with the hub. Its read loop handles
// incoming envelopes, its write loop sends the queued messages and its ping
// loop checks an idle connection is still alive.
type client struct {
	hub       *Hub
	conn      msgConn
	send      chan []byte // queued messages, sent by writeLoop
	id        string
	ip        string
	lastSeen  atomic.Int64 // unix nanoseconds; set by the read and ping loops
	cancel    context.CancelFunc
	limiter   *rate.Limiter // limits the messages the client sends
	closeOnce sync.Once
}

// errUnsupportedData is returned by readLoop (and a tcpConn's Write) for a
// binary message.
var errUnsupportedData = errors.New("binary frames are not supported")

// newClient creates a client for conn. cancel cancels the context the client
// runs with; the client calls it when it finishes or is closed. The client's
// messages are limited by websocket.message_rate_limit if it is enabled.
func newClient(
	h *Hub, conn msgConn, ip string, cancel context.CancelFunc,
) *client {
	limiter := rate.NewLimiter(rate.Inf, 0)
	if rule := h.conf.MessageRateLimit; rule.Enabled && rule.PerSecond > 0 {
		limiter = rate.NewLimiter(
			rate.Limit(rule.PerSecond),
			max(rule.Burst, 1),
		)
	}

	c := &client{
		hub:     h,
		conn:    conn,
		id:      rand.Text()[:12],
		ip:      ip,
		send:    make(chan []byte, clientSendBuffer),
		cancel:  cancel,
		limiter: limiter,
	}
	c.lastSeen.Store(time.Now().UnixNano())

	return c
}

// run starts the writer and keepalive goroutines and reads until the
// connection ends. It returns once the connection is finished.
func (c *client) run(ctx context.Context) {
	var wg sync.WaitGroup
	wg.Add(2)
	go func() { defer wg.Done(); c.writeLoop(ctx) }()
	go func() { defer wg.Done(); c.pingLoop(ctx) }()

	c.sendEnvelope(Envelope{
		Type: TypeWelcome,
		Payload: mustRaw(map[string]any{
			"clientId": c.id,
		}),
	})

	err := c.readLoop(ctx)
	switch {
	case err == nil:
	case websocket.CloseStatus(err) != -1:
		// The peer closed the connection, or close() closed it.
	case errors.Is(err, errUnsupportedData):
		c.close(websocket.StatusUnsupportedData, "text frames only")
	case errors.Is(err, context.Canceled):
		// The client was cancelled: by close() (shutdown or a slow
		// consumer), after a failed write or ping, or by the caller.
		c.conn.CloseNow() //nolint:errcheck,gosec // finished with; it may already be closed
	default:
		c.hub.logger().Debug("read error", "client", c.id, "error", err)
		c.conn.CloseNow() //nolint:errcheck,gosec // the connection has failed anyway
	}

	c.cancel()
	wg.Wait()
}

// readLoop reads messages until the connection ends, and returns why.
// Messages over the rate limit or that aren't valid JSON envelopes get an
// error reply; the rest are passed to handle.
func (c *client) readLoop(ctx context.Context) error {
	for {
		messageType, data, err := c.conn.Read(ctx)
		if err != nil {
			return err
		}
		c.logMessage(ctx, "message received", data)
		c.lastSeen.Store(time.Now().UnixNano())
		if messageType != websocket.MessageText {
			return errUnsupportedData
		}
		// Check if the client is rate limited for early discard
		if !c.limiter.Allow() {
			c.sendError("", "rate limited")
			continue
		}
		// Unpack the envelope and pass to handle to process if valid
		var env Envelope
		if err := json.Unmarshal(data, &env); err != nil {
			c.sendError("", "invalid JSON envelope")
			continue
		}
		c.handle(env)
	}
}

// handle answers one envelope from the client: ping and echo, or an error
// for any other type.
func (c *client) handle(env Envelope) {
	switch env.Type {
	case TypePing:
		c.sendEnvelope(Envelope{
			Type:    TypePong,
			ID:      env.ID,
			Payload: mustRaw(map[string]any{"time": time.Now().UTC()}),
		})

	case TypeEcho:
		c.sendEnvelope(Envelope{
			Type:    TypeEcho,
			ID:      env.ID,
			Payload: env.Payload,
		})

	default:
		c.sendError(env.ID, "unknown message type")
	}
}

// writeLoop sends queued messages, each within the write timeout (no limit if
// it is 0 or less), until ctx is cancelled. A failed write cancels the client.
func (c *client) writeLoop(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case msg := <-c.send:
			writeCtx, cancel := withTimeout(ctx, c.hub.conf.WriteTimeout)
			err := c.conn.Write(writeCtx, websocket.MessageText, msg)
			cancel()
			if err != nil {
				c.hub.logger().Debug("write failed", "client", c.id, "error", err)
				c.cancel()
				return
			}
			c.logMessage(ctx, "message sent", msg)
		}
	}
}

// pingLoop pings the connection when it has been idle for the ping interval;
// no pong within the pong timeout (no limit if it is 0 or less) drops the
// connection. Ping needs the concurrent reader in readLoop to see pongs.
func (c *client) pingLoop(ctx context.Context) {
	pingInterval := c.hub.conf.PingInterval
	pongTimeout := c.hub.conf.PongTimeout

	if pingInterval <= 0 {
		return
	}

	ticker := time.NewTicker(pingInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			// Don't ping if the lastSeen is more recent than the ping interval
			lastSeen := time.Unix(0, c.lastSeen.Load())
			if time.Now().Before(lastSeen.Add(pingInterval)) {
				continue
			}
			pingCtx, cancel := withTimeout(ctx, pongTimeout)
			err := c.conn.Ping(pingCtx)
			cancel()
			if err != nil {
				c.hub.logger().Debug("ping failed", "client", c.id, "error", err)
				c.conn.CloseNow() //nolint:errcheck,gosec // dropping the client anyway
				c.cancel()
				return
			}
			c.lastSeen.Store(time.Now().UnixNano())
		}
	}
}

// trySend queues msg without blocking; false means the buffer is full.
func (c *client) trySend(msg []byte) bool {
	select {
	case c.send <- msg:
		return true
	default:
		return false
	}
}

// sendEnvelope queues env for the client, closing the client as a slow
// consumer if its queue is full.
func (c *client) sendEnvelope(env Envelope) {
	encoded, err := json.Marshal(env)
	if err != nil {
		c.hub.logger().Error("marshal", "client", c.id, "error", err)
		return
	}
	if !c.trySend(encoded) {
		c.hub.logger().Warn("dropping slow client", "client", c.id, "ip", c.ip)
		go c.close(websocket.StatusPolicyViolation, "slow consumer")
	}
}

// sendError queues an error envelope replying to the message with id (none
// if empty).
func (c *client) sendError(id, msg string) {
	c.sendEnvelope(Envelope{
		Type:  TypeError,
		ID:    id,
		Error: msg,
	})
}

// close performs a graceful close handshake once, then cancels the client.
// It may block for the handshake, so call it from its own goroutine unless
// the read loop has already exited.
func (c *client) close(code websocket.StatusCode, reason string) {
	c.closeOnce.Do(func() {
		c.conn.Close(code, reason) //nolint:errcheck,gosec // the client is cancelled anyway
		c.cancel()
	})
}

// logMessage logs a message sent to or received from the client at debug
// level, with its size; at most logMessageLimit bytes of it are logged. It
// does nothing unless debug logging is enabled.
func (c *client) logMessage(ctx context.Context, event string, data []byte) {
	logger := c.hub.logger()
	if !logger.Enabled(ctx, slog.LevelDebug) {
		return
	}
	logged, truncated := data, false
	if len(logged) > logMessageLimit {
		logged, truncated = logged[:logMessageLimit], true
	}
	logger.LogAttrs(
		ctx,
		slog.LevelDebug,
		event,
		slog.String("client", c.id),
		slog.Int("bytes", len(data)),
		slog.String("message", string(logged)),
		slog.Bool("truncated", truncated),
	)
}

// withTimeout returns ctx with the timeout applied, or ctx itself if the
// timeout is 0 or less (no limit).
func withTimeout(
	ctx context.Context, timeout time.Duration,
) (context.Context, context.CancelFunc) {
	if timeout <= 0 {
		return ctx, func() {}
	}
	return context.WithTimeout(ctx, timeout)
}