406 lines
12 KiB
Go
406 lines
12 KiB
Go
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}
|
|
}
|