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