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/server.go
+++ b/ws/server.go
@@ -0,0 +1,405 @@
+package ws
+
+import (
+	"context"
+	"errors"
+	"net"
+	"net/http"
+	"net/url"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+// ServerConfig is independent of HTTP body limits and HTTP request deadlines.
+// Context owns accepted sessions; it must outlive HTTP handlers. Prepare runs
+// before upgrading and may authenticate or create per-connection handlers.
+type ServerConfig struct {
+	Session            SessionConfig                                          `json:"session"`
+	MaxConnections     int                                                    `json:"maxConnections"`
+	HandshakeTimeout   time.Duration                                          `json:"handshakeTimeout"`
+	AllowedOrigins     []string                                               `json:"allowedOrigins"`
+	RequireOrigin      bool                                                   `json:"requireOrigin"`
+	Subprotocols       []string                                               `json:"subprotocols"`
+	RequireSubprotocol bool                                                   `json:"requireSubprotocol"`
+	EnableCompression  bool                                                   `json:"enableCompression"`
+	HandshakeRate      float64                                                `json:"handshakeRate"` // Per second; zero disables rate limiting.
+	HandshakeBurst     int                                                    `json:"handshakeBurst"`
+	Context            context.Context                                        `json:"-"`
+	Prepare            func(context.Context, *http.Request) (Handlers, error) `json:"-"`
+}
+
+func DefaultServerConfig() ServerConfig {
+	return ServerConfig{Session: DefaultSessionConfig(), MaxConnections: 4096, HandshakeTimeout: 10 * time.Second}
+}
+
+func (cfg ServerConfig) Validate() error { _, err := normalizeServerConfig(cfg); return err }
+
+func normalizeServerConfig(cfg ServerConfig) (ServerConfig, error) {
+	var err error
+	cfg.Session, err = normalizeSessionConfig(cfg.Session)
+	if err != nil {
+		return cfg, err
+	}
+	if cfg.MaxConnections < 0 || cfg.MaxConnections > 1<<20 || cfg.HandshakeTimeout < 0 || cfg.HandshakeRate < 0 || cfg.HandshakeRate > 1e9 || cfg.HandshakeRate != cfg.HandshakeRate || cfg.HandshakeBurst < 0 || cfg.HandshakeBurst > 1<<20 {
+		return cfg, errors.New("ws: invalid server limits")
+	}
+	if cfg.MaxConnections == 0 {
+		cfg.MaxConnections = 4096
+	}
+	if cfg.HandshakeTimeout == 0 {
+		cfg.HandshakeTimeout = 10 * time.Second
+	}
+	if cfg.HandshakeRate > 0 && cfg.HandshakeBurst == 0 {
+		cfg.HandshakeBurst = 1
+	}
+	if cfg.HandshakeRate == 0 && cfg.HandshakeBurst != 0 {
+		return cfg, errors.New("ws: handshake burst requires a rate")
+	}
+	if len(cfg.AllowedOrigins) > 1024 {
+		return cfg, errors.New("ws: too many allowed origins")
+	}
+	cfg.AllowedOrigins = append([]string(nil), cfg.AllowedOrigins...)
+	for i, v := range cfg.AllowedOrigins {
+		o, ok := canonicalOrigin(v)
+		if !ok {
+			return cfg, errors.New("ws: invalid allowed origin")
+		}
+		cfg.AllowedOrigins[i] = o
+	}
+	if len(cfg.Subprotocols) > 64 {
+		return cfg, errors.New("ws: too many subprotocols")
+	}
+	cfg.Subprotocols = append([]string(nil), cfg.Subprotocols...)
+	seen := make(map[string]bool)
+	for _, v := range cfg.Subprotocols {
+		if !protocolToken(v) || seen[v] {
+			return cfg, errors.New("ws: invalid or duplicate subprotocol")
+		}
+		seen[v] = true
+	}
+	if cfg.RequireSubprotocol && len(cfg.Subprotocols) == 0 {
+		return cfg, errors.New("ws: required subprotocol list is empty")
+	}
+	if cfg.Context == nil {
+		cfg.Context = context.Background()
+	}
+	return cfg, nil
+}
+
+func protocolToken(s string) bool {
+	if len(s) == 0 || len(s) > 256 {
+		return false
+	}
+	for _, r := range s {
+		if r <= 32 || r >= 127 || strings.ContainsRune("()<>@,;:\\\"/[]?={}", r) {
+			return false
+		}
+	}
+	return true
+}
+
+func canonicalOrigin(raw string) (string, bool) {
+	if len(raw) > 4096 || strings.ContainsAny(raw, " \t\r\n,\\?#") {
+		return "", false
+	}
+	u, err := url.Parse(raw)
+	if err != nil || u == nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.Path != "" || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || u.Opaque != "" || strings.Contains(u.Host, "*") || strings.HasSuffix(u.Host, ":") {
+		return "", false
+	}
+	host := strings.ToLower(u.Host)
+	// net/url accepts some unbracketed colon hosts. An Origin must contain a
+	// serialized DNS host or a bracketed, valid IPv6 address, never that ambiguity.
+	name := u.Hostname()
+	if strings.Contains(name, ":") {
+		if !strings.HasPrefix(u.Host, "[") || net.ParseIP(name) == nil {
+			return "", false
+		}
+	} else {
+		if strings.ContainsAny(u.Host, "[]") {
+			return "", false
+		}
+		for _, r := range name {
+			if !(r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '.' || r == '-' || r == '_') {
+				return "", false
+			}
+		}
+	}
+	if port := u.Port(); port != "" {
+		n, err := strconv.Atoi(port)
+		if err != nil || n < 1 || n > 65535 {
+			return "", false
+		}
+	}
+	if u.Scheme == "http" && u.Port() == "80" {
+		host = strings.TrimSuffix(host, ":80")
+	}
+	if u.Scheme == "https" && u.Port() == "443" {
+		host = strings.TrimSuffix(host, ":443")
+	}
+	return u.Scheme + "://" + host, true
+}
+
+// RejectError lets Prepare select a public HTTP 4xx status. Error details are
+// deliberately never written into the handshake response.
+type RejectError struct{ StatusCode int }
+
+func (e *RejectError) Error() string { return "ws: handshake rejected" }
+
+type ServerStats struct {
+	ActiveConnections, Handshakes int
+	Accepted, Rejected            uint64
+}
+
+// Server owns all upgraded sessions, including their callbacks. Use Stop and
+// Drain with lifecycle.Runtime, or Shutdown before releasing shared resources.
+type Server struct {
+	cfg                ServerConfig
+	mu                 sync.Mutex
+	stopped            bool
+	active             int // includes reserved handshakes; protected by the same admission lock
+	handshakes         int
+	sessions           map[string]*Session
+	done               chan struct{}
+	doneClosed         bool
+	accepted, rejected uint64
+	tokens             float64
+	refill             time.Time
+	stopParent         func() bool
+}
+
+func NewServer(cfg ServerConfig) (*Server, error) {
+	cfg, err := normalizeServerConfig(cfg)
+	if err != nil {
+		return nil, err
+	}
+	s := &Server{cfg: cfg, sessions: make(map[string]*Session), done: make(chan struct{}), tokens: float64(cfg.HandshakeBurst), refill: time.Now()}
+	s.stopParent = context.AfterFunc(cfg.Context, func() { _ = s.Stop(context.Background()) })
+	return s, nil
+}
+
+func (s *Server) Handler(h Handlers) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = s.Accept(w, r, h) })
+}
+func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { _, _ = s.Accept(w, r, Handlers{}) }
+
+func (s *Server) originAllowed(r *http.Request) bool {
+	values := r.Header.Values("Origin")
+	if len(values) == 0 {
+		return !s.cfg.RequireOrigin
+	}
+	if len(values) != 1 {
+		return false
+	}
+	origin, ok := canonicalOrigin(values[0])
+	if !ok {
+		return false
+	}
+	if len(s.cfg.AllowedOrigins) > 0 {
+		for _, v := range s.cfg.AllowedOrigins {
+			if v == origin {
+				return true
+			}
+		}
+		return false
+	}
+	scheme := "http"
+	if r.TLS != nil {
+		scheme = "https"
+	}
+	self, ok := canonicalOrigin(scheme + "://" + r.Host)
+	return ok && origin == self
+}
+
+func (s *Server) admit() int {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.stopped || s.cfg.Context.Err() != nil {
+		s.rejected++
+		return http.StatusServiceUnavailable
+	}
+	if s.active >= s.cfg.MaxConnections {
+		s.rejected++
+		return http.StatusServiceUnavailable
+	}
+	if s.cfg.HandshakeRate > 0 {
+		now := time.Now()
+		s.tokens = min(float64(s.cfg.HandshakeBurst), s.tokens+now.Sub(s.refill).Seconds()*s.cfg.HandshakeRate)
+		s.refill = now
+		if s.tokens < 1 {
+			s.rejected++
+			return http.StatusTooManyRequests
+		}
+		s.tokens--
+	}
+	s.active++
+	s.handshakes++
+	return 0
+}
+
+func (s *Server) finishLocked() {
+	if s.stopped && s.active == 0 && !s.doneClosed {
+		s.doneClosed = true
+		close(s.done)
+	}
+}
+
+// Accept writes the handshake response (including failures). It returns after
+// transferring the connection to Server; callers must not reuse the writer or
+// retain request-scoped state in callbacks. Session uses ServerConfig.Context.
+func (s *Server) Accept(w http.ResponseWriter, r *http.Request, h Handlers) (session *Session, retErr error) {
+	if status := s.admit(); status != 0 {
+		http.Error(w, http.StatusText(status), status)
+		return nil, &HandshakeError{StatusCode: status}
+	}
+	transferred := false
+	defer func() {
+		if !transferred {
+			s.mu.Lock()
+			s.active--
+			s.handshakes--
+			s.rejected++
+			s.finishLocked()
+			s.mu.Unlock()
+		}
+	}()
+	fail := func(status int, cause error) (*Session, error) {
+		http.Error(w, http.StatusText(status), status)
+		return nil, &HandshakeError{StatusCode: status, Cause: cause}
+	}
+	if r.Method != http.MethodGet || !websocket.IsWebSocketUpgrade(r) {
+		return fail(http.StatusBadRequest, nil)
+	}
+	if !s.originAllowed(r) {
+		return fail(http.StatusForbidden, nil)
+	}
+	if _, ok := w.(http.Hijacker); !ok {
+		return fail(http.StatusNotImplemented, ErrUnsupportedTransport)
+	}
+	if s.cfg.RequireSubprotocol {
+		matched := false
+		for _, offered := range websocket.Subprotocols(r) {
+			for _, supported := range s.cfg.Subprotocols {
+				if offered == supported {
+					matched = true
+				}
+			}
+		}
+		if !matched {
+			return fail(http.StatusBadRequest, nil)
+		}
+	}
+	ctx, cancel := context.WithTimeout(r.Context(), s.cfg.HandshakeTimeout)
+	defer cancel()
+	if s.cfg.Prepare != nil {
+		var err error
+		h, err = prepareSafely(s.cfg.Prepare, ctx, r)
+		if err != nil {
+			status := http.StatusUnauthorized
+			var reject *RejectError
+			if errors.As(err, &reject) && reject != nil && reject.StatusCode >= 400 && reject.StatusCode < 500 {
+				status = reject.StatusCode
+			}
+			return fail(status, nil)
+		}
+	}
+	if ctx.Err() != nil {
+		return fail(http.StatusRequestTimeout, ctx.Err())
+	}
+	deadline, _ := ctx.Deadline()
+	remaining := time.Until(deadline)
+	if remaining <= 0 {
+		return fail(http.StatusRequestTimeout, context.DeadlineExceeded)
+	}
+	upgrader := websocket.Upgrader{HandshakeTimeout: remaining, Subprotocols: s.cfg.Subprotocols, EnableCompression: s.cfg.EnableCompression, CheckOrigin: func(*http.Request) bool { return true }, Error: func(w http.ResponseWriter, _ *http.Request, status int, _ error) {
+		http.Error(w, http.StatusText(status), status)
+	}}
+	conn, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		return nil, &HandshakeError{Cause: errors.New("upgrade failed")}
+	}
+	session = newSession(conn, s.cfg.Session, h, s.cfg.Context)
+	s.mu.Lock()
+	s.handshakes--
+	s.sessions[session.ID()] = session
+	s.accepted++
+	stopped := s.stopped
+	transferred = true
+	s.mu.Unlock()
+	session.start()
+	if stopped {
+		_ = session.Close(websocket.CloseGoingAway, "")
+	}
+	go func() {
+		<-session.Done()
+		s.mu.Lock()
+		delete(s.sessions, session.ID())
+		s.active--
+		s.finishLocked()
+		s.mu.Unlock()
+	}()
+	return session, nil
+}
+
+func prepareSafely(fn func(context.Context, *http.Request) (Handlers, error), ctx context.Context, r *http.Request) (h Handlers, err error) {
+	defer func() {
+		if recover() != nil {
+			err = errors.New("ws: preparation failed")
+		}
+	}()
+	return fn(ctx, r)
+}
+
+// Stop closes admission and begins bounded graceful closing of existing
+// sessions. It never waits for callbacks and is safe to call from a callback.
+func (s *Server) Stop(context.Context) error {
+	s.mu.Lock()
+	s.stopped = true
+	all := make([]*Session, 0, len(s.sessions))
+	for _, v := range s.sessions {
+		all = append(all, v)
+	}
+	s.finishLocked()
+	s.mu.Unlock()
+	for _, v := range all {
+		_ = v.Close(websocket.CloseGoingAway, "")
+	}
+	return nil
+}
+
+// Drain returns nil only when all handshakes, I/O and callbacks have exited.
+// On timeout it aborts sockets, but does not claim blocked user callbacks ended.
+func (s *Server) Drain(ctx context.Context) error {
+	_ = s.Stop(ctx)
+	select {
+	case <-s.done:
+		s.stopParent()
+		return nil
+	default:
+	}
+	select {
+	case <-s.done:
+		s.stopParent()
+		return nil
+	case <-ctx.Done():
+		s.mu.Lock()
+		all := make([]*Session, 0, len(s.sessions))
+		for _, v := range s.sessions {
+			all = append(all, v)
+		}
+		s.mu.Unlock()
+		for _, v := range all {
+			v.Abort()
+		}
+		return ctx.Err()
+	}
+}
+func (s *Server) Shutdown(ctx context.Context) error { return s.Drain(ctx) }
+func (s *Server) Stats() ServerStats {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return ServerStats{ActiveConnections: len(s.sessions), Handshakes: s.handshakes, Accepted: s.accepted, Rejected: s.rejected}
+}