overflow_service_websocket/server/server.go

package server

import (
	"fmt"
	"log"
	"net/http"
	"strings"
	"sync"

	"git.loafle.net/overflow/overflow_service_websocket/protocol"
	"github.com/gorilla/websocket"
)

type (
	// OnConnectionFunc is callback function that used when client is connected
	OnConnectionFunc func(Client)
)

// Server is the websocket server,
// listens on the config's port, the critical part is the event OnConnection
type Server interface {
	GetOptions() *Options
	RegistProtocol(protocol string, h protocol.Handler)
	ProtocolHandler(protocol string) protocol.Handler
	HTTPHandler() http.Handler
	HandleConnection(*http.Request, Connection)
	OnConnection(cb OnConnectionFunc)
	IsConnected(clientID string) bool
	GetClient(clientID string) Client
	Disconnect(clientID string) error
}

type server struct {
	options               *Options
	clients               map[string]Client
	clientMTX             sync.Mutex
	onConnectionListeners []OnConnectionFunc
	protocols             map[string]protocol.Handler
}

var _ Server = &server{}

var defaultServer = newServer(nil)

// server implementation

// New creates a websocket server and returns it
func New(o *Options) Server {
	return newServer(o)
}

// newServer creates a websocket server and returns it
func newServer(o *Options) Server {

	s := &server{
		clients:               make(map[string]Client, 100),
		onConnectionListeners: make([]OnConnectionFunc, 0),
		protocols:             make(map[string]protocol.Handler, 0),
	}

	if nil == o {
		o = &Options{}
	}

	o.Validate()

	s.options = o
	return s
}

func RegistProtocol(protocol string, h protocol.Handler) {
	defaultServer.RegistProtocol(protocol, h)
}

func (s *server) RegistProtocol(protocol string, h protocol.Handler) {
	s.protocols[strings.ToLower(protocol)] = h
}

func ProtocolHandler(protocol string) protocol.Handler {
	return defaultServer.ProtocolHandler(protocol)
}

func (s *server) ProtocolHandler(protocol string) protocol.Handler {
	return s.protocols[protocol]
}

// GetOptions is function that get option values
func GetOptions() *Options {
	return defaultServer.GetOptions()
}

func (s *server) GetOptions() *Options {
	return s.options
}

// Handler is the function that used on http request
func HTTPHandler() http.Handler {
	return defaultServer.HTTPHandler()
}

func (s *server) HTTPHandler() http.Handler {
	o := s.options

	upgrader := websocket.Upgrader{
		ReadBufferSize:  o.ReadBufferSize,
		WriteBufferSize: o.WriteBufferSize,
		Error:           o.OnError,
		CheckOrigin:     o.OnCheckOrigin,
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, w.Header())
		if err != nil {
			http.Error(w, "Websocket Error: "+err.Error(), http.StatusServiceUnavailable)
			return
		}
		s.HandleConnection(r, conn)
	})
}

func HandleConnection(r *http.Request, conn Connection) {
	defaultServer.HandleConnection(r, conn)
}

func (s *server) HandleConnection(r *http.Request, conn Connection) {
	clientID := s.options.IDGenerator(r)
	c := newClient(s, r, conn, clientID)
	err := s.addClient(clientID, c)
	if nil != err {
		log.Println(fmt.Errorf("%v", err))
		return
	}

	for i := range s.onConnectionListeners {
		s.onConnectionListeners[i](c)
	}

	c.initialize()
}

// OnConnection is function that add the callback when client is connected to default Server
func OnConnection(cb OnConnectionFunc) {
	defaultServer.OnConnection(cb)
}

func (s *server) OnConnection(cb OnConnectionFunc) {
	s.onConnectionListeners = append(s.onConnectionListeners, cb)
}

// IsConnected is function that check client is connect
func IsConnected(clientID string) bool {
	return defaultServer.IsConnected(clientID)
}

func (s *server) IsConnected(clientID string) bool {
	soc := s.clients[clientID]
	return soc != nil
}

// GetClient is function that return client instance
func GetClient(clientID string) Client {
	return defaultServer.GetClient(clientID)
}

func (s *server) GetClient(clientID string) Client {
	return s.clients[clientID]
}

// Disconnect is function that disconnect a client
func Disconnect(clientID string) error {
	return defaultServer.Disconnect(clientID)
}

func (s *server) Disconnect(clientID string) error {
	c := s.clients[clientID]

	if nil == c {
		return nil
	}

	delete(s.clients, clientID)

	return c.destroy()
}

func (s *server) addClient(clientID string, c Client) error {
	s.clientMTX.Lock()
	if s.clients[clientID] != nil {
		return fmt.Errorf("ID of Client[%s] is exist already", clientID)
	}
	s.clients[clientID] = c
	s.clientMTX.Unlock()
	return nil
}