kra-new/app/system/internal/integration/websocket/server.go

package websocket

import (
	"errors"
	"net/http"
	"strings"
	"sync"
	"time"

	melody "github.com/olahol/melody"
	"kra/app/system/internal/conf"
	platformws "kra/pkg/websocket"
)

// Server is the system-owned WebSocket endpoint. Business modules can attach
// handlers and publish messages without depending on Gin or Melody directly.
type Server struct {
	mu                 sync.RWMutex
	current            *platformws.Server
	path               string
	messageHandlers    []func(*melody.Session, []byte)
	binaryHandlers     []func(*melody.Session, []byte)
	connectHandlers    []func(*melody.Session)
	disconnectHandlers []func(*melody.Session)
	closed             bool
}

func New(runtime *conf.Runtime) (*Server, func(), error) {
	s := &Server{}
	s.Replace(runtime)
	var unsubscribe func()
	if runtime != nil {
		unsubscribe = runtime.Subscribe(func(_ *conf.Data, _ *conf.AdminBackend) { s.Replace(runtime) })
	}
	return s, func() {
		if unsubscribe != nil {
			unsubscribe()
		}
		s.mu.Lock()
		s.closed = true
		current := s.current
		s.current = nil
		s.mu.Unlock()
		if current != nil {
			_ = current.Close()
		}
	}, nil
}

func (s *Server) Replace(runtime *conf.Runtime) {
	var config *conf.AdminBackend_WebSocket
	if runtime != nil && runtime.Admin() != nil {
		config = runtime.Admin().Websocket
	}
	if config == nil {
		config = &conf.AdminBackend_WebSocket{}
	}
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return
	}
	previous := s.current
	if !config.Enabled {
		s.current = nil
		s.path = ""
		s.mu.Unlock()
		if previous != nil {
			_ = previous.Close()
		}
		return
	}
	next := platformws.New(platformws.Config{
		WriteWait:                 duration(config.WriteWait, 10*time.Second),
		PongWait:                  duration(config.PongWait, 60*time.Second),
		PingPeriod:                duration(config.PingPeriod, 54*time.Second),
		MaxMessageSize:            config.MaxMessageSize,
		MessageBufferSize:         int(config.MessageBufferSize),
		ConcurrentMessageHandling: config.ConcurrentMessageHandling,
		AllowOrigins:              config.AllowOrigins,
	})
	for _, handler := range s.messageHandlers {
		next.OnMessage(handler)
	}
	for _, handler := range s.binaryHandlers {
		next.OnBinaryMessage(handler)
	}
	for _, handler := range s.connectHandlers {
		next.OnConnect(handler)
	}
	for _, handler := range s.disconnectHandlers {
		next.OnDisconnect(handler)
	}
	s.path = strings.TrimSpace(config.Path)
	if s.path == "" {
		s.path = "/ws"
	} else if !strings.HasPrefix(s.path, "/") {
		s.path = "/" + s.path
	}
	s.current = next
	s.mu.Unlock()
	if previous != nil {
		_ = previous.Close()
	}
}

func duration(value interface{ AsDuration() time.Duration }, fallback time.Duration) time.Duration {
	if value == nil {
		return fallback
	}
	if result := value.AsDuration(); result > 0 {
		return result
	}
	return fallback
}

func (s *Server) Enabled() bool {
	if s == nil {
		return false
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.current != nil
}
func (s *Server) Path() string {
	if s == nil {
		return ""
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.path
}
func (s *Server) HandleRequest(w http.ResponseWriter, r *http.Request) error {
	if s == nil {
		return errors.New("websocket server is disabled")
	}
	s.mu.RLock()
	current := s.current
	s.mu.RUnlock()
	if current == nil {
		return errors.New("websocket server is disabled")
	}
	return current.HandleRequest(w, r)
}
func (s *Server) HandleRequestWithKeys(w http.ResponseWriter, r *http.Request, keys map[string]any) error {
	current, err := s.active()
	if err != nil {
		return err
	}
	return current.HandleRequestWithKeys(w, r, keys)
}
func (s *Server) Broadcast(message []byte) error {
	if s == nil {
		return errors.New("websocket server is disabled")
	}
	s.mu.RLock()
	current := s.current
	s.mu.RUnlock()
	if current == nil {
		return errors.New("websocket server is disabled")
	}
	return current.Broadcast(message)
}
func (s *Server) BroadcastBinary(message []byte) error {
	if s == nil {
		return errors.New("websocket server is disabled")
	}
	s.mu.RLock()
	current := s.current
	s.mu.RUnlock()
	if current == nil {
		return errors.New("websocket server is disabled")
	}
	return current.BroadcastBinary(message)
}
func (s *Server) Sessions() ([]*melody.Session, error) {
	current, err := s.active()
	if err != nil {
		return nil, err
	}
	return current.Sessions()
}
func (s *Server) Len() int {
	current, err := s.active()
	if err != nil {
		return 0
	}
	return current.Len()
}
func (s *Server) Send(session *melody.Session, message []byte) error {
	current, err := s.active()
	if err != nil {
		return err
	}
	return current.Send(session, message)
}
func (s *Server) SendBinary(session *melody.Session, message []byte) error {
	current, err := s.active()
	if err != nil {
		return err
	}
	return current.SendBinary(session, message)
}
func (s *Server) active() (*platformws.Server, error) {
	if s == nil {
		return nil, errors.New("websocket server is disabled")
	}
	s.mu.RLock()
	current := s.current
	s.mu.RUnlock()
	if current == nil {
		return nil, errors.New("websocket server is disabled")
	}
	return current, nil
}
func (s *Server) OnMessage(handler func(*melody.Session, []byte)) {
	if s == nil || handler == nil {
		return
	}
	s.mu.Lock()
	s.messageHandlers = append(s.messageHandlers, handler)
	current := s.current
	if current != nil {
		current.OnMessage(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)
	if s.current != nil {
		s.current.OnBinaryMessage(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)
	if s.current != nil {
		s.current.OnConnect(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)
	if s.current != nil {
		s.current.OnDisconnect(handler)
	}
	s.mu.Unlock()
}