proto-socket/go/ws_server.go

package toki_socket

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"net/http"
	"sync"

	"google.golang.org/protobuf/proto"
	"nhooyr.io/websocket"
)

type WsServer struct {
	host              string
	port              int
	path              string
	tlsConfig         *tls.Config
	newClient         func(*websocket.Conn) *WsClient
	clients           []*WsClient
	mu                sync.Mutex
	httpSrv           *http.Server
	listener          net.Listener
	started           bool
	OnClientConnected func(*WsClient)
}

func NewWsServer(host string, port int, path string, newClient func(*websocket.Conn) *WsClient) *WsServer {
	return &WsServer{
		host:              host,
		port:              port,
		path:              path,
		newClient:         newClient,
		OnClientConnected: func(*WsClient) {},
	}
}

func NewWsServerTLS(host string, port int, path string, tlsCfg *tls.Config, newClient func(*websocket.Conn) *WsClient) *WsServer {
	s := NewWsServer(host, port, path, newClient)
	s.tlsConfig = tlsCfg
	return s
}

func (s *WsServer) Started() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.started
}

func (s *WsServer) IsSecure() bool {
	return s.tlsConfig != nil
}

func (s *WsServer) Clients() []*WsClient {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]*WsClient{}, s.clients...)
}

func (s *WsServer) Start(ctx context.Context) error {
	mux := http.NewServeMux()
	mux.HandleFunc(s.path, s.handleWebSocket)
	s.httpSrv = &http.Server{
		Addr:      fmt.Sprintf("%s:%d", s.host, s.port),
		Handler:   mux,
		TLSConfig: s.tlsConfig,
	}

	ln, err := net.Listen("tcp", s.httpSrv.Addr)
	if err != nil {
		return err
	}
	if s.tlsConfig != nil {
		ln = tls.NewListener(ln, s.tlsConfig)
	}

	s.mu.Lock()
	s.listener = ln
	s.started = true
	s.mu.Unlock()

	go func() {
		<-ctx.Done()
		_ = s.Stop()
	}()
	go func() {
		err := s.httpSrv.Serve(ln)
		if err != nil && !errors.Is(err, http.ErrServerClosed) && !errors.Is(err, net.ErrClosed) {
			_ = s.Stop()
		}
	}()
	return nil
}

func (s *WsServer) Stop() error {
	s.mu.Lock()
	if !s.started {
		s.mu.Unlock()
		return nil
	}
	s.started = false
	srv := s.httpSrv
	clients := append([]*WsClient{}, s.clients...)
	s.clients = nil
	s.mu.Unlock()

	var err error
	if srv != nil {
		err = srv.Close()
	}
	for _, client := range clients {
		_ = client.Close()
	}
	return err
}

func (s *WsServer) Broadcast(m proto.Message) error {
	clients := s.Clients()
	var errs []error
	for _, client := range clients {
		if err := client.Send(m); err != nil {
			errs = append(errs, err)
		}
	}
	return errors.Join(errs...)
}

func (s *WsServer) handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := websocket.Accept(w, r, nil)
	if err != nil {
		return
	}
	client := s.newClient(conn)
	client.AddDisconnectListener(func(c *WsClient) {
		s.removeClient(c)
	})

	s.mu.Lock()
	s.clients = append(s.clients, client)
	onConnected := s.OnClientConnected
	s.mu.Unlock()
	onConnected(client)
}

func (s *WsServer) removeClient(client *WsClient) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i, item := range s.clients {
		if item == client {
			s.clients = append(s.clients[:i], s.clients[i+1:]...)
			return
		}
	}
}