utils/websocket/websocket.go

package websocket

import (
	"github.com/olahol/melody"
	"net/http"
)

const (
	groupIDKey = "group-id"
)

type HandleConnectFunc func(groupID string)
type HandleDisconnectFunc func(groupID string)
type HandleErrorFunc func(groupID string, err error)
type HandleCloseFunc func(groupID string, i int, s string) error
type HandlePongFunc func(groupID string)
type HandleMessageFunc func(groupID string, message []byte)

var managerInstance *Manager

func Init() {
	if managerInstance == nil {
		melodyInstance := melody.New()
		melodyInstance.Upgrader.CheckOrigin = func(r *http.Request) bool { return true }
		managerInstance = &Manager{melodyInstance: melodyInstance}
	}
}

func Destroy() {
	if managerInstance != nil {
		err := managerInstance.melodyInstance.Close()
		if err != nil {
			panic(err)
		}
	}

	managerInstance = nil
}

func GetInstance() *Manager {
	return managerInstance
}

type Manager struct {
	melodyInstance *melody.Melody
}

func (m *Manager) HandleConnect(handleConnectFunc HandleConnectFunc) {
	m.melodyInstance.HandleConnect(func(session *melody.Session) {
		if handleConnectFunc != nil {
			handleConnectFunc(session.Keys[groupIDKey].(string))
		}
	})
}

func (m *Manager) HandleDisconnect(handleDisconnectFunc HandleDisconnectFunc) {
	m.melodyInstance.HandleDisconnect(func(session *melody.Session) {
		if handleDisconnectFunc != nil {
			handleDisconnectFunc(session.Keys[groupIDKey].(string))
		}
	})
}

func (m *Manager) HandleError(handleErrorFunc HandleErrorFunc) {
	m.melodyInstance.HandleError(func(session *melody.Session, err error) {
		if handleErrorFunc != nil {
			handleErrorFunc(session.Keys[groupIDKey].(string), err)
		}
	})
}

func (m *Manager) HandleClose(handleCloseFunc HandleCloseFunc) {
	m.melodyInstance.HandleClose(func(session *melody.Session, i int, s string) error {
		if handleCloseFunc != nil {
			err := handleCloseFunc(session.Keys[groupIDKey].(string), i, s)
			if err != nil {
				return err
			}
		}

		return nil
	})
}

func (m *Manager) HandlePong(handlePongFunc HandlePongFunc) {
	m.melodyInstance.HandlePong(func(session *melody.Session) {
		if handlePongFunc != nil {
			handlePongFunc(session.Keys[groupIDKey].(string))
		}
	})
}

func (m *Manager) HandleRequest(groupID string, w http.ResponseWriter, r *http.Request) error {
	err := m.melodyInstance.HandleRequestWithKeys(w, r, map[string]interface{}{groupIDKey: groupID})
	if err != nil {
		return err
	}

	return nil
}

func (m *Manager) HandleMessage(handleMessageFunc HandleMessageFunc) {
	m.melodyInstance.HandleMessage(func(session *melody.Session, bytes []byte) {
		if handleMessageFunc != nil {
			handleMessageFunc(session.Keys[groupIDKey].(string), bytes)
		}
	})
}

func (m *Manager) BroadCast(groupID string, msg []byte) error {
	return m.melodyInstance.BroadcastFilter(msg, func(session *melody.Session) bool {
		if session.Keys[groupIDKey] != groupID {
			return false
		}

		return true
	})
}