From 9a70fe93acbc946c5a283adce611603c99d0f95b Mon Sep 17 00:00:00 2001
From: Robbie Trencheny <me@robbiet.us>
Date: Tue, 28 Mar 2017 14:00:51 -0700
Subject: [PATCH] Add closeHandler support

---
 melody.go  | 42 ++++++++++++++++++++++++++++++++++++++++++
 session.go |  5 +++++
 2 files changed, 47 insertions(+)

diff --git a/melody.go b/melody.go
index 3c16334..546b2b3 100644
--- a/melody.go
+++ b/melody.go
@@ -27,8 +27,29 @@ const (
 	CloseTLSHandshake            = 1015
 )
 
+// Duplicate of codes from gorilla/websocket for convenience.
+var validReceivedCloseCodes = map[int]bool{
+	// see http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number
+
+	CloseNormalClosure:           true,
+	CloseGoingAway:               true,
+	CloseProtocolError:           true,
+	CloseUnsupportedData:         true,
+	CloseNoStatusReceived:        false,
+	CloseAbnormalClosure:         false,
+	CloseInvalidFramePayloadData: true,
+	ClosePolicyViolation:         true,
+	CloseMessageTooBig:           true,
+	CloseMandatoryExtension:      true,
+	CloseInternalServerErr:       true,
+	CloseServiceRestart:          true,
+	CloseTryAgainLater:           true,
+	CloseTLSHandshake:            false,
+}
+
 type handleMessageFunc func(*Session, []byte)
 type handleErrorFunc func(*Session, error)
+type handleCloseFunc func(*Session, int, string) error
 type handleSessionFunc func(*Session)
 type filterFunc func(*Session) bool
 
@@ -41,6 +62,7 @@ type Melody struct {
 	messageSentHandler       handleMessageFunc
 	messageSentHandlerBinary handleMessageFunc
 	errorHandler             handleErrorFunc
+	closeHandler             handleCloseFunc
 	connectHandler           handleSessionFunc
 	disconnectHandler        handleSessionFunc
 	pongHandler              handleSessionFunc
@@ -66,6 +88,7 @@ func New() *Melody {
 		messageSentHandler:       func(*Session, []byte) {},
 		messageSentHandlerBinary: func(*Session, []byte) {},
 		errorHandler:             func(*Session, error) {},
+		closeHandler:             func(*Session, int, string) error { return nil },
 		connectHandler:           func(*Session) {},
 		disconnectHandler:        func(*Session) {},
 		pongHandler:              func(*Session) {},
@@ -113,6 +136,25 @@ func (m *Melody) HandleError(fn func(*Session, error)) {
 	m.errorHandler = fn
 }
 
+// HandleClose sets the handler for close messages received from the session.
+// The code argument to h is the received close code or CloseNoStatusReceived
+// if the close message is empty. The default close handler sends a close frame
+// back to the session.
+//
+// The application must read the connection to process close messages as
+// described in the section on Control Frames above.
+//
+// The connection read methods return a CloseError when a close frame is
+// received. Most applications should handle close messages as part of their
+// normal error handling. Applications should only set a close handler when the
+// application must perform some action before sending a close frame back to
+// the session.
+func (m *Melody) HandleClose(fn func(*Session, int, string) error) {
+	if fn != nil {
+		m.closeHandler = fn
+	}
+}
+
 // HandleRequest upgrades http requests to websocket connections and dispatches them to be handled by the melody instance.
 func (m *Melody) HandleRequest(w http.ResponseWriter, r *http.Request) error {
 	return m.HandleRequestWithKeys(w, r, nil)
diff --git a/session.go b/session.go
index f728c22..ed49f34 100644
--- a/session.go
+++ b/session.go
@@ -115,6 +115,11 @@ func (s *Session) readPump() {
 		return nil
 	})
 
+	s.conn.SetCloseHandler(func(code int, text string) error {
+		s.melody.closeHandler(s, code, text)
+		return nil
+	})
+
 	for {
 		t, message, err := s.conn.ReadMessage()