github.com/ezoic/ws@v1.0.4-0.20220713205711-5c1d69e074c5/wsutil/handler.go

package wsutil

import (
	"errors"
	"io"
	"io/ioutil"
	"strconv"

	"github.com/ezoic/pool/pbytes"
	"github.com/ezoic/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)
}