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
}

// CreateServer for a specific HOST and PORT
func CreateServer(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 {

	// 1. Build URI parser
	uriScheme, err := uri.FromString(uriStr)
	if err != nil {
		return fmt.Errorf("Cannot build URI: %s", err)
	}

	// 2. Create controller
	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

	/* (1) Listen socket
	---------------------------------------------------------*/
	// 1. Build full url
	url := fmt.Sprintf("%s:%d", s.addr, s.port)

	// 2. Bind socket to listen
	s.sock, err = net.Listen("tcp", url)
	if err != nil {
		return fmt.Errorf("Listen socket: %s", err)
	}

	defer s.sock.Close()

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

	// 3. Launch scheduler
	go s.scheduler()

	/* (2) For each incoming connection (client)
	---------------------------------------------------------*/
	for {

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

		go func() {

			// 2. Try to create client
			cli, err := buildClient(sock, s.ctl, s.ch)
			if err != nil {
				fmt.Printf(" - %s\n", err)
				return
			}

			// 3. Register client
			s.ch.register <- cli

		}()

	}

	return nil

}

// Scheduler schedules clients registration and broadcast
func (s *Server) scheduler() {

	for {

		select {

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

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

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

		}

	}

}