feat(ws): release generic WebSocket client and server v0.31.0-preview

2026-10-29 19:20:12 +08:00
parent 19ee9746e9
commit 3272974bc1
58 changed files with 7654 additions and 21 deletions
--- a/ws/client.go
+++ b/ws/client.go
@@ -0,0 +1,633 @@
+package ws
+
+import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"errors"
+	"io"
+	"math/rand"
+	"net"
+	"net/http"
+	"net/url"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+var ErrConnectionLimit = errors.New("ws: connection limit reached")
+
+// ReconnectConfig controls Run only. Dial never retries. Zero MaxAttempts and
+// MaxElapsed mean unlimited retries/time when Enabled is explicitly true.
+// A stable connection resets both retry budgets; short flaps do not.
+type ReconnectConfig struct {
+	Enabled     bool          `json:"enabled"`
+	MinBackoff  time.Duration `json:"minBackoff"`
+	MaxBackoff  time.Duration `json:"maxBackoff"`
+	StableAfter time.Duration `json:"stableAfter"`
+	MaxAttempts int           `json:"maxAttempts"`
+	MaxElapsed  time.Duration `json:"maxElapsed"`
+}
+
+type ClientConfig struct {
+	Session            SessionConfig `json:"session"`
+	HandshakeTimeout   time.Duration `json:"handshakeTimeout"`
+	MaxConnections     int           `json:"maxConnections"`
+	Subprotocols       []string      `json:"subprotocols"`
+	RequireSubprotocol bool          `json:"requireSubprotocol"`
+	EnableCompression  bool          `json:"enableCompression"`
+	// ProxyMode is environment (default), direct, or url. ProxyURL may contain
+	// credentials and must not be logged. Explicit proxy URLs ignore NO_PROXY.
+	// Supported proxy schemes are http, socks5, and socks5h (remote DNS).
+	ProxyMode string          `json:"proxyMode"`
+	ProxyURL  string          `json:"proxyUrl"`
+	Reconnect ReconnectConfig `json:"reconnect"`
+	Headers   http.Header     `json:"-"`
+	// HeaderProvider is called before each handshake. Return a fresh header map;
+	// the client snapshots it. Provider values replace static Headers values.
+	// Implementations must honor context cancellation and support concurrent calls.
+	HeaderProvider func(context.Context) (http.Header, error) `json:"-"`
+	// Prepare creates per-session handlers after a successful handshake. When
+	// set, its result replaces Dial/Run handlers. It must honor cancellation.
+	Prepare   func(context.Context) (Handlers, error) `json:"-"`
+	TLSConfig *tls.Config                             `json:"-"`
+	// ShouldReconnect overrides default error classification, not budgets or
+	// explicit client/context shutdown. It must be safe for concurrent Run calls.
+	ShouldReconnect func(error) bool `json:"-"`
+}
+
+func (cfg ClientConfig) Validate() error { _, _, err := normalizeClientConfig(cfg); return err }
+
+func normalizeClientConfig(cfg ClientConfig) (ClientConfig, func(*http.Request) (*url.URL, error), error) {
+	var err error
+	cfg.Session, err = normalizeSessionConfig(cfg.Session)
+	if err != nil {
+		return cfg, nil, err
+	}
+	if cfg.HandshakeTimeout < 0 || cfg.MaxConnections < 0 || cfg.MaxConnections > 65536 || cfg.Reconnect.MinBackoff < 0 || cfg.Reconnect.MaxBackoff < 0 || cfg.Reconnect.StableAfter < 0 || cfg.Reconnect.MaxElapsed < 0 || cfg.Reconnect.MaxAttempts < 0 {
+		return cfg, nil, errors.New("ws: invalid client limits or durations")
+	}
+	if cfg.HandshakeTimeout == 0 {
+		cfg.HandshakeTimeout = 10 * time.Second
+	}
+	if cfg.MaxConnections == 0 {
+		cfg.MaxConnections = 64
+	}
+	if cfg.Reconnect.MinBackoff == 0 {
+		cfg.Reconnect.MinBackoff = 100 * time.Millisecond
+	}
+	if cfg.Reconnect.MaxBackoff == 0 {
+		cfg.Reconnect.MaxBackoff = 30 * time.Second
+	}
+	if cfg.Reconnect.StableAfter == 0 {
+		cfg.Reconnect.StableAfter = 30 * time.Second
+	}
+	if cfg.Reconnect.MinBackoff > cfg.Reconnect.MaxBackoff {
+		return cfg, nil, errors.New("ws: minBackoff exceeds maxBackoff")
+	}
+	cfg.Headers, err = clientHeaders(cfg.Headers)
+	if err != nil {
+		return cfg, nil, err
+	}
+	if len(cfg.Subprotocols) > 32 {
+		return cfg, nil, errors.New("ws: too many subprotocols")
+	}
+	seen := map[string]bool{}
+	for _, protocol := range cfg.Subprotocols {
+		if !headerToken(protocol) || len(protocol) > 256 || seen[protocol] {
+			return cfg, nil, errors.New("ws: invalid or duplicate subprotocol")
+		}
+		seen[protocol] = true
+	}
+	if cfg.RequireSubprotocol && len(cfg.Subprotocols) == 0 {
+		return cfg, nil, errors.New("ws: required subprotocol list is empty")
+	}
+	cfg.Subprotocols = append([]string(nil), cfg.Subprotocols...)
+	if cfg.TLSConfig != nil {
+		cfg.TLSConfig = cfg.TLSConfig.Clone()
+	}
+	if cfg.ProxyMode != "url" && cfg.ProxyURL != "" {
+		return cfg, nil, errors.New("ws: proxyUrl requires url proxyMode")
+	}
+	var proxy func(*http.Request) (*url.URL, error)
+	switch cfg.ProxyMode {
+	case "", "environment":
+		proxy = func(r *http.Request) (*url.URL, error) {
+			u, err := http.ProxyFromEnvironment(r)
+			if err != nil || u == nil {
+				return u, err
+			}
+			return supportedProxy(u)
+		}
+	case "direct":
+	case "url":
+		u, e := url.Parse(cfg.ProxyURL)
+		if e != nil || u == nil || u.Hostname() == "" || u.Opaque != "" || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || (u.Path != "" && u.Path != "/") || !validPort(u) {
+			return cfg, nil, errors.New("ws: invalid proxy URL")
+		}
+		u, e = supportedProxy(u)
+		if e != nil {
+			return cfg, nil, e
+		}
+		proxy = http.ProxyURL(u)
+	default:
+		return cfg, nil, errors.New("ws: invalid proxyMode")
+	}
+	return cfg, proxy, nil
+}
+
+func supportedProxy(u *url.URL) (*url.URL, error) {
+	copyURL := *u
+	switch copyURL.Scheme {
+	case "http", "socks5":
+	case "socks5h":
+		copyURL.Scheme = "socks5"
+	default:
+		return nil, errors.New("ws: proxy scheme must be http, socks5 or socks5h")
+	}
+	return &copyURL, nil
+}
+
+func validPort(u *url.URL) bool {
+	if strings.HasSuffix(u.Host, ":") {
+		return false
+	}
+	if p := u.Port(); p != "" {
+		n, err := strconv.Atoi(p)
+		return err == nil && n > 0 && n <= 65535
+	}
+	return true
+}
+
+func headerToken(s string) bool {
+	if s == "" {
+		return false
+	}
+	for i := 0; i < len(s); i++ {
+		c := s[i]
+		if c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' || strings.ContainsRune("!#$%&'*+-.^_`|~", rune(c)) {
+			continue
+		}
+		return false
+	}
+	return true
+}
+func clientHeaders(input http.Header) (http.Header, error) {
+	if len(input) > 128 {
+		return nil, errors.New("ws: too many handshake headers")
+	}
+	out := make(http.Header, len(input))
+	size := 0
+	for key, values := range input {
+		name := http.CanonicalHeaderKey(key)
+		if !headerToken(key) || len(key) > 256 {
+			return nil, errors.New("ws: invalid handshake header name")
+		}
+		switch name {
+		case "Host", "Connection", "Upgrade", "Sec-Websocket-Key", "Sec-Websocket-Version", "Sec-Websocket-Extensions", "Sec-Websocket-Protocol", "Content-Length", "Transfer-Encoding", "Trailer":
+			return nil, errors.New("ws: reserved handshake header")
+		}
+		if _, ok := out[name]; ok {
+			return nil, errors.New("ws: duplicate handshake header")
+		}
+		for _, value := range values {
+			size += len(key) + len(value)
+			for i := 0; i < len(value); i++ {
+				if value[i] < 32 && value[i] != '\t' || value[i] == 127 {
+					return nil, errors.New("ws: invalid handshake header value")
+				}
+			}
+		}
+		if size > 64<<10 {
+			return nil, errors.New("ws: handshake headers exceed limit")
+		}
+		out[name] = append([]string(nil), values...)
+	}
+	return out, nil
+}
+
+// Client owns its sessions. Concurrent Dial and Run calls are independent and
+// share only limits and lifecycle. One Run owns one reconnecting connection;
+// every successful handshake produces a new Session with a new ID.
+type Client struct {
+	config                                      ClientConfig
+	dialer                                      websocket.Dialer
+	ctx                                         context.Context
+	cancel                                      context.CancelFunc
+	mu                                          sync.Mutex
+	closed                                      bool
+	active                                      int
+	connections                                 int
+	dialAttempts, handshakeFailures, reconnects uint64
+	sessions                                    map[string]*Session
+	releases                                    map[string]func()
+	done                                        chan struct{}
+}
+
+func NewClient(cfg ClientConfig) (*Client, error) {
+	cfg, proxy, err := normalizeClientConfig(cfg)
+	if err != nil {
+		return nil, err
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	return &Client{config: cfg, dialer: websocket.Dialer{Proxy: proxy, HandshakeTimeout: cfg.HandshakeTimeout, TLSClientConfig: cfg.TLSConfig, Subprotocols: cfg.Subprotocols, EnableCompression: cfg.EnableCompression}, ctx: ctx, cancel: cancel, sessions: make(map[string]*Session), releases: make(map[string]func()), done: make(chan struct{})}, nil
+}
+
+func (c *Client) begin(connection bool) error {
+	for {
+		c.mu.Lock()
+		if c.closed {
+			c.mu.Unlock()
+			return ErrClosed
+		}
+		if connection && c.connections >= c.config.MaxConnections {
+			// Done may wake a caller before its registry watcher gets CPU time.
+			// Retire finished sessions synchronously before reporting overload.
+			var release func()
+			for id, s := range c.sessions {
+				select {
+				case <-s.Done():
+					release = c.releases[id]
+				default:
+				}
+				if release != nil {
+					break
+				}
+			}
+			c.mu.Unlock()
+			if release != nil {
+				release()
+				continue
+			}
+			return ErrConnectionLimit
+		}
+		c.active++
+		if connection {
+			c.connections++
+			c.dialAttempts++
+		}
+		c.mu.Unlock()
+		return nil
+	}
+}
+func (c *Client) end(connection bool, id string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if id != "" {
+		delete(c.sessions, id)
+		delete(c.releases, id)
+	}
+	c.active--
+	if connection {
+		c.connections--
+	}
+	if c.closed && c.active == 0 {
+		close(c.done)
+	}
+}
+
+// Dial performs exactly one handshake. Its timeout ends at the handshake;
+// ctx continues to own the resulting session's lifetime. No message is replayed.
+func (c *Client) Dial(ctx context.Context, endpoint string, handlers Handlers) (*Session, error) {
+	return c.dial(ctx, endpoint, handlers, 0)
+}
+func (c *Client) dial(ctx context.Context, endpoint string, handlers Handlers, budget time.Duration) (*Session, error) {
+	if ctx == nil {
+		return nil, errors.New("ws: nil context")
+	}
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	u, err := url.Parse(endpoint)
+	if err != nil || u == nil || u.Hostname() == "" || u.User != nil || u.Fragment != "" || u.Opaque != "" || (u.Scheme != "ws" && u.Scheme != "wss") || !validPort(u) {
+		return nil, errors.New("ws: invalid endpoint")
+	}
+	if err = c.begin(true); err != nil {
+		return nil, err
+	}
+	transferred := false
+	sessionCtx, cancelSession := context.WithCancel(ctx)
+	stopSessionShutdown := context.AfterFunc(c.ctx, cancelSession)
+	defer func() {
+		if !transferred {
+			c.mu.Lock()
+			c.handshakeFailures++
+			c.mu.Unlock()
+			stopSessionShutdown()
+			cancelSession()
+			c.end(true, "")
+		}
+	}()
+	handshakeTimeout := c.config.HandshakeTimeout
+	if budget > 0 && budget < handshakeTimeout {
+		handshakeTimeout = budget
+	}
+	hctx, cancel := context.WithTimeout(ctx, handshakeTimeout)
+	stop := context.AfterFunc(c.ctx, cancel)
+	defer func() { stop(); cancel() }()
+	headers := c.config.Headers.Clone()
+	if c.config.HeaderProvider != nil {
+		provided, e := provideHeaders(c.config.HeaderProvider, hctx)
+		if e != nil {
+			return nil, &OpError{Op: "headers", Cause: e}
+		}
+		provided, e = clientHeaders(provided)
+		if e != nil {
+			return nil, &OpError{Op: "headers", Cause: e}
+		}
+		for k, v := range provided {
+			headers[k] = v
+		}
+	}
+	dialer := c.dialer
+	// Gorilla applies the context deadline to handshake I/O, but canceling a
+	// deadline-less context after TCP connect does not interrupt its HTTP read.
+	// Retain the raw socket only through the handshake so Close/Stop cancels
+	// stalled HTTP and proxy handshakes immediately, without timing out sessions.
+	var stopHandshake func() bool
+	dialer.NetDialContext = func(dialctx context.Context, network, address string) (net.Conn, error) {
+		conn, err := (&net.Dialer{Timeout: c.config.HandshakeTimeout, KeepAlive: 30 * time.Second}).DialContext(dialctx, network, address)
+		if err == nil {
+			stopHandshake = context.AfterFunc(hctx, func() { _ = conn.Close() })
+		}
+		return conn, err
+	}
+	if dialer.TLSClientConfig != nil {
+		dialer.TLSClientConfig = dialer.TLSClientConfig.Clone()
+	}
+	conn, response, err := dialer.DialContext(hctx, endpoint, headers)
+	if stopHandshake != nil {
+		stopHandshake()
+	}
+	if err == nil && hctx.Err() != nil {
+		_ = conn.Close()
+		err = hctx.Err()
+	}
+	if err != nil {
+		status := 0
+		if response != nil {
+			status = response.StatusCode
+			if response.Body != nil {
+				_ = response.Body.Close()
+			}
+		}
+		if ctx.Err() != nil {
+			return nil, ctx.Err()
+		}
+		if c.ctx.Err() != nil {
+			return nil, ErrClosed
+		}
+		if hctx.Err() != nil {
+			err = hctx.Err()
+		} else if deadline, ok := hctx.Deadline(); ok && !time.Now().Before(deadline) {
+			var timed net.Error
+			if errors.As(err, &timed) && timed.Timeout() {
+				err = errors.Join(context.DeadlineExceeded, err)
+			}
+		}
+		return nil, &HandshakeError{StatusCode: status, Cause: err}
+	}
+	if c.config.RequireSubprotocol && conn.Subprotocol() == "" {
+		_ = conn.Close()
+		return nil, &HandshakeError{StatusCode: http.StatusSwitchingProtocols, Cause: websocket.ErrBadHandshake}
+	}
+	if c.config.Prepare != nil {
+		handlers, err = prepareClient(c.config.Prepare, sessionCtx)
+		if err != nil {
+			_ = conn.Close()
+			return nil, &OpError{Op: "prepare", Cause: err}
+		}
+	}
+	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		_ = conn.Close()
+		return nil, ErrClosed
+	}
+	s := newSession(conn, c.config.Session, handlers, sessionCtx)
+	c.sessions[s.ID()] = s
+	var releaseOnce sync.Once
+	release := func() { releaseOnce.Do(func() { cancelSession(); c.end(true, s.ID()) }) }
+	c.releases[s.ID()] = release
+	transferred = true
+	stopSessionShutdown()
+	c.mu.Unlock()
+	s.start()
+	go func() { <-s.Done(); release() }()
+	return s, nil
+}
+
+func provideHeaders(fn func(context.Context) (http.Header, error), ctx context.Context) (headers http.Header, err error) {
+	defer func() {
+		if recover() != nil {
+			err = errors.New("ws: header provider panicked")
+		}
+	}()
+	return fn(ctx)
+}
+func prepareClient(fn func(context.Context) (Handlers, error), ctx context.Context) (handlers Handlers, err error) {
+	defer func() {
+		if recover() != nil {
+			err = errors.New("ws: prepare callback panicked")
+		}
+	}()
+	return fn(ctx)
+}
+
+// Run blocks until canceled, closed, or retry policy/budget stops it. Reconnect
+// must be enabled explicitly. Handlers are reused; OnOpen must restore any
+// application subscriptions. Credentials are refreshed by HeaderProvider.
+func (c *Client) Run(ctx context.Context, endpoint string, handlers Handlers) error {
+	if ctx == nil {
+		return errors.New("ws: nil context")
+	}
+	if err := c.begin(false); err != nil {
+		return err
+	}
+	defer c.end(false, "")
+	runctx, cancel := context.WithCancel(ctx)
+	stop := context.AfterFunc(c.ctx, cancel)
+	defer func() { stop(); cancel() }()
+	attempts := 0
+	episode := time.Now()
+	var previousError error = context.DeadlineExceeded
+	repeat := false
+	for {
+		var remaining time.Duration
+		if c.config.Reconnect.Enabled && c.config.Reconnect.MaxElapsed > 0 {
+			remaining = c.config.Reconnect.MaxElapsed - time.Since(episode)
+			if remaining <= 0 {
+				return previousError
+			}
+		}
+		started := time.Now()
+		if repeat {
+			c.mu.Lock()
+			c.reconnects++
+			c.mu.Unlock()
+		}
+		repeat = true
+		s, err := c.dial(runctx, endpoint, handlers, remaining)
+		if err == nil {
+			started = time.Now()
+			err = s.Wait(runctx)
+			if time.Since(started) >= c.config.Reconnect.StableAfter {
+				attempts = 0
+				episode = time.Now()
+			}
+		}
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+		if c.ctx.Err() != nil {
+			return nil
+		}
+		if !c.config.Reconnect.Enabled || !c.retryable(err) {
+			return err
+		}
+		previousError = err
+		rc := c.config.Reconnect
+		if rc.MaxAttempts > 0 && attempts >= rc.MaxAttempts {
+			return err
+		}
+		delay := retryBackoff(rc, attempts)
+		if rc.MaxElapsed > 0 && time.Since(episode)+delay >= rc.MaxElapsed {
+			return err
+		}
+		attempts++
+		timer := time.NewTimer(delay)
+		select {
+		case <-runctx.Done():
+			timer.Stop()
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
+			return nil
+		case <-timer.C:
+		}
+	}
+}
+
+func retryBackoff(cfg ReconnectConfig, attempt int) time.Duration {
+	d := cfg.MinBackoff
+	for i := 0; i < attempt && d < cfg.MaxBackoff; i++ {
+		if d > cfg.MaxBackoff/2 {
+			d = cfg.MaxBackoff
+		} else {
+			d *= 2
+		}
+	}
+	if d > 1 {
+		d = d/2 + time.Duration(rand.Int63n(int64(d-d/2))+1)
+	}
+	return d
+}
+func (c *Client) retryable(err error) bool {
+	if err == nil || errors.Is(err, ErrClosed) || errors.Is(err, context.Canceled) {
+		return false
+	}
+	if c.config.ShouldReconnect != nil {
+		return c.config.ShouldReconnect(err)
+	}
+	var handshake *HandshakeError
+	if errors.As(err, &handshake) {
+		switch handshake.StatusCode {
+		case 429, 500, 502, 503, 504:
+			return true
+		case 0:
+		default:
+			return false
+		}
+		if errors.Is(err, websocket.ErrBadHandshake) {
+			return false
+		}
+	}
+	var unknown x509.UnknownAuthorityError
+	var invalid x509.CertificateInvalidError
+	var hostname x509.HostnameError
+	if errors.As(err, &unknown) || errors.As(err, &invalid) || errors.As(err, &hostname) {
+		return false
+	}
+	var closeError *CloseError
+	if errors.As(err, &closeError) {
+		switch closeError.Code {
+		case 1001, 1006, 1011, 1012, 1013:
+			return true
+		default:
+			return false
+		}
+	}
+	var network net.Error
+	if errors.As(err, &network) || errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
+		return true
+	}
+	return handshake != nil || errors.Is(err, ErrNotConnected)
+}
+
+// Close is idempotent and nonblocking: it rejects new work, cancels handshakes
+// and reconnect waits, and requests session close. Drain waits for all work,
+// including user callbacks. A callback that ignores cancellation cannot be killed.
+func (c *Client) Close() error {
+	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		return nil
+	}
+	c.closed = true
+	sessions := make([]*Session, 0, len(c.sessions))
+	for _, s := range c.sessions {
+		sessions = append(sessions, s)
+	}
+	if c.active == 0 {
+		close(c.done)
+	}
+	c.mu.Unlock()
+	c.cancel()
+	for _, s := range sessions {
+		_ = s.Close(1001, "client shutdown")
+	}
+	return nil
+}
+func (c *Client) Stop(context.Context) error { return c.Close() }
+func (c *Client) Drain(ctx context.Context) error {
+	if ctx == nil {
+		return errors.New("ws: nil context")
+	}
+	select {
+	case <-c.done:
+		return nil
+	case <-ctx.Done():
+		c.mu.Lock()
+		sessions := make([]*Session, 0, len(c.sessions))
+		for _, s := range c.sessions {
+			sessions = append(sessions, s)
+		}
+		c.mu.Unlock()
+		for _, s := range sessions {
+			s.Abort()
+		}
+		return ctx.Err()
+	}
+}
+func (c *Client) Shutdown(ctx context.Context) error { _ = c.Close(); return c.Drain(ctx) }
+
+type ClientStats struct {
+	ActiveConnections int
+	Handshakes        int
+	Running           int
+	Closed            bool
+	// DialAttempts includes admitted attempts canceled before the HTTP handshake.
+	DialAttempts      uint64
+	HandshakeFailures uint64
+	// Reconnects counts subsequent Dial attempts made by Run, successful or not.
+	Reconnects uint64
+}
+
+func (c *Client) Stats() ClientStats {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return ClientStats{ActiveConnections: len(c.sessions), Handshakes: c.connections - len(c.sessions), Running: c.active - c.connections, Closed: c.closed, DialAttempts: c.dialAttempts, HandshakeFailures: c.handshakeFailures, Reconnects: c.reconnects}
+}