Added websocket hub for handling websocket clients, with a wsserve command
which runs a plain TCP server for receiving websocket hub clients, with server config, CSRF and rate limit middleware. - internal/websocket: Hub pushing messages to WebSocket clients (ServeHTTP, example.com/ws upgrade with optional CSRF check) and plain TCP clients (ServeTCP, newline-delimited JSON envelopes); ping/pong keepalive, per-client message rate limit, slow consumer handling, graceful Shutdown, debug logging of messages sent and received - internal/websocket/tcpserver: standalone TCP server for a hub, with accept back-off, a per-IP connection limit (server.websocket.connection_limit) and graceful shutdown - internal/cli/wsserve: `wsserve` command running the TCP server, with --listen/--host/--port overrides - internal/middleware: keyed token bucket rate limiter, real client IP from trusted proxies - internal/httpx: RFC 9457 problem responses, request IDs in context - internal/config: server section (host, port, shutdown timeout, rate limiter idle TTL, websocket); util.ByteSize settings (e.g. "32KiB"); lists of sections as single keys; empty lists stored as nil; isExplicitlySet for printLoadedConfig; JobSearchConfigSource renamed to JSCSourcesSource; TestSections checks every section's methods - internal/util: ByteSize, parsed and formatted with go-humanize
This commit is contained in:
bwilliams committed
1 parent
803a8929e6
commit
c9f42b62ad
42 files changed
+3619
-57
No files matched your search
@@ -0,0 +1,300 @@
|
||||
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.log.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.log.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.log.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.log.Error("marshal", "client", c.id, "error", err)
|
||||
return
|
||||
}
|
||||
if !c.trySend(encoded) {
|
||||
c.hub.log.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) {
|
||||
if !c.hub.log.Enabled(ctx, slog.LevelDebug) {
|
||||
return
|
||||
}
|
||||
logged, truncated := data, false
|
||||
if len(logged) > logMessageLimit {
|
||||
logged, truncated = logged[:logMessageLimit], true
|
||||
}
|
||||
c.hub.log.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)
|
||||
}
|
||||
Reference in new issue
Block a user