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