utils/websocket/websocket.go

package websocket

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

const (
	groupIDKey           = "group-id"
	connectionContextKey = "connection-context"
)

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

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), session.Keys[connectionContextKey])
		}
	})
}

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

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

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, session.Keys[connectionContextKey])
			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), session.Keys[connectionContextKey])
		}
	})
}

func (m *Manager) HandleRequest(groupID string, w http.ResponseWriter, r *http.Request, opts ...ConnectionOption) error {
	sessionMap := map[string]interface{}{
		groupIDKey: groupID,
	}

	for _, opt := range opts {
		opt(sessionMap)
	}

	err := m.melodyInstance.HandleRequestWithKeys(w, r, sessionMap)
	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, session.Keys[connectionContextKey])
		}
	})
}

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
	})
}

type ConnectionOption func(sessionMap map[string]any)

func WithConnectionContext(context any) ConnectionOption {
	return func(sessionMap map[string]any) {
		sessionMap[connectionContextKey] = context
	}
}