569 lines
15 KiB
Go
569 lines
15 KiB
Go
// Package ws provides bounded, managed WebSocket connections.
|
|
package ws
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"strconv"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
"unicode/utf8"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type MessageType int
|
|
|
|
const (
|
|
Text MessageType = websocket.TextMessage
|
|
Binary MessageType = websocket.BinaryMessage
|
|
)
|
|
|
|
type Message struct {
|
|
Type MessageType
|
|
Data []byte
|
|
}
|
|
type Handler func(context.Context, *Session, Message) error
|
|
|
|
// Handlers execute serially for each session. OnClose runs after all I/O and
|
|
// message callbacks have stopped; it must not call Wait on its own session.
|
|
type Handlers struct {
|
|
OnOpen func(context.Context, *Session) error
|
|
OnMessage Handler
|
|
OnClose func(*Session, error)
|
|
}
|
|
|
|
type SessionStats struct {
|
|
MessagesSent, MessagesReceived, BytesSent, BytesReceived uint64
|
|
QueuedMessages int
|
|
QueuedBytes int64
|
|
ReadQueuedMessages int
|
|
ReadQueuedBytes int64
|
|
}
|
|
|
|
type outgoing struct {
|
|
message Message
|
|
ctx context.Context
|
|
result chan error
|
|
}
|
|
type controlFrame struct {
|
|
kind int
|
|
data []byte
|
|
}
|
|
|
|
var nextSessionID atomic.Uint64
|
|
|
|
// Session owns exactly one reader and one writer. Sends are ordered by queue
|
|
// admission. Message callbacks never run on the reader goroutine.
|
|
type Session struct {
|
|
conn *websocket.Conn
|
|
cfg SessionConfig
|
|
handlers Handlers
|
|
id string
|
|
parent context.Context
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
done, stopped, closingCh, peerClosed, closeSent chan struct{}
|
|
startOnce, stopOnce, peerOnce sync.Once
|
|
wg sync.WaitGroup
|
|
mu sync.Mutex
|
|
closing bool
|
|
closeCode int
|
|
closeReason string
|
|
err error
|
|
stats SessionStats
|
|
out chan outgoing
|
|
in chan Message
|
|
controls chan controlFrame
|
|
}
|
|
|
|
// newSession requires a validated config. Ownership transfers before start so
|
|
// callers can register the session without racing completion.
|
|
func newSession(conn *websocket.Conn, cfg SessionConfig, handlers Handlers, parent context.Context) *Session {
|
|
if parent == nil {
|
|
parent = context.Background()
|
|
}
|
|
ctx, cancel := context.WithCancel(parent)
|
|
cfg = cfg.defaults()
|
|
return &Session{conn: conn, cfg: cfg, handlers: handlers, id: "ws-" + strconv.FormatUint(nextSessionID.Add(1), 10),
|
|
parent: parent, ctx: ctx, cancel: cancel, done: make(chan struct{}), stopped: make(chan struct{}),
|
|
closingCh: make(chan struct{}), peerClosed: make(chan struct{}), closeSent: make(chan struct{}),
|
|
out: make(chan outgoing, cfg.WriteQueueSize), in: make(chan Message, cfg.ReadQueueSize), controls: make(chan controlFrame, 8)}
|
|
}
|
|
|
|
func (s *Session) ID() string { return s.id }
|
|
func (s *Session) Context() context.Context { return s.ctx }
|
|
func (s *Session) Done() <-chan struct{} { return s.done }
|
|
func (s *Session) Err() error { s.mu.Lock(); defer s.mu.Unlock(); return s.err }
|
|
func (s *Session) Subprotocol() string { return s.conn.Subprotocol() }
|
|
func (s *Session) Stats() SessionStats { s.mu.Lock(); defer s.mu.Unlock(); return s.stats }
|
|
|
|
func (s *Session) start() {
|
|
s.startOnce.Do(func() {
|
|
s.wg.Add(4)
|
|
go func() { defer s.wg.Done(); s.readLoop() }()
|
|
go func() { defer s.wg.Done(); s.writeLoop() }()
|
|
go func() { defer s.wg.Done(); s.dispatchLoop() }()
|
|
go func() { defer s.wg.Done(); s.monitor() }()
|
|
go func() {
|
|
s.wg.Wait()
|
|
for {
|
|
select {
|
|
case item := <-s.out:
|
|
s.complete(item, ErrClosed)
|
|
default:
|
|
goto drained
|
|
}
|
|
}
|
|
drained:
|
|
// Release any discarded inbound payloads before publishing completion.
|
|
for {
|
|
select {
|
|
case msg := <-s.in:
|
|
s.releaseRead(msg)
|
|
default:
|
|
goto notified
|
|
}
|
|
}
|
|
notified:
|
|
if s.handlers.OnClose != nil {
|
|
func() {
|
|
defer func() {
|
|
if p := recover(); p != nil {
|
|
s.recordError(&OpError{Op: "on_close", Cause: fmt.Errorf("callback panic: %v", p)})
|
|
}
|
|
}()
|
|
s.handlers.OnClose(s, s.Err())
|
|
}()
|
|
}
|
|
close(s.done)
|
|
}()
|
|
})
|
|
}
|
|
|
|
// Send waits for the actual write. A cancellation before writing skips the
|
|
// message. Cancellation during writing terminates the connection; delivery may
|
|
// already have occurred and the caller must not automatically replay it.
|
|
func (s *Session) Send(ctx context.Context, msg Message) error {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
item, err := s.enqueue(ctx, msg, true)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
select {
|
|
case err := <-item.result:
|
|
return err
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
}
|
|
}
|
|
|
|
// TrySend only acknowledges queue admission, not delivery. The payload is copied.
|
|
func (s *Session) TrySend(msg Message) error {
|
|
_, err := s.enqueue(context.Background(), msg, false)
|
|
return err
|
|
}
|
|
func (s *Session) SendJSON(ctx context.Context, value any) error {
|
|
data, err := json.Marshal(value)
|
|
if err != nil {
|
|
return &OpError{Op: "encode", Cause: err}
|
|
}
|
|
return s.Send(ctx, Message{Type: Text, Data: data})
|
|
}
|
|
|
|
func (s *Session) enqueue(ctx context.Context, msg Message, wait bool) (outgoing, error) {
|
|
if err := ctx.Err(); err != nil {
|
|
return outgoing{}, err
|
|
}
|
|
if msg.Type != Text && msg.Type != Binary {
|
|
return outgoing{}, ErrInvalidMessage
|
|
}
|
|
if int64(len(msg.Data)) > s.cfg.MaxMessageBytes {
|
|
return outgoing{}, ErrMessageTooLarge
|
|
}
|
|
if msg.Type == Text && !utf8.Valid(msg.Data) {
|
|
return outgoing{}, ErrInvalidMessage
|
|
}
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if s.closing {
|
|
return outgoing{}, ErrClosed
|
|
}
|
|
if s.stats.QueuedMessages >= s.cfg.WriteQueueSize || int64(len(msg.Data)) > s.cfg.WriteQueueBytes-s.stats.QueuedBytes {
|
|
return outgoing{}, ErrQueueFull
|
|
}
|
|
item := outgoing{message: Message{Type: msg.Type, Data: append([]byte(nil), msg.Data...)}, ctx: ctx}
|
|
if wait {
|
|
item.result = make(chan error, 1)
|
|
}
|
|
select {
|
|
case s.out <- item:
|
|
s.stats.QueuedMessages++
|
|
s.stats.QueuedBytes += int64(len(msg.Data))
|
|
return item, nil
|
|
default:
|
|
return outgoing{}, ErrQueueFull
|
|
}
|
|
}
|
|
|
|
func validCloseCode(code int) bool {
|
|
switch code {
|
|
case 1000, 1001, 1002, 1003, 1007, 1008, 1009, 1010, 1011, 1012, 1013, 1014:
|
|
return true
|
|
}
|
|
return code >= 3000 && code <= 4999
|
|
}
|
|
|
|
// Close stops admission and begins a bounded graceful close, draining already
|
|
// admitted sends. It is idempotent and never waits for a callback to finish.
|
|
func (s *Session) Close(code int, reason string) error {
|
|
if !validCloseCode(code) || len(reason) > 123 || !utf8.ValidString(reason) {
|
|
return ErrInvalidMessage
|
|
}
|
|
s.beginClose(code, reason, nil)
|
|
return nil
|
|
}
|
|
func (s *Session) beginClose(code int, reason string, cause error) {
|
|
s.mu.Lock()
|
|
if cause != nil && s.err == nil {
|
|
s.err = cause
|
|
}
|
|
if !s.closing {
|
|
s.closing = true
|
|
s.closeCode = code
|
|
s.closeReason = reason
|
|
close(s.closingCh)
|
|
}
|
|
s.mu.Unlock()
|
|
s.cancel()
|
|
}
|
|
func (s *Session) recordError(err error) {
|
|
s.mu.Lock()
|
|
if s.err == nil {
|
|
s.err = err
|
|
}
|
|
s.mu.Unlock()
|
|
}
|
|
func (s *Session) terminate(err error) {
|
|
s.beginClose(1001, "", err)
|
|
s.stopOnce.Do(func() { close(s.stopped); _ = s.conn.Close() })
|
|
}
|
|
|
|
// Abort immediately closes the socket; Wait still waits for callbacks to return.
|
|
func (s *Session) Abort() { s.terminate(ErrClosed) }
|
|
func (s *Session) Wait(ctx context.Context) error {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
select {
|
|
case <-s.done:
|
|
return s.Err()
|
|
default:
|
|
}
|
|
select {
|
|
case <-s.done:
|
|
return s.Err()
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
}
|
|
}
|
|
|
|
func (s *Session) monitor() {
|
|
select {
|
|
case <-s.stopped:
|
|
return
|
|
case <-s.parent.Done():
|
|
s.beginClose(1001, "", s.parent.Err())
|
|
case <-s.closingCh:
|
|
}
|
|
timer := time.NewTimer(s.cfg.CloseTimeout)
|
|
defer timer.Stop()
|
|
select {
|
|
case <-s.stopped:
|
|
return
|
|
case <-timer.C:
|
|
s.terminate(&OpError{Op: "close", Cause: context.DeadlineExceeded})
|
|
return
|
|
case <-s.closeSent:
|
|
}
|
|
select {
|
|
case <-s.stopped:
|
|
return
|
|
case <-timer.C:
|
|
s.terminate(&OpError{Op: "close", Cause: context.DeadlineExceeded})
|
|
case <-s.peerClosed:
|
|
s.terminate(nil)
|
|
}
|
|
}
|
|
|
|
func (s *Session) readLoop() {
|
|
s.conn.SetReadLimit(s.cfg.MaxMessageBytes)
|
|
if !s.cfg.DisableHeartbeat {
|
|
_ = s.conn.SetReadDeadline(time.Now().Add(s.cfg.PongTimeout))
|
|
s.conn.SetPongHandler(func(string) error { return s.conn.SetReadDeadline(time.Now().Add(s.cfg.PongTimeout)) })
|
|
}
|
|
s.conn.SetPingHandler(func(data string) error {
|
|
select {
|
|
case s.controls <- controlFrame{websocket.PongMessage, []byte(data)}:
|
|
return nil
|
|
default:
|
|
return ErrQueueFull
|
|
}
|
|
})
|
|
s.conn.SetCloseHandler(func(code int, _ string) error {
|
|
var cause error
|
|
s.mu.Lock()
|
|
locallyClosing := s.closing
|
|
s.mu.Unlock()
|
|
// A remote going-away frame can signal a restarting service. Preserve
|
|
// its code for Client.Run's retry policy; our own close acknowledgement
|
|
// remains a normal local shutdown.
|
|
if code != 1000 && code != 1005 && (code != 1001 || !locallyClosing) {
|
|
cause = &CloseError{Code: code}
|
|
}
|
|
reply := code
|
|
if !validCloseCode(reply) {
|
|
reply = 1000
|
|
}
|
|
s.beginClose(reply, "", cause)
|
|
s.peerOnce.Do(func() { close(s.peerClosed) })
|
|
return &websocket.CloseError{Code: code}
|
|
})
|
|
for {
|
|
kind, reader, err := s.conn.NextReader()
|
|
if err != nil {
|
|
select {
|
|
case <-s.peerClosed:
|
|
return
|
|
case <-s.stopped:
|
|
return
|
|
default:
|
|
}
|
|
var ce *websocket.CloseError
|
|
if errors.As(err, &ce) {
|
|
s.terminate(&CloseError{Code: ce.Code, Cause: err})
|
|
} else if errors.Is(err, websocket.ErrReadLimit) {
|
|
s.terminate(ErrMessageTooLarge)
|
|
} else {
|
|
s.terminate(&OpError{Op: "read", Cause: err})
|
|
}
|
|
return
|
|
}
|
|
data, err := io.ReadAll(io.LimitReader(reader, s.cfg.MaxMessageBytes+1))
|
|
if err != nil {
|
|
if errors.Is(err, websocket.ErrReadLimit) {
|
|
s.beginClose(1009, "", ErrMessageTooLarge)
|
|
} else {
|
|
s.terminate(&OpError{Op: "read", Cause: err})
|
|
}
|
|
return
|
|
}
|
|
if int64(len(data)) > s.cfg.MaxMessageBytes {
|
|
s.beginClose(1009, "", ErrMessageTooLarge)
|
|
return
|
|
}
|
|
if kind == websocket.TextMessage && !utf8.Valid(data) {
|
|
s.beginClose(1007, "", ErrInvalidMessage)
|
|
return
|
|
}
|
|
s.mu.Lock()
|
|
if s.closing {
|
|
s.mu.Unlock()
|
|
continue
|
|
}
|
|
if s.stats.ReadQueuedMessages >= s.cfg.ReadQueueSize || int64(len(data)) > s.cfg.ReadQueueBytes-s.stats.ReadQueuedBytes {
|
|
s.mu.Unlock()
|
|
s.beginClose(1008, "", ErrQueueFull)
|
|
return
|
|
}
|
|
msg := Message{Type: MessageType(kind), Data: data}
|
|
select {
|
|
case s.in <- msg:
|
|
s.stats.ReadQueuedMessages++
|
|
s.stats.ReadQueuedBytes += int64(len(data))
|
|
s.stats.MessagesReceived++
|
|
s.stats.BytesReceived += uint64(len(data))
|
|
s.mu.Unlock()
|
|
default:
|
|
s.mu.Unlock()
|
|
s.beginClose(1008, "", ErrQueueFull)
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *Session) releaseRead(msg Message) {
|
|
s.mu.Lock()
|
|
s.stats.ReadQueuedMessages--
|
|
s.stats.ReadQueuedBytes -= int64(len(msg.Data))
|
|
s.mu.Unlock()
|
|
}
|
|
func (s *Session) callback(fn func() error) (err error) {
|
|
defer func() {
|
|
if p := recover(); p != nil {
|
|
err = fmt.Errorf("callback panic: %v", p)
|
|
}
|
|
}()
|
|
return fn()
|
|
}
|
|
func (s *Session) dispatchLoop() {
|
|
if s.ctx.Err() != nil {
|
|
return
|
|
}
|
|
if s.handlers.OnOpen != nil {
|
|
if err := s.callback(func() error { return s.handlers.OnOpen(s.ctx, s) }); err != nil {
|
|
s.beginClose(1011, "", &OpError{Op: "on_open", Cause: err})
|
|
return
|
|
}
|
|
}
|
|
for {
|
|
select {
|
|
case <-s.ctx.Done():
|
|
return
|
|
case msg := <-s.in:
|
|
if s.ctx.Err() != nil {
|
|
s.releaseRead(msg)
|
|
return
|
|
}
|
|
var err error
|
|
if s.handlers.OnMessage != nil {
|
|
err = s.callback(func() error { return s.handlers.OnMessage(s.ctx, s, msg) })
|
|
}
|
|
s.releaseRead(msg)
|
|
if err != nil {
|
|
s.beginClose(1011, "", &OpError{Op: "on_message", Cause: err})
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *Session) complete(item outgoing, err error) {
|
|
s.mu.Lock()
|
|
s.stats.QueuedMessages--
|
|
s.stats.QueuedBytes -= int64(len(item.message.Data))
|
|
if err == nil {
|
|
s.stats.MessagesSent++
|
|
s.stats.BytesSent += uint64(len(item.message.Data))
|
|
}
|
|
s.mu.Unlock()
|
|
if item.result != nil {
|
|
item.result <- err
|
|
}
|
|
}
|
|
func (s *Session) write(item outgoing) error {
|
|
if err := item.ctx.Err(); err != nil {
|
|
s.complete(item, err)
|
|
return nil
|
|
}
|
|
deadline := time.Now().Add(s.cfg.WriteTimeout)
|
|
if d, ok := item.ctx.Deadline(); ok && d.Before(deadline) {
|
|
deadline = d
|
|
}
|
|
if err := s.conn.SetWriteDeadline(deadline); err != nil {
|
|
s.complete(item, &OpError{Op: "write", Cause: err})
|
|
return err
|
|
}
|
|
// Only the writer touches Gorilla's deadline field. Cancellation interrupts
|
|
// the underlying socket directly; joining the callback prevents a late expiry
|
|
// from corrupting the following message's deadline.
|
|
canceled := make(chan struct{})
|
|
stop := context.AfterFunc(item.ctx, func() { _ = s.conn.UnderlyingConn().SetWriteDeadline(time.Now()); close(canceled) })
|
|
err := s.conn.WriteMessage(int(item.message.Type), item.message.Data)
|
|
if !stop() {
|
|
<-canceled
|
|
}
|
|
if err != nil {
|
|
if item.ctx.Err() != nil {
|
|
err = item.ctx.Err()
|
|
} else if d, ok := item.ctx.Deadline(); ok && !time.Now().Before(d) {
|
|
// The socket timer can win the race with context's timer.
|
|
err = context.DeadlineExceeded
|
|
}
|
|
wrapped := &OpError{Op: "write", Cause: err}
|
|
s.complete(item, wrapped)
|
|
return wrapped
|
|
}
|
|
s.complete(item, nil)
|
|
return nil
|
|
}
|
|
func (s *Session) writeControl(frame controlFrame) error {
|
|
return s.conn.WriteControl(frame.kind, frame.data, time.Now().Add(s.cfg.WriteTimeout))
|
|
}
|
|
func (s *Session) writeLoop() {
|
|
var ticks <-chan time.Time
|
|
if !s.cfg.DisableHeartbeat {
|
|
ticker := time.NewTicker(s.cfg.PingInterval)
|
|
defer ticker.Stop()
|
|
ticks = ticker.C
|
|
}
|
|
for {
|
|
select {
|
|
case <-s.stopped:
|
|
return
|
|
case <-s.closingCh:
|
|
s.finishWrites()
|
|
return
|
|
default:
|
|
}
|
|
select {
|
|
case <-s.stopped:
|
|
return
|
|
case <-s.closingCh:
|
|
s.finishWrites()
|
|
return
|
|
case item := <-s.out:
|
|
if err := s.write(item); err != nil {
|
|
s.terminate(err)
|
|
return
|
|
}
|
|
case frame := <-s.controls:
|
|
if err := s.writeControl(frame); err != nil {
|
|
s.terminate(&OpError{Op: "control", Cause: err})
|
|
return
|
|
}
|
|
case <-ticks:
|
|
if err := s.writeControl(controlFrame{kind: websocket.PingMessage}); err != nil {
|
|
s.terminate(&OpError{Op: "ping", Cause: err})
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
func (s *Session) finishWrites() {
|
|
for {
|
|
select {
|
|
case <-s.stopped:
|
|
return
|
|
case item := <-s.out:
|
|
select {
|
|
case <-s.peerClosed:
|
|
s.complete(item, ErrClosed)
|
|
continue
|
|
default:
|
|
}
|
|
if err := s.write(item); err != nil {
|
|
s.terminate(err)
|
|
return
|
|
}
|
|
default:
|
|
goto closeFrame
|
|
}
|
|
}
|
|
closeFrame:
|
|
s.mu.Lock()
|
|
code, reason := s.closeCode, s.closeReason
|
|
s.mu.Unlock()
|
|
if err := s.writeControl(controlFrame{websocket.CloseMessage, websocket.FormatCloseMessage(code, reason)}); err != nil {
|
|
s.terminate(&OpError{Op: "close", Cause: err})
|
|
return
|
|
}
|
|
close(s.closeSent)
|
|
}
|