kra-new/pkg/websocket/melody.go

// 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 snapshotHandlers[T any](mu *sync.RWMutex, handlers []T) []T {
	mu.RLock()
	defer mu.RUnlock()
	return append([]T(nil), handlers...)
}

func (s *Server) messageHandlerSnapshot() []func(*melody.Session, []byte) {
	return snapshotHandlers(&s.mu, s.messageHandlers)
}

func (s *Server) binaryHandlerSnapshot() []func(*melody.Session, []byte) {
	return snapshotHandlers(&s.mu, s.binaryHandlers)
}

func (s *Server) connectHandlerSnapshot() []func(*melody.Session) {
	return snapshotHandlers(&s.mu, s.connectHandlers)
}

func (s *Server) disconnectHandlerSnapshot() []func(*melody.Session) {
	return snapshotHandlers(&s.mu, s.disconnectHandlers)
}
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()
}