go-omni/ws/client.go

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}
}