utils/websocket/websocket.go

package websocket

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

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(opts ...InitOption) {
	if managerInstance == nil {
		melodyInstance := melody.New()

		options := new(InitOptions)
		for _, opt := range opts {
			opt(options)
		}

		if options.writeWaitSec != 0 {
			melodyInstance.Config.WriteWait = time.Duration(options.writeWaitSec) * time.Second
		}

		if options.pongWaitSec != 0 {
			melodyInstance.Config.PongWait = time.Duration(options.pongWaitSec) * time.Second
		}

		if options.pingPeriodSec != 0 {
			melodyInstance.Config.PingPeriod = time.Duration(options.pingPeriodSec) * time.Second
		}

		if options.maxMessageSize != 0 {
			melodyInstance.Config.MaxMessageSize = options.maxMessageSize
		}

		if options.messageBufferSize != 0 {
			melodyInstance.Config.MessageBufferSize = options.messageBufferSize
		}

		melodyInstance.Config.ConcurrentMessageHandling = options.concurrentMessageHandling

		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 InitOption func(*InitOptions)

type InitOptions struct {
	writeWaitSec              int64
	pongWaitSec               int64
	pingPeriodSec             int64
	maxMessageSize            int64
	messageBufferSize         int
	concurrentMessageHandling bool
}

func InitWithWriteWaitSec(writeWaitSec int64) InitOption {
	return func(options *InitOptions) {
		options.writeWaitSec = writeWaitSec
	}
}

func InitWithPongWaitSec(pongWaitSec int64) InitOption {
	return func(options *InitOptions) {
		options.pongWaitSec = pongWaitSec
	}
}

func InitWithPingPeriodSec(pingPeriodSec int64) InitOption {
	return func(options *InitOptions) {
		options.pingPeriodSec = pingPeriodSec
	}
}

func InitWithMaxMessageSize(maxMessageSize int64) InitOption {
	return func(options *InitOptions) {
		options.maxMessageSize = maxMessageSize
	}
}

func InitWithMaxMessageBufferSize(messageBufferSize int) InitOption {
	return func(options *InitOptions) {
		options.messageBufferSize = messageBufferSize
	}
}

func InitWithConcurrentMessageHandling(concurrentMessageHandling bool) InitOption {
	return func(options *InitOptions) {
		options.concurrentMessageHandling = concurrentMessageHandling
	}
}

type ConnectionOption func(sessionMap map[string]any)

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