Allow sending custom close messages

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()
 			}

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)
+}

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{}) {