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) }