melody/melody.go

package melody

import (
	"github.com/gorilla/websocket"
	"net/http"
)

type handleMessageFunc func(*Session, []byte)
type handleErrorFunc func(*Session, error)
type handleSessionFunc func(*Session)
type filterFunc func(*Session) bool

type Melody struct {
	Config            *Config
	Upgrader          *websocket.Upgrader
	MessageHandler    handleMessageFunc
	ErrorHandler      handleErrorFunc
	ConnectHandler    handleSessionFunc
	DisconnectHandler handleSessionFunc
	hub               *hub
}

func Default() *Melody {
	upgrader := &websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}

	hub := newHub()

	go hub.run()

	return &Melody{
		Config:            newConfig(),
		Upgrader:          upgrader,
		MessageHandler:    func(*Session, []byte) {},
		ErrorHandler:      func(*Session, error) {},
		ConnectHandler:    func(*Session) {},
		DisconnectHandler: func(*Session) {},
		hub:               hub,
	}
}

func (m *Melody) HandleConnect(fn handleSessionFunc) {
	m.ConnectHandler = fn
}

func (m *Melody) HandleDisconnect(fn handleSessionFunc) {
	m.DisconnectHandler = fn
}

func (m *Melody) HandleMessage(fn handleMessageFunc) {
	m.MessageHandler = fn
}

func (m *Melody) HandleError(fn handleErrorFunc) {
	m.ErrorHandler = fn
}

func (m *Melody) HandleRequest(w http.ResponseWriter, r *http.Request) error {
	conn, err := m.Upgrader.Upgrade(w, r, nil)

	if err != nil {
		return err
	}

	session := newSession(m.Config, conn)

	m.hub.register <- session

	go m.ConnectHandler(session)

	go session.writePump(m.ErrorHandler)

	session.readPump(m.MessageHandler, m.ErrorHandler)

	m.hub.unregister <- session

	go m.DisconnectHandler(session)

	return nil
}

func (m *Melody) Broadcast(msg []byte) {
	message := &envelope{t: websocket.TextMessage, msg: msg}
	m.hub.broadcast <- message
}

func (m *Melody) BroadcastFilter(fn filterFunc, msg []byte) {
	message := &envelope{t: websocket.TextMessage, msg: msg, filter: fn}
	m.hub.broadcast <- message
}