From 9fbf129ff2a3cf2467f7d0021de2eb4d3aecc109 Mon Sep 17 00:00:00 2001
From: Gary Burd <gary@beagledreams.com>
Date: Wed, 2 Nov 2016 13:04:59 -0700
Subject: [PATCH] Add Conn.[Set]CloseHandler

---
 conn.go | 31 ++++++++++++++++++++++++++++---
 1 file changed, 28 insertions(+), 3 deletions(-)

diff --git a/conn.go b/conn.go
index cc04ed9..b7a97ba 100644
--- a/conn.go
+++ b/conn.go
@@ -247,6 +247,7 @@ type Conn struct {
 	readMaskKey   [4]byte
 	handlePong    func(string) error
 	handlePing    func(string) error
+	handleClose   func(int, string) error
 	readErrCount  int
 	messageReader *messageReader // the current low-level reader
 
@@ -277,6 +278,7 @@ func newConn(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int)
 		writeBuf:               make([]byte, writeBufferSize+maxFrameHeaderSize),
 		enableWriteCompression: true,
 	}
+	c.SetCloseHandler(nil)
 	c.SetPingHandler(nil)
 	c.SetPongHandler(nil)
 	return c
@@ -804,11 +806,9 @@ func (c *Conn) advanceFrame() (int, error) {
 			return noFrame, err
 		}
 	case CloseMessage:
-		echoMessage := []byte{}
 		closeCode := CloseNoStatusReceived
 		closeText := ""
 		if len(payload) >= 2 {
-			echoMessage = payload[:2]
 			closeCode = int(binary.BigEndian.Uint16(payload))
 			if !isValidReceivedCloseCode(closeCode) {
 				return noFrame, c.handleProtocolError("invalid close code")
@@ -818,7 +818,9 @@ func (c *Conn) advanceFrame() (int, error) {
 				return noFrame, c.handleProtocolError("invalid utf8 payload in close frame")
 			}
 		}
-		c.WriteControl(CloseMessage, echoMessage, time.Now().Add(writeWait))
+		if err := c.handleClose(closeCode, closeText); err != nil {
+			return noFrame, err
+		}
 		return noFrame, &CloseError{Code: closeCode, Text: closeText}
 	}
 
@@ -946,6 +948,29 @@ func (c *Conn) SetReadLimit(limit int64) {
 	c.readLimit = limit
 }
 
+// CloseHandler returns the current close handler
+func (c *Conn) CloseHandler() func(code int, text string) error {
+	return c.handleClose
+}
+
+// SetCloseHandler sets the handler for close messages received from the peer.
+// 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 peer.
+func (c *Conn) SetCloseHandler(h func(code int, text string) error) {
+	if h == nil {
+		h = func(code int, text string) error {
+			message := []byte{}
+			if code != CloseNoStatusReceived {
+				message = FormatCloseMessage(code, "")
+			}
+			c.WriteControl(CloseMessage, message, time.Now().Add(writeWait))
+			return nil
+		}
+	}
+	c.handleClose = h
+}
+
 // PingHandler returns the current ping handler
 func (c *Conn) PingHandler() func(appData string) error {
 	return c.handlePing