forked from mirror/websocket
Add Conn.[Set]CloseHandler
This commit is contained in:
parent
80a0029a65
commit
9fbf129ff2
31
conn.go
31
conn.go
|
@ -247,6 +247,7 @@ type Conn struct {
|
||||||
readMaskKey [4]byte
|
readMaskKey [4]byte
|
||||||
handlePong func(string) error
|
handlePong func(string) error
|
||||||
handlePing func(string) error
|
handlePing func(string) error
|
||||||
|
handleClose func(int, string) error
|
||||||
readErrCount int
|
readErrCount int
|
||||||
messageReader *messageReader // the current low-level reader
|
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),
|
writeBuf: make([]byte, writeBufferSize+maxFrameHeaderSize),
|
||||||
enableWriteCompression: true,
|
enableWriteCompression: true,
|
||||||
}
|
}
|
||||||
|
c.SetCloseHandler(nil)
|
||||||
c.SetPingHandler(nil)
|
c.SetPingHandler(nil)
|
||||||
c.SetPongHandler(nil)
|
c.SetPongHandler(nil)
|
||||||
return c
|
return c
|
||||||
|
@ -804,11 +806,9 @@ func (c *Conn) advanceFrame() (int, error) {
|
||||||
return noFrame, err
|
return noFrame, err
|
||||||
}
|
}
|
||||||
case CloseMessage:
|
case CloseMessage:
|
||||||
echoMessage := []byte{}
|
|
||||||
closeCode := CloseNoStatusReceived
|
closeCode := CloseNoStatusReceived
|
||||||
closeText := ""
|
closeText := ""
|
||||||
if len(payload) >= 2 {
|
if len(payload) >= 2 {
|
||||||
echoMessage = payload[:2]
|
|
||||||
closeCode = int(binary.BigEndian.Uint16(payload))
|
closeCode = int(binary.BigEndian.Uint16(payload))
|
||||||
if !isValidReceivedCloseCode(closeCode) {
|
if !isValidReceivedCloseCode(closeCode) {
|
||||||
return noFrame, c.handleProtocolError("invalid close code")
|
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")
|
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}
|
return noFrame, &CloseError{Code: closeCode, Text: closeText}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -946,6 +948,29 @@ func (c *Conn) SetReadLimit(limit int64) {
|
||||||
c.readLimit = limit
|
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
|
// PingHandler returns the current ping handler
|
||||||
func (c *Conn) PingHandler() func(appData string) error {
|
func (c *Conn) PingHandler() func(appData string) error {
|
||||||
return c.handlePing
|
return c.handlePing
|
||||||
|
|
Loading…
Reference in New Issue