1package melody 2 3import ( 4 "errors" 5 "net/http" 6 "sync" 7 8 "github.com/gorilla/websocket" 9) 10 11// Close codes defined in RFC 6455, section 11.7. 12// Duplicate of codes from gorilla/websocket for convenience. 13const ( 14 CloseNormalClosure = 1000 15 CloseGoingAway = 1001 16 CloseProtocolError = 1002 17 CloseUnsupportedData = 1003 18 CloseNoStatusReceived = 1005 19 CloseAbnormalClosure = 1006 20 CloseInvalidFramePayloadData = 1007 21 ClosePolicyViolation = 1008 22 CloseMessageTooBig = 1009 23 CloseMandatoryExtension = 1010 24 CloseInternalServerErr = 1011 25 CloseServiceRestart = 1012 26 CloseTryAgainLater = 1013 27 CloseTLSHandshake = 1015 28) 29 30// Duplicate of codes from gorilla/websocket for convenience. 31var validReceivedCloseCodes = map[int]bool{ 32 // see http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number 33 34 CloseNormalClosure: true, 35 CloseGoingAway: true, 36 CloseProtocolError: true, 37 CloseUnsupportedData: true, 38 CloseNoStatusReceived: false, 39 CloseAbnormalClosure: false, 40 CloseInvalidFramePayloadData: true, 41 ClosePolicyViolation: true, 42 CloseMessageTooBig: true, 43 CloseMandatoryExtension: true, 44 CloseInternalServerErr: true, 45 CloseServiceRestart: true, 46 CloseTryAgainLater: true, 47 CloseTLSHandshake: false, 48} 49 50type handleMessageFunc func(*Session, []byte) 51type handleErrorFunc func(*Session, error) 52type handleCloseFunc func(*Session, int, string) error 53type handleSessionFunc func(*Session) 54type filterFunc func(*Session) bool 55 56// Melody implements a websocket manager. 57type Melody struct { 58 Config *Config 59 Upgrader *websocket.Upgrader 60 messageHandler handleMessageFunc 61 messageHandlerBinary handleMessageFunc 62 messageSentHandler handleMessageFunc 63 messageSentHandlerBinary handleMessageFunc 64 errorHandler handleErrorFunc 65 closeHandler handleCloseFunc 66 connectHandler handleSessionFunc 67 disconnectHandler handleSessionFunc 68 pongHandler handleSessionFunc 69 hub *hub 70} 71 72// New creates a new melody instance with default Upgrader and Config. 73func New() *Melody { 74 upgrader := &websocket.Upgrader{ 75 ReadBufferSize: 1024, 76 WriteBufferSize: 1024, 77 CheckOrigin: func(r *http.Request) bool { return true }, 78 } 79 80 hub := newHub() 81 82 go hub.run() 83 84 return &Melody{ 85 Config: newConfig(), 86 Upgrader: upgrader, 87 messageHandler: func(*Session, []byte) {}, 88 messageHandlerBinary: func(*Session, []byte) {}, 89 messageSentHandler: func(*Session, []byte) {}, 90 messageSentHandlerBinary: func(*Session, []byte) {}, 91 errorHandler: func(*Session, error) {}, 92 closeHandler: nil, 93 connectHandler: func(*Session) {}, 94 disconnectHandler: func(*Session) {}, 95 pongHandler: func(*Session) {}, 96 hub: hub, 97 } 98} 99 100// HandleConnect fires fn when a session connects. 101func (m *Melody) HandleConnect(fn func(*Session)) { 102 m.connectHandler = fn 103} 104 105// HandleDisconnect fires fn when a session disconnects. 106func (m *Melody) HandleDisconnect(fn func(*Session)) { 107 m.disconnectHandler = fn 108} 109 110// HandlePong fires fn when a pong is received from a session. 111func (m *Melody) HandlePong(fn func(*Session)) { 112 m.pongHandler = fn 113} 114 115// HandleMessage fires fn when a text message comes in. 116func (m *Melody) HandleMessage(fn func(*Session, []byte)) { 117 m.messageHandler = fn 118} 119 120// HandleMessageBinary fires fn when a binary message comes in. 121func (m *Melody) HandleMessageBinary(fn func(*Session, []byte)) { 122 m.messageHandlerBinary = fn 123} 124 125// HandleSentMessage fires fn when a text message is successfully sent. 126func (m *Melody) HandleSentMessage(fn func(*Session, []byte)) { 127 m.messageSentHandler = fn 128} 129 130// HandleSentMessageBinary fires fn when a binary message is successfully sent. 131func (m *Melody) HandleSentMessageBinary(fn func(*Session, []byte)) { 132 m.messageSentHandlerBinary = fn 133} 134 135// HandleError fires fn when a session has an error. 136func (m *Melody) HandleError(fn func(*Session, error)) { 137 m.errorHandler = fn 138} 139 140// HandleClose sets the handler for close messages received from the session. 141// The code argument to h is the received close code or CloseNoStatusReceived 142// if the close message is empty. The default close handler sends a close frame 143// back to the session. 144// 145// The application must read the connection to process close messages as 146// described in the section on Control Frames above. 147// 148// The connection read methods return a CloseError when a close frame is 149// received. Most applications should handle close messages as part of their 150// normal error handling. Applications should only set a close handler when the 151// application must perform some action before sending a close frame back to 152// the session. 153func (m *Melody) HandleClose(fn func(*Session, int, string) error) { 154 if fn != nil { 155 m.closeHandler = fn 156 } 157} 158 159// HandleRequest upgrades http requests to websocket connections and dispatches them to be handled by the melody instance. 160func (m *Melody) HandleRequest(w http.ResponseWriter, r *http.Request) error { 161 return m.HandleRequestWithKeys(w, r, nil) 162} 163 164// HandleRequestWithKeys does the same as HandleRequest but populates session.Keys with keys. 165func (m *Melody) HandleRequestWithKeys(w http.ResponseWriter, r *http.Request, keys map[string]interface{}) error { 166 if m.hub.closed() { 167 return errors.New("melody instance is closed") 168 } 169 170 conn, err := m.Upgrader.Upgrade(w, r, nil) 171 172 if err != nil { 173 return err 174 } 175 176 session := &Session{ 177 Request: r, 178 Keys: keys, 179 conn: conn, 180 output: make(chan *envelope, m.Config.MessageBufferSize), 181 melody: m, 182 open: true, 183 rwmutex: &sync.RWMutex{}, 184 } 185 186 m.hub.register <- session 187 188 m.connectHandler(session) 189 190 go session.writePump() 191 192 session.readPump() 193 194 if !m.hub.closed() { 195 m.hub.unregister <- session 196 } 197 198 session.close() 199 200 m.disconnectHandler(session) 201 202 return nil 203} 204 205// Broadcast broadcasts a text message to all sessions. 206func (m *Melody) Broadcast(msg []byte) error { 207 if m.hub.closed() { 208 return errors.New("melody instance is closed") 209 } 210 211 message := &envelope{t: websocket.TextMessage, msg: msg} 212 m.hub.broadcast <- message 213 214 return nil 215} 216 217// BroadcastFilter broadcasts a text message to all sessions that fn returns true for. 218func (m *Melody) BroadcastFilter(msg []byte, fn func(*Session) bool) error { 219 if m.hub.closed() { 220 return errors.New("melody instance is closed") 221 } 222 223 message := &envelope{t: websocket.TextMessage, msg: msg, filter: fn} 224 m.hub.broadcast <- message 225 226 return nil 227} 228 229// BroadcastOthers broadcasts a text message to all sessions except session s. 230func (m *Melody) BroadcastOthers(msg []byte, s *Session) error { 231 return m.BroadcastFilter(msg, func(q *Session) bool { 232 return s != q 233 }) 234} 235 236// BroadcastMultiple broadcasts a text message to multiple sessions given in the sessions slice. 237func (m *Melody) BroadcastMultiple(msg []byte, sessions []*Session) error { 238 for _, sess := range sessions { 239 if writeErr := sess.Write(msg); writeErr != nil { 240 return writeErr 241 } 242 } 243 return nil 244} 245 246// BroadcastBinary broadcasts a binary message to all sessions. 247func (m *Melody) BroadcastBinary(msg []byte) error { 248 if m.hub.closed() { 249 return errors.New("melody instance is closed") 250 } 251 252 message := &envelope{t: websocket.BinaryMessage, msg: msg} 253 m.hub.broadcast <- message 254 255 return nil 256} 257 258// BroadcastBinaryFilter broadcasts a binary message to all sessions that fn returns true for. 259func (m *Melody) BroadcastBinaryFilter(msg []byte, fn func(*Session) bool) error { 260 if m.hub.closed() { 261 return errors.New("melody instance is closed") 262 } 263 264 message := &envelope{t: websocket.BinaryMessage, msg: msg, filter: fn} 265 m.hub.broadcast <- message 266 267 return nil 268} 269 270// BroadcastBinaryOthers broadcasts a binary message to all sessions except session s. 271func (m *Melody) BroadcastBinaryOthers(msg []byte, s *Session) error { 272 return m.BroadcastBinaryFilter(msg, func(q *Session) bool { 273 return s != q 274 }) 275} 276 277// Close closes the melody instance and all connected sessions. 278func (m *Melody) Close() error { 279 if m.hub.closed() { 280 return errors.New("melody instance is already closed") 281 } 282 283 m.hub.exit <- &envelope{t: websocket.CloseMessage, msg: []byte{}} 284 285 return nil 286} 287 288// CloseWithMsg closes the melody instance with the given close payload and all connected sessions. 289// Use the FormatCloseMessage function to format a proper close message payload. 290func (m *Melody) CloseWithMsg(msg []byte) error { 291 if m.hub.closed() { 292 return errors.New("melody instance is already closed") 293 } 294 295 m.hub.exit <- &envelope{t: websocket.CloseMessage, msg: msg} 296 297 return nil 298} 299 300// Len return the number of connected sessions. 301func (m *Melody) Len() int { 302 return m.hub.len() 303} 304 305// IsClosed returns the status of the melody instance. 306func (m *Melody) IsClosed() bool { 307 return m.hub.closed() 308} 309 310// FormatCloseMessage formats closeCode and text as a WebSocket close message. 311func FormatCloseMessage(closeCode int, text string) []byte { 312 return websocket.FormatCloseMessage(closeCode, text) 313} 314