ws/server.go

package websocket

import (
	"fmt"
	"net"

	"git.xdrm.io/go/ws/internal/uri"
)

// All channels that a server features
type serverChannelSet struct {
	register   chan *client
	unregister chan *client
	broadcast  chan Message
}

// Server is a websocket server
type Server struct {
	sock net.Listener // listen socket
	addr []byte       // server listening ip/host
	port uint16       // server listening port

	clients map[net.Conn]*client

	ctl ControllerSet // controllers

	ch serverChannelSet
}

// NewServer creates a server
func NewServer(host string, port uint16) *Server {
	return &Server{
		addr: []byte(host),
		port: port,

		clients: make(map[net.Conn]*client, 0),

		ctl: ControllerSet{
			Def: nil,
			URI: make([]*Controller, 0),
		},

		ch: serverChannelSet{
			register:   make(chan *client, 1),
			unregister: make(chan *client, 1),
			broadcast:  make(chan Message, 1),
		},
	}
}

// BindDefault binds a default controller
// it will be called if the URI does not
// match another controller
func (s *Server) BindDefault(f ControllerFunc) {
	s.ctl.Def = &Controller{
		URI: nil,
		Fun: f,
	}
}

// Bind a controller to an URI scheme
func (s *Server) Bind(uriStr string, f ControllerFunc) error {
	uriScheme, err := uri.FromString(uriStr)
	if err != nil {
		return fmt.Errorf("cannot build URI: %w", err)
	}

	s.ctl.URI = append(s.ctl.URI, &Controller{
		URI: uriScheme,
		Fun: f,
	})
	return nil

}

// Launch the websocket server
func (s *Server) Launch() error {
	var (
		err error
		url = fmt.Sprintf("%s:%d", s.addr, s.port)
	)

	s.sock, err = net.Listen("tcp", url)
	if err != nil {
		return fmt.Errorf("listen: %w", err)
	}
	defer s.sock.Close()

	fmt.Printf("+ listening on %s\n", url)
	go s.schedule()

	for {
		sock, err := s.sock.Accept()
		if err != nil {
			break
		}

		go func() {
			cli, err := newClient(sock, s.ctl, s.ch)
			if err != nil {
				fmt.Printf(" - %s\n", err)
				return
			}

			s.ch.register <- cli
		}()
	}
	return nil
}

// schedule client registration and broadcast
func (s *Server) schedule() {
	for {
		select {

		case client := <-s.ch.register:
			s.clients[client.io.sock] = client

		case client := <-s.ch.unregister:
			delete(s.clients, client.io.sock)

		case msg := <-s.ch.broadcast:
			for _, c := range s.clients {
				c.ch.send <- msg
			}

		}
	}
}