well-goknown/vendor/github.com/gobwas/ws/wsutil/handler.go

220 lines
6.1 KiB
Go
Raw Normal View History

package wsutil
import (
"errors"
"io"
"io/ioutil"
"strconv"
"github.com/gobwas/pool/pbytes"
"github.com/gobwas/ws"
)
// ClosedError returned when peer has closed the connection with appropriate
// code and a textual reason.
type ClosedError struct {
Code ws.StatusCode
Reason string
}
// Error implements error interface.
func (err ClosedError) Error() string {
return "ws closed: " + strconv.FormatUint(uint64(err.Code), 10) + " " + err.Reason
}
// ControlHandler contains logic of handling control frames.
//
// The intentional way to use it is to read the next frame header from the
// connection, optionally check its validity via ws.CheckHeader() and if it is
// not a ws.OpText of ws.OpBinary (or ws.OpContinuation) pass it to Handle()
// method.
//
// That is, passed header should be checked to get rid of unexpected errors.
//
// The Handle() method will read out all control frame payload (if any) and
// write necessary bytes as a rfc compatible response.
type ControlHandler struct {
Src io.Reader
Dst io.Writer
State ws.State
// DisableSrcCiphering disables unmasking payload data read from Src.
// It is useful when wsutil.Reader is used or when frame payload already
// pulled and ciphered out from the connection (and introduced by
// bytes.Reader, for example).
DisableSrcCiphering bool
}
// ErrNotControlFrame is returned by ControlHandler to indicate that given
// header could not be handled.
var ErrNotControlFrame = errors.New("not a control frame")
// Handle handles control frames regarding to the c.State and writes responses
// to the c.Dst when needed.
//
// It returns ErrNotControlFrame when given header is not of ws.OpClose,
// ws.OpPing or ws.OpPong operation code.
func (c ControlHandler) Handle(h ws.Header) error {
switch h.OpCode {
case ws.OpPing:
return c.HandlePing(h)
case ws.OpPong:
return c.HandlePong(h)
case ws.OpClose:
return c.HandleClose(h)
}
return ErrNotControlFrame
}
// HandlePing handles ping frame and writes specification compatible response
// to the c.Dst.
func (c ControlHandler) HandlePing(h ws.Header) error {
if h.Length == 0 {
// The most common case when ping is empty.
// Note that when sending masked frame the mask for empty payload is
// just four zero bytes.
return ws.WriteHeader(c.Dst, ws.Header{
Fin: true,
OpCode: ws.OpPong,
Masked: c.State.ClientSide(),
})
}
// In other way reply with Pong frame with copied payload.
p := pbytes.GetLen(int(h.Length) + ws.HeaderSize(ws.Header{
Length: h.Length,
Masked: c.State.ClientSide(),
}))
defer pbytes.Put(p)
// Deal with ciphering i/o:
// Masking key is used to mask the "Payload data" defined in the same
// section as frame-payload-data, which includes "Extension data" and
// "Application data".
//
// See https://tools.ietf.org/html/rfc6455#section-5.3
//
// NOTE: We prefer ControlWriter with preallocated buffer to
// ws.WriteHeader because it performs one syscall instead of two.
w := NewControlWriterBuffer(c.Dst, c.State, ws.OpPong, p)
r := c.Src
if c.State.ServerSide() && !c.DisableSrcCiphering {
r = NewCipherReader(r, h.Mask)
}
_, err := io.Copy(w, r)
if err == nil {
err = w.Flush()
}
return err
}
// HandlePong handles pong frame by discarding it.
func (c ControlHandler) HandlePong(h ws.Header) error {
if h.Length == 0 {
return nil
}
buf := pbytes.GetLen(int(h.Length))
defer pbytes.Put(buf)
// Discard pong message according to the RFC6455:
// A Pong frame MAY be sent unsolicited. This serves as a
// unidirectional heartbeat. A response to an unsolicited Pong frame
// is not expected.
_, err := io.CopyBuffer(ioutil.Discard, c.Src, buf)
return err
}
// HandleClose handles close frame, makes protocol validity checks and writes
// specification compatible response to the c.Dst.
func (c ControlHandler) HandleClose(h ws.Header) error {
if h.Length == 0 {
err := ws.WriteHeader(c.Dst, ws.Header{
Fin: true,
OpCode: ws.OpClose,
Masked: c.State.ClientSide(),
})
if err != nil {
return err
}
// Due to RFC, we should interpret the code as no status code
// received:
// If this Close control frame contains no status code, _The WebSocket
// Connection Close Code_ is considered to be 1005.
//
// See https://tools.ietf.org/html/rfc6455#section-7.1.5
return ClosedError{
Code: ws.StatusNoStatusRcvd,
}
}
// Prepare bytes both for reading reason and sending response.
p := pbytes.GetLen(int(h.Length) + ws.HeaderSize(ws.Header{
Length: h.Length,
Masked: c.State.ClientSide(),
}))
defer pbytes.Put(p)
// Get the subslice to read the frame payload out.
subp := p[:h.Length]
r := c.Src
if c.State.ServerSide() && !c.DisableSrcCiphering {
r = NewCipherReader(r, h.Mask)
}
if _, err := io.ReadFull(r, subp); err != nil {
return err
}
code, reason := ws.ParseCloseFrameData(subp)
if err := ws.CheckCloseFrameData(code, reason); err != nil {
// Here we could not use the prepared bytes because there is no
// guarantee that it may fit our protocol error closure code and a
// reason.
c.closeWithProtocolError(err)
return err
}
// Deal with ciphering i/o:
// Masking key is used to mask the "Payload data" defined in the same
// section as frame-payload-data, which includes "Extension data" and
// "Application data".
//
// See https://tools.ietf.org/html/rfc6455#section-5.3
//
// NOTE: We prefer ControlWriter with preallocated buffer to
// ws.WriteHeader because it performs one syscall instead of two.
w := NewControlWriterBuffer(c.Dst, c.State, ws.OpClose, p)
// RFC6455#5.5.1:
// If an endpoint receives a Close frame and did not previously
// send a Close frame, the endpoint MUST send a Close frame in
// response. (When sending a Close frame in response, the endpoint
// typically echoes the status code it received.)
_, err := w.Write(p[:2])
if err != nil {
return err
}
if err := w.Flush(); err != nil {
return err
}
return ClosedError{
Code: code,
Reason: reason,
}
}
func (c ControlHandler) closeWithProtocolError(reason error) error {
f := ws.NewCloseFrame(ws.NewCloseFrameBody(
ws.StatusProtocolError, reason.Error(),
))
if c.State.ClientSide() {
ws.MaskFrameInPlace(f)
}
return ws.WriteFrame(c.Dst, f)
}