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() }