// Package websocket provides a small, framework-neutral wrapper around Melody. package websocket import ( "net/http" "net/url" "strings" "sync" "time" melody "github.com/olahol/melody" ) type Config struct { WriteWait time.Duration PongWait time.Duration PingPeriod time.Duration MaxMessageSize int64 MessageBufferSize int ConcurrentMessageHandling bool AllowOrigins []string } type Session = melody.Session // Hub is the shared WebSocket boundary consumed by business modules. The // application composition root supplies the runtime-configured implementation. type Hub interface { Enabled() bool Broadcast([]byte) error BroadcastBinary([]byte) error Sessions() ([]*Session, error) Len() int Send(*Session, []byte) error SendBinary(*Session, []byte) error OnMessage(func(*Session, []byte)) OnBinaryMessage(func(*Session, []byte)) OnConnect(func(*Session)) OnDisconnect(func(*Session)) } func (c Config) defaults() Config { if c.WriteWait <= 0 { c.WriteWait = 10 * time.Second } if c.PongWait <= 0 { c.PongWait = 60 * time.Second } if c.PingPeriod <= 0 { c.PingPeriod = (c.PongWait * 9) / 10 } if c.MaxMessageSize <= 0 { c.MaxMessageSize = 1 << 20 } if c.MessageBufferSize <= 0 { c.MessageBufferSize = 256 } return c } type Server struct { m *melody.Melody mu sync.RWMutex closed bool messageHandlers []func(*melody.Session, []byte) binaryHandlers []func(*melody.Session, []byte) connectHandlers []func(*melody.Session) disconnectHandlers []func(*melody.Session) } func New(config Config) *Server { c := config.defaults() m := melody.New() m.Config.WriteWait = c.WriteWait m.Config.PongWait = c.PongWait m.Config.PingPeriod = c.PingPeriod m.Config.MaxMessageSize = c.MaxMessageSize m.Config.MessageBufferSize = c.MessageBufferSize m.Config.ConcurrentMessageHandling = c.ConcurrentMessageHandling allowed := append([]string(nil), c.AllowOrigins...) m.Upgrader.CheckOrigin = func(r *http.Request) bool { if len(allowed) == 0 { origin := r.Header.Get("Origin") if origin == "" { return true } parsed, err := url.Parse(origin) return err == nil && strings.EqualFold(parsed.Host, r.Host) } origin := r.Header.Get("Origin") for _, value := range allowed { if value == "*" || strings.EqualFold(strings.TrimSpace(value), origin) { return true } } return false } s := &Server{m: m} m.HandleMessage(func(session *melody.Session, message []byte) { for _, handler := range s.messageHandlerSnapshot() { handler(session, message) } }) m.HandleMessageBinary(func(session *melody.Session, message []byte) { for _, handler := range s.binaryHandlerSnapshot() { handler(session, message) } }) m.HandleConnect(func(session *melody.Session) { for _, handler := range s.connectHandlerSnapshot() { handler(session) } }) m.HandleDisconnect(func(session *melody.Session) { for _, handler := range s.disconnectHandlerSnapshot() { handler(session) } }) return s } func (s *Server) Melody() *melody.Melody { if s == nil { return nil } return s.m } func (s *Server) Enabled() bool { if s == nil || s.m == nil { return false } s.mu.RLock() defer s.mu.RUnlock() return !s.closed } func (s *Server) HandleRequest(w http.ResponseWriter, r *http.Request) error { return s.m.HandleRequest(w, r) } func (s *Server) HandleRequestWithKeys(w http.ResponseWriter, r *http.Request, keys map[string]any) error { return s.m.HandleRequestWithKeys(w, r, keys) } func (s *Server) Broadcast(message []byte) error { return s.m.Broadcast(message) } func (s *Server) BroadcastBinary(message []byte) error { return s.m.BroadcastBinary(message) } func (s *Server) Sessions() ([]*melody.Session, error) { return s.m.Sessions() } func (s *Server) Len() int { return s.m.Len() } func (s *Server) OnMessage(handler func(*melody.Session, []byte)) { if s == nil || handler == nil { return } s.mu.Lock() s.messageHandlers = append(s.messageHandlers, handler) s.mu.Unlock() } func (s *Server) OnBinaryMessage(handler func(*melody.Session, []byte)) { if s == nil || handler == nil { return } s.mu.Lock() s.binaryHandlers = append(s.binaryHandlers, handler) s.mu.Unlock() } func (s *Server) OnConnect(handler func(*melody.Session)) { if s == nil || handler == nil { return } s.mu.Lock() s.connectHandlers = append(s.connectHandlers, handler) s.mu.Unlock() } func (s *Server) OnDisconnect(handler func(*melody.Session)) { if s == nil || handler == nil { return } s.mu.Lock() s.disconnectHandlers = append(s.disconnectHandlers, handler) s.mu.Unlock() } func (s *Server) messageHandlerSnapshot() []func(*melody.Session, []byte) { s.mu.RLock() defer s.mu.RUnlock() handlers := make([]func(*melody.Session, []byte), len(s.messageHandlers)) copy(handlers, s.messageHandlers) return handlers } func (s *Server) binaryHandlerSnapshot() []func(*melody.Session, []byte) { s.mu.RLock() defer s.mu.RUnlock() handlers := make([]func(*melody.Session, []byte), len(s.binaryHandlers)) copy(handlers, s.binaryHandlers) return handlers } func (s *Server) connectHandlerSnapshot() []func(*melody.Session) { s.mu.RLock() defer s.mu.RUnlock() handlers := make([]func(*melody.Session), len(s.connectHandlers)) copy(handlers, s.connectHandlers) return handlers } func (s *Server) disconnectHandlerSnapshot() []func(*melody.Session) { s.mu.RLock() defer s.mu.RUnlock() handlers := make([]func(*melody.Session), len(s.disconnectHandlers)) copy(handlers, s.disconnectHandlers) return handlers } func (s *Server) Send(session *melody.Session, message []byte) error { if session == nil { return melody.ErrClosed } return session.Write(message) } func (s *Server) SendBinary(session *melody.Session, message []byte) error { if session == nil { return melody.ErrClosed } return session.WriteBinary(message) } func (s *Server) Close() error { if s == nil || s.m == nil { return nil } s.mu.Lock() if s.closed { s.mu.Unlock() return nil } s.closed = true s.mu.Unlock() return s.m.Close() }