websocket 增加多分组 fork https://github.com/olahol/melody
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
92 lines
1.9 KiB
92 lines
1.9 KiB
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
|
|
}
|