diff --git a/hub.go b/hub.go index fba1653..edc6337 100644 --- a/hub.go +++ b/hub.go @@ -9,7 +9,7 @@ type hub struct { broadcast chan *envelope register chan *Session unregister chan *Session - exit chan bool + exit chan *envelope open bool rwmutex *sync.RWMutex } @@ -20,7 +20,7 @@ func newHub() *hub { broadcast: make(chan *envelope), register: make(chan *Session), unregister: make(chan *Session), - exit: make(chan bool), + exit: make(chan *envelope), open: true, rwmutex: &sync.RWMutex{}, } @@ -52,9 +52,10 @@ loop: } } h.rwmutex.RUnlock() - case <-h.exit: + case m := <-h.exit: h.rwmutex.Lock() for s := range h.sessions { + s.writeMessage(m) delete(h.sessions, s) s.Close() } diff --git a/melody.go b/melody.go index 4e838d6..3c16334 100644 --- a/melody.go +++ b/melody.go @@ -2,9 +2,29 @@ package melody import ( "errors" - "github.com/gorilla/websocket" "net/http" "sync" + + "github.com/gorilla/websocket" +) + +// Close codes defined in RFC 6455, section 11.7. +// Duplicate of codes from gorilla/websocket for convenience. +const ( + CloseNormalClosure = 1000 + CloseGoingAway = 1001 + CloseProtocolError = 1002 + CloseUnsupportedData = 1003 + CloseNoStatusReceived = 1005 + CloseAbnormalClosure = 1006 + CloseInvalidFramePayloadData = 1007 + ClosePolicyViolation = 1008 + CloseMessageTooBig = 1009 + CloseMandatoryExtension = 1010 + CloseInternalServerErr = 1011 + CloseServiceRestart = 1012 + CloseTryAgainLater = 1013 + CloseTLSHandshake = 1015 ) type handleMessageFunc func(*Session, []byte) @@ -207,7 +227,19 @@ func (m *Melody) Close() error { return errors.New("Melody instance is already closed.") } - m.hub.exit <- true + m.hub.exit <- &envelope{t: websocket.CloseMessage, msg: []byte{}} + + return nil +} + +// CloseWithMsg closes the melody instance with the given close payload and all connected sessions. +// Use the FormatCloseMessage function to format a proper close message payload. +func (m *Melody) CloseWithMsg(msg []byte) error { + if m.hub.closed() { + return errors.New("Melody instance is already closed.") + } + + m.hub.exit <- &envelope{t: websocket.CloseMessage, msg: msg} return nil } @@ -216,3 +248,8 @@ func (m *Melody) Close() error { func (m *Melody) Len() int { return m.hub.len() } + +// FormatCloseMessage formats closeCode and text as a WebSocket close message. +func FormatCloseMessage(closeCode int, text string) []byte { + return websocket.FormatCloseMessage(closeCode, text) +} diff --git a/session.go b/session.go index 00aab67..f728c22 100644 --- a/session.go +++ b/session.go @@ -2,10 +2,11 @@ package melody import ( "errors" - "github.com/gorilla/websocket" "net/http" "sync" "time" + + "github.com/gorilla/websocket" ) // Session wrapper around websocket connections. @@ -165,6 +166,18 @@ func (s *Session) Close() error { return nil } +// CloseWithMsg closes the session with the provided payload. +// Use the FormatCloseMessage function to format a proper close message payload. +func (s *Session) CloseWithMsg(msg []byte) error { + if s.closed() { + return errors.New("Session is already closed.") + } + + s.writeMessage(&envelope{t: websocket.CloseMessage, msg: msg}) + + return nil +} + // Set is used to store a new key/value pair exclusivelly for this session. // It also lazy initializes s.Keys if it was not used previously. func (s *Session) Set(key string, value interface{}) {