TUN-3403: Unit test for origin/proxy to test serving HTTP and Websocket

vendor/github.com/gobwas/ws/wsutil/cipher.go

@@ -0,0 +1,72 @@
+package wsutil
+
+import (
+	"io"
+
+	"github.com/gobwas/pool/pbytes"
+	"github.com/gobwas/ws"
+)
+
+// CipherReader implements io.Reader that applies xor-cipher to the bytes read
+// from source.
+// It could help to unmask WebSocket frame payload on the fly.
+type CipherReader struct {
+	r    io.Reader
+	mask [4]byte
+	pos  int
+}
+
+// NewCipherReader creates xor-cipher reader from r with given mask.
+func NewCipherReader(r io.Reader, mask [4]byte) *CipherReader {
+	return &CipherReader{r, mask, 0}
+}
+
+// Reset resets CipherReader to read from r with given mask.
+func (c *CipherReader) Reset(r io.Reader, mask [4]byte) {
+	c.r = r
+	c.mask = mask
+	c.pos = 0
+}
+
+// Read implements io.Reader interface. It applies mask given during
+// initialization to every read byte.
+func (c *CipherReader) Read(p []byte) (n int, err error) {
+	n, err = c.r.Read(p)
+	ws.Cipher(p[:n], c.mask, c.pos)
+	c.pos += n
+	return
+}
+
+// CipherWriter implements io.Writer that applies xor-cipher to the bytes
+// written to the destination writer. It does not modify the original bytes.
+type CipherWriter struct {
+	w    io.Writer
+	mask [4]byte
+	pos  int
+}
+
+// NewCipherWriter creates xor-cipher writer to w with given mask.
+func NewCipherWriter(w io.Writer, mask [4]byte) *CipherWriter {
+	return &CipherWriter{w, mask, 0}
+}
+
+// Reset reset CipherWriter to write to w with given mask.
+func (c *CipherWriter) Reset(w io.Writer, mask [4]byte) {
+	c.w = w
+	c.mask = mask
+	c.pos = 0
+}
+
+// Write implements io.Writer interface. It applies masking during
+// initialization to every sent byte. It does not modify original slice.
+func (c *CipherWriter) Write(p []byte) (n int, err error) {
+	cp := pbytes.GetLen(len(p))
+	defer pbytes.Put(cp)
+
+	copy(cp, p)
+	ws.Cipher(cp, c.mask, c.pos)
+	n, err = c.w.Write(cp)
+	c.pos += n
+
+	return
+}

vendor/github.com/gobwas/ws/wsutil/dialer.go

@@ -0,0 +1,146 @@
+package wsutil
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"io"
+	"io/ioutil"
+	"net"
+	"net/http"
+
+	"github.com/gobwas/ws"
+)
+
+// DebugDialer is a wrapper around ws.Dialer. It tracks i/o of WebSocket
+// handshake. That is, it gives ability to receive copied HTTP request and
+// response bytes that made inside Dialer.Dial().
+//
+// Note that it must not be used in production applications that requires
+// Dial() to be efficient.
+type DebugDialer struct {
+	// Dialer contains WebSocket connection establishment options.
+	Dialer ws.Dialer
+
+	// OnRequest and OnResponse are the callbacks that will be called with the
+	// HTTP request and response respectively.
+	OnRequest, OnResponse func([]byte)
+}
+
+// Dial connects to the url host and upgrades connection to WebSocket. It makes
+// it by calling d.Dialer.Dial().
+func (d *DebugDialer) Dial(ctx context.Context, urlstr string) (conn net.Conn, br *bufio.Reader, hs ws.Handshake, err error) {
+	// Need to copy Dialer to prevent original object mutation.
+	dialer := d.Dialer
+	var (
+		reqBuf bytes.Buffer
+		resBuf bytes.Buffer
+
+		resContentLength int64
+	)
+	userWrap := dialer.WrapConn
+	dialer.WrapConn = func(c net.Conn) net.Conn {
+		if userWrap != nil {
+			c = userWrap(c)
+		}
+
+		// Save the pointer to the raw connection.
+		conn = c
+
+		var (
+			r io.Reader = conn
+			w io.Writer = conn
+		)
+		if d.OnResponse != nil {
+			r = &prefetchResponseReader{
+				source:        conn,
+				buffer:        &resBuf,
+				contentLength: &resContentLength,
+			}
+		}
+		if d.OnRequest != nil {
+			w = io.MultiWriter(conn, &reqBuf)
+		}
+		return rwConn{conn, r, w}
+	}
+
+	_, br, hs, err = dialer.Dial(ctx, urlstr)
+
+	if onRequest := d.OnRequest; onRequest != nil {
+		onRequest(reqBuf.Bytes())
+	}
+	if onResponse := d.OnResponse; onResponse != nil {
+		// We must split response inside buffered bytes from other received
+		// bytes from server.
+		p := resBuf.Bytes()
+		n := bytes.Index(p, headEnd)
+		h := n + len(headEnd)         // Head end index.
+		n = h + int(resContentLength) // Body end index.
+
+		onResponse(p[:n])
+
+		if br != nil {
+			// If br is non-nil, then it mean two things. First is that
+			// handshake is OK and server has sent additional bytes – probably
+			// immediate sent frames (or weird but possible response body).
+			// Second, the bad one, is that br buffer's source is now rwConn
+			// instance from above WrapConn call. It is incorrect, so we must
+			// fix it.
+			var r io.Reader = conn
+			if len(p) > h {
+				// Buffer contains more than just HTTP headers bytes.
+				r = io.MultiReader(
+					bytes.NewReader(p[h:]),
+					conn,
+				)
+			}
+			br.Reset(r)
+			// Must make br.Buffered() to be non-zero.
+			br.Peek(len(p[h:]))
+		}
+	}
+
+	return conn, br, hs, err
+}
+
+type rwConn struct {
+	net.Conn
+
+	r io.Reader
+	w io.Writer
+}
+
+func (rwc rwConn) Read(p []byte) (int, error) {
+	return rwc.r.Read(p)
+}
+func (rwc rwConn) Write(p []byte) (int, error) {
+	return rwc.w.Write(p)
+}
+
+var headEnd = []byte("\r\n\r\n")
+
+type prefetchResponseReader struct {
+	source io.Reader // Original connection source.
+	reader io.Reader // Wrapped reader used to read from by clients.
+	buffer *bytes.Buffer
+
+	contentLength *int64
+}
+
+func (r *prefetchResponseReader) Read(p []byte) (int, error) {
+	if r.reader == nil {
+		resp, err := http.ReadResponse(bufio.NewReader(
+			io.TeeReader(r.source, r.buffer),
+		), nil)
+		if err == nil {
+			*r.contentLength, _ = io.Copy(ioutil.Discard, resp.Body)
+			resp.Body.Close()
+		}
+		bts := r.buffer.Bytes()
+		r.reader = io.MultiReader(
+			bytes.NewReader(bts),
+			r.source,
+		)
+	}
+	return r.reader.Read(p)
+}

vendor/github.com/gobwas/ws/wsutil/handler.go

@@ -0,0 +1,219 @@
+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)
+}

vendor/github.com/gobwas/ws/wsutil/helper.go

@@ -0,0 +1,279 @@
+package wsutil
+
+import (
+	"bytes"
+	"io"
+	"io/ioutil"
+
+	"github.com/gobwas/ws"
+)
+
+// Message represents a message from peer, that could be presented in one or
+// more frames. That is, it contains payload of all message fragments and
+// operation code of initial frame for this message.
+type Message struct {
+	OpCode  ws.OpCode
+	Payload []byte
+}
+
+// ReadMessage is a helper function that reads next message from r. It appends
+// received message(s) to the third argument and returns the result of it and
+// an error if some failure happened. That is, it probably could receive more
+// than one message when peer sending fragmented message in multiple frames and
+// want to send some control frame between fragments. Then returned slice will
+// contain those control frames at first, and then result of gluing fragments.
+//
+// TODO(gobwas): add DefaultReader with buffer size options.
+func ReadMessage(r io.Reader, s ws.State, m []Message) ([]Message, error) {
+	rd := Reader{
+		Source:    r,
+		State:     s,
+		CheckUTF8: true,
+		OnIntermediate: func(hdr ws.Header, src io.Reader) error {
+			bts, err := ioutil.ReadAll(src)
+			if err != nil {
+				return err
+			}
+			m = append(m, Message{hdr.OpCode, bts})
+			return nil
+		},
+	}
+	h, err := rd.NextFrame()
+	if err != nil {
+		return m, err
+	}
+	var p []byte
+	if h.Fin {
+		// No more frames will be read. Use fixed sized buffer to read payload.
+		p = make([]byte, h.Length)
+		// It is not possible to receive io.EOF here because Reader does not
+		// return EOF if frame payload was successfully fetched.
+		// Thus we consistent here with io.Reader behavior.
+		_, err = io.ReadFull(&rd, p)
+	} else {
+		// Frame is fragmented, thus use ioutil.ReadAll behavior.
+		var buf bytes.Buffer
+		_, err = buf.ReadFrom(&rd)
+		p = buf.Bytes()
+	}
+	if err != nil {
+		return m, err
+	}
+	return append(m, Message{h.OpCode, p}), nil
+}
+
+// ReadClientMessage reads next message from r, considering that caller
+// represents server side.
+// It is a shortcut for ReadMessage(r, ws.StateServerSide, m)
+func ReadClientMessage(r io.Reader, m []Message) ([]Message, error) {
+	return ReadMessage(r, ws.StateServerSide, m)
+}
+
+// ReadServerMessage reads next message from r, considering that caller
+// represents client side.
+// It is a shortcut for ReadMessage(r, ws.StateClientSide, m)
+func ReadServerMessage(r io.Reader, m []Message) ([]Message, error) {
+	return ReadMessage(r, ws.StateClientSide, m)
+}
+
+// ReadData is a helper function that reads next data (non-control) message
+// from rw.
+// It takes care on handling all control frames. It will write response on
+// control frames to the write part of rw. It blocks until some data frame
+// will be received.
+//
+// Note this may handle and write control frames into the writer part of a
+// given io.ReadWriter.
+func ReadData(rw io.ReadWriter, s ws.State) ([]byte, ws.OpCode, error) {
+	return readData(rw, s, ws.OpText|ws.OpBinary)
+}
+
+// ReadClientData reads next data message from rw, considering that caller
+// represents server side. It is a shortcut for ReadData(rw, ws.StateServerSide).
+//
+// Note this may handle and write control frames into the writer part of a
+// given io.ReadWriter.
+func ReadClientData(rw io.ReadWriter) ([]byte, ws.OpCode, error) {
+	return ReadData(rw, ws.StateServerSide)
+}
+
+// ReadClientText reads next text message from rw, considering that caller
+// represents server side. It is a shortcut for ReadData(rw, ws.StateServerSide).
+// It discards received binary messages.
+//
+// Note this may handle and write control frames into the writer part of a
+// given io.ReadWriter.
+func ReadClientText(rw io.ReadWriter) ([]byte, error) {
+	p, _, err := readData(rw, ws.StateServerSide, ws.OpText)
+	return p, err
+}
+
+// ReadClientBinary reads next binary message from rw, considering that caller
+// represents server side. It is a shortcut for ReadData(rw, ws.StateServerSide).
+// It discards received text messages.
+//
+// Note this may handle and write control frames into the writer part of a given
+//  io.ReadWriter.
+func ReadClientBinary(rw io.ReadWriter) ([]byte, error) {
+	p, _, err := readData(rw, ws.StateServerSide, ws.OpBinary)
+	return p, err
+}
+
+// ReadServerData reads next data message from rw, considering that caller
+// represents client side. It is a shortcut for ReadData(rw, ws.StateClientSide).
+//
+// Note this may handle and write control frames into the writer part of a
+// given io.ReadWriter.
+func ReadServerData(rw io.ReadWriter) ([]byte, ws.OpCode, error) {
+	return ReadData(rw, ws.StateClientSide)
+}
+
+// ReadServerText reads next text message from rw, considering that caller
+// represents client side. It is a shortcut for ReadData(rw, ws.StateClientSide).
+// It discards received binary messages.
+//
+// Note this may handle and write control frames into the writer part of a given
+//  io.ReadWriter.
+func ReadServerText(rw io.ReadWriter) ([]byte, error) {
+	p, _, err := readData(rw, ws.StateClientSide, ws.OpText)
+	return p, err
+}
+
+// ReadServerBinary reads next binary message from rw, considering that caller
+// represents client side. It is a shortcut for ReadData(rw, ws.StateClientSide).
+// It discards received text messages.
+//
+// Note this may handle and write control frames into the writer part of a
+// given io.ReadWriter.
+func ReadServerBinary(rw io.ReadWriter) ([]byte, error) {
+	p, _, err := readData(rw, ws.StateClientSide, ws.OpBinary)
+	return p, err
+}
+
+// WriteMessage is a helper function that writes message to the w. It
+// constructs single frame with given operation code and payload.
+// It uses given state to prepare side-dependent things, like cipher
+// payload bytes from client to server. It will not mutate p bytes if
+// cipher must be made.
+//
+// If you want to write message in fragmented frames, use Writer instead.
+func WriteMessage(w io.Writer, s ws.State, op ws.OpCode, p []byte) error {
+	return writeFrame(w, s, op, true, p)
+}
+
+// WriteServerMessage writes message to w, considering that caller
+// represents server side.
+func WriteServerMessage(w io.Writer, op ws.OpCode, p []byte) error {
+	return WriteMessage(w, ws.StateServerSide, op, p)
+}
+
+// WriteServerText is the same as WriteServerMessage with
+// ws.OpText.
+func WriteServerText(w io.Writer, p []byte) error {
+	return WriteServerMessage(w, ws.OpText, p)
+}
+
+// WriteServerBinary is the same as WriteServerMessage with
+// ws.OpBinary.
+func WriteServerBinary(w io.Writer, p []byte) error {
+	return WriteServerMessage(w, ws.OpBinary, p)
+}
+
+// WriteClientMessage writes message to w, considering that caller
+// represents client side.
+func WriteClientMessage(w io.Writer, op ws.OpCode, p []byte) error {
+	return WriteMessage(w, ws.StateClientSide, op, p)
+}
+
+// WriteClientText is the same as WriteClientMessage with
+// ws.OpText.
+func WriteClientText(w io.Writer, p []byte) error {
+	return WriteClientMessage(w, ws.OpText, p)
+}
+
+// WriteClientBinary is the same as WriteClientMessage with
+// ws.OpBinary.
+func WriteClientBinary(w io.Writer, p []byte) error {
+	return WriteClientMessage(w, ws.OpBinary, p)
+}
+
+// HandleClientControlMessage handles control frame from conn and writes
+// response when needed.
+//
+// It considers that caller represents server side.
+func HandleClientControlMessage(conn io.Writer, msg Message) error {
+	return HandleControlMessage(conn, ws.StateServerSide, msg)
+}
+
+// HandleServerControlMessage handles control frame from conn and writes
+// response when needed.
+//
+// It considers that caller represents client side.
+func HandleServerControlMessage(conn io.Writer, msg Message) error {
+	return HandleControlMessage(conn, ws.StateClientSide, msg)
+}
+
+// HandleControlMessage handles message which was read by ReadMessage()
+// functions.
+//
+// That is, it is expected, that payload is already unmasked and frame header
+// were checked by ws.CheckHeader() call.
+func HandleControlMessage(conn io.Writer, state ws.State, msg Message) error {
+	return (ControlHandler{
+		DisableSrcCiphering: true,
+		Src:                 bytes.NewReader(msg.Payload),
+		Dst:                 conn,
+		State:               state,
+	}).Handle(ws.Header{
+		Length: int64(len(msg.Payload)),
+		OpCode: msg.OpCode,
+		Fin:    true,
+		Masked: state.ServerSide(),
+	})
+}
+
+// ControlFrameHandler returns FrameHandlerFunc for handling control frames.
+// For more info see ControlHandler docs.
+func ControlFrameHandler(w io.Writer, state ws.State) FrameHandlerFunc {
+	return func(h ws.Header, r io.Reader) error {
+		return (ControlHandler{
+			DisableSrcCiphering: true,
+			Src:                 r,
+			Dst:                 w,
+			State:               state,
+		}).Handle(h)
+	}
+}
+
+func readData(rw io.ReadWriter, s ws.State, want ws.OpCode) ([]byte, ws.OpCode, error) {
+	controlHandler := ControlFrameHandler(rw, s)
+	rd := Reader{
+		Source:          rw,
+		State:           s,
+		CheckUTF8:       true,
+		SkipHeaderCheck: false,
+		OnIntermediate:  controlHandler,
+	}
+	for {
+		hdr, err := rd.NextFrame()
+		if err != nil {
+			return nil, 0, err
+		}
+		if hdr.OpCode.IsControl() {
+			if err := controlHandler(hdr, &rd); err != nil {
+				return nil, 0, err
+			}
+			continue
+		}
+		if hdr.OpCode&want == 0 {
+			if err := rd.Discard(); err != nil {
+				return nil, 0, err
+			}
+			continue
+		}
+
+		bts, err := ioutil.ReadAll(&rd)
+
+		return bts, hdr.OpCode, err
+	}
+}

vendor/github.com/gobwas/ws/wsutil/reader.go

@@ -0,0 +1,257 @@
+package wsutil
+
+import (
+	"errors"
+	"io"
+	"io/ioutil"
+
+	"github.com/gobwas/ws"
+)
+
+// ErrNoFrameAdvance means that Reader's Read() method was called without
+// preceding NextFrame() call.
+var ErrNoFrameAdvance = errors.New("no frame advance")
+
+// FrameHandlerFunc handles parsed frame header and its body represented by
+// io.Reader.
+//
+// Note that reader represents already unmasked body.
+type FrameHandlerFunc func(ws.Header, io.Reader) error
+
+// Reader is a wrapper around source io.Reader which represents WebSocket
+// connection. It contains options for reading messages from source.
+//
+// Reader implements io.Reader, which Read() method reads payload of incoming
+// WebSocket frames. It also takes care on fragmented frames and possibly
+// intermediate control frames between them.
+//
+// Note that Reader's methods are not goroutine safe.
+type Reader struct {
+	Source io.Reader
+	State  ws.State
+
+	// SkipHeaderCheck disables checking header bits to be RFC6455 compliant.
+	SkipHeaderCheck bool
+
+	// CheckUTF8 enables UTF-8 checks for text frames payload. If incoming
+	// bytes are not valid UTF-8 sequence, ErrInvalidUTF8 returned.
+	CheckUTF8 bool
+
+	// TODO(gobwas): add max frame size limit here.
+
+	OnContinuation FrameHandlerFunc
+	OnIntermediate FrameHandlerFunc
+
+	opCode ws.OpCode        // Used to store message op code on fragmentation.
+	frame  io.Reader        // Used to as frame reader.
+	raw    io.LimitedReader // Used to discard frames without cipher.
+	utf8   UTF8Reader       // Used to check UTF8 sequences if CheckUTF8 is true.
+}
+
+// NewReader creates new frame reader that reads from r keeping given state to
+// make some protocol validity checks when it needed.
+func NewReader(r io.Reader, s ws.State) *Reader {
+	return &Reader{
+		Source: r,
+		State:  s,
+	}
+}
+
+// NewClientSideReader is a helper function that calls NewReader with r and
+// ws.StateClientSide.
+func NewClientSideReader(r io.Reader) *Reader {
+	return NewReader(r, ws.StateClientSide)
+}
+
+// NewServerSideReader is a helper function that calls NewReader with r and
+// ws.StateServerSide.
+func NewServerSideReader(r io.Reader) *Reader {
+	return NewReader(r, ws.StateServerSide)
+}
+
+// Read implements io.Reader. It reads the next message payload into p.
+// It takes care on fragmented messages.
+//
+// The error is io.EOF only if all of message bytes were read.
+// If an io.EOF happens during reading some but not all the message bytes
+// Read() returns io.ErrUnexpectedEOF.
+//
+// The error is ErrNoFrameAdvance if no NextFrame() call was made before
+// reading next message bytes.
+func (r *Reader) Read(p []byte) (n int, err error) {
+	if r.frame == nil {
+		if !r.fragmented() {
+			// Every new Read() must be preceded by NextFrame() call.
+			return 0, ErrNoFrameAdvance
+		}
+		// Read next continuation or intermediate control frame.
+		_, err := r.NextFrame()
+		if err != nil {
+			return 0, err
+		}
+		if r.frame == nil {
+			// We handled intermediate control and now got nothing to read.
+			return 0, nil
+		}
+	}
+
+	n, err = r.frame.Read(p)
+	if err != nil && err != io.EOF {
+		return
+	}
+	if err == nil && r.raw.N != 0 {
+		return
+	}
+
+	switch {
+	case r.raw.N != 0:
+		err = io.ErrUnexpectedEOF
+
+	case r.fragmented():
+		err = nil
+		r.resetFragment()
+
+	case r.CheckUTF8 && !r.utf8.Valid():
+		n = r.utf8.Accepted()
+		err = ErrInvalidUTF8
+
+	default:
+		r.reset()
+		err = io.EOF
+	}
+
+	return
+}
+
+// Discard discards current message unread bytes.
+// It discards all frames of fragmented message.
+func (r *Reader) Discard() (err error) {
+	for {
+		_, err = io.Copy(ioutil.Discard, &r.raw)
+		if err != nil {
+			break
+		}
+		if !r.fragmented() {
+			break
+		}
+		if _, err = r.NextFrame(); err != nil {
+			break
+		}
+	}
+	r.reset()
+	return err
+}
+
+// NextFrame prepares r to read next message. It returns received frame header
+// and non-nil error on failure.
+//
+// Note that next NextFrame() call must be done after receiving or discarding
+// all current message bytes.
+func (r *Reader) NextFrame() (hdr ws.Header, err error) {
+	hdr, err = ws.ReadHeader(r.Source)
+	if err == io.EOF && r.fragmented() {
+		// If we are in fragmented state EOF means that is was totally
+		// unexpected.
+		//
+		// NOTE: This is necessary to prevent callers such that
+		// ioutil.ReadAll to receive some amount of bytes without an error.
+		// ReadAll() ignores an io.EOF error, thus caller may think that
+		// whole message fetched, but actually only part of it.
+		err = io.ErrUnexpectedEOF
+	}
+	if err == nil && !r.SkipHeaderCheck {
+		err = ws.CheckHeader(hdr, r.State)
+	}
+	if err != nil {
+		return hdr, err
+	}
+
+	// Save raw reader to use it on discarding frame without ciphering and
+	// other streaming checks.
+	r.raw = io.LimitedReader{r.Source, hdr.Length}
+
+	frame := io.Reader(&r.raw)
+	if hdr.Masked {
+		frame = NewCipherReader(frame, hdr.Mask)
+	}
+	if r.fragmented() {
+		if hdr.OpCode.IsControl() {
+			if cb := r.OnIntermediate; cb != nil {
+				err = cb(hdr, frame)
+			}
+			if err == nil {
+				// Ensure that src is empty.
+				_, err = io.Copy(ioutil.Discard, &r.raw)
+			}
+			return
+		}
+	} else {
+		r.opCode = hdr.OpCode
+	}
+	if r.CheckUTF8 && (hdr.OpCode == ws.OpText || (r.fragmented() && r.opCode == ws.OpText)) {
+		r.utf8.Source = frame
+		frame = &r.utf8
+	}
+
+	// Save reader with ciphering and other streaming checks.
+	r.frame = frame
+
+	if hdr.OpCode == ws.OpContinuation {
+		if cb := r.OnContinuation; cb != nil {
+			err = cb(hdr, frame)
+		}
+	}
+
+	if hdr.Fin {
+		r.State = r.State.Clear(ws.StateFragmented)
+	} else {
+		r.State = r.State.Set(ws.StateFragmented)
+	}
+
+	return
+}
+
+func (r *Reader) fragmented() bool {
+	return r.State.Fragmented()
+}
+
+func (r *Reader) resetFragment() {
+	r.raw = io.LimitedReader{}
+	r.frame = nil
+	// Reset source of the UTF8Reader, but not the state.
+	r.utf8.Source = nil
+}
+
+func (r *Reader) reset() {
+	r.raw = io.LimitedReader{}
+	r.frame = nil
+	r.utf8 = UTF8Reader{}
+	r.opCode = 0
+}
+
+// NextReader prepares next message read from r. It returns header that
+// describes the message and io.Reader to read message's payload. It returns
+// non-nil error when it is not possible to read message's initial frame.
+//
+// Note that next NextReader() on the same r should be done after reading all
+// bytes from previously returned io.Reader. For more performant way to discard
+// message use Reader and its Discard() method.
+//
+// Note that it will not handle any "intermediate" frames, that possibly could
+// be received between text/binary continuation frames. That is, if peer sent
+// text/binary frame with fin flag "false", then it could send ping frame, and
+// eventually remaining part of text/binary frame with fin "true" – with
+// NextReader() the ping frame will be dropped without any notice. To handle
+// this rare, but possible situation (and if you do not know exactly which
+// frames peer could send), you could use Reader with OnIntermediate field set.
+func NextReader(r io.Reader, s ws.State) (ws.Header, io.Reader, error) {
+	rd := &Reader{
+		Source: r,
+		State:  s,
+	}
+	header, err := rd.NextFrame()
+	if err != nil {
+		return header, nil, err
+	}
+	return header, rd, nil
+}

vendor/github.com/gobwas/ws/wsutil/upgrader.go

@@ -0,0 +1,68 @@
+package wsutil
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+	"io/ioutil"
+	"net/http"
+
+	"github.com/gobwas/ws"
+)
+
+// DebugUpgrader is a wrapper around ws.Upgrader. It tracks I/O of a
+// WebSocket handshake.
+//
+// Note that it must not be used in production applications that requires
+// Upgrade() to be efficient.
+type DebugUpgrader struct {
+	// Upgrader contains upgrade to WebSocket options.
+	Upgrader ws.Upgrader
+
+	// OnRequest and OnResponse are the callbacks that will be called with the
+	// HTTP request and response respectively.
+	OnRequest, OnResponse func([]byte)
+}
+
+// Upgrade calls Upgrade() on underlying ws.Upgrader and tracks I/O on conn.
+func (d *DebugUpgrader) Upgrade(conn io.ReadWriter) (hs ws.Handshake, err error) {
+	var (
+		// Take the Reader and Writer parts from conn to be probably replaced
+		// below.
+		r io.Reader = conn
+		w io.Writer = conn
+	)
+	if onRequest := d.OnRequest; onRequest != nil {
+		var buf bytes.Buffer
+		// First, we must read the entire request.
+		req, err := http.ReadRequest(bufio.NewReader(
+			io.TeeReader(conn, &buf),
+		))
+		if err == nil {
+			// Fulfill the buffer with the response body.
+			io.Copy(ioutil.Discard, req.Body)
+			req.Body.Close()
+		}
+		onRequest(buf.Bytes())
+
+		r = io.MultiReader(
+			&buf, conn,
+		)
+	}
+
+	if onResponse := d.OnResponse; onResponse != nil {
+		var buf bytes.Buffer
+		// Intercept the response stream written by the Upgrade().
+		w = io.MultiWriter(
+			conn, &buf,
+		)
+		defer func() {
+			onResponse(buf.Bytes())
+		}()
+	}
+
+	return d.Upgrader.Upgrade(struct {
+		io.Reader
+		io.Writer
+	}{r, w})
+}

vendor/github.com/gobwas/ws/wsutil/utf8.go

@@ -0,0 +1,140 @@
+package wsutil
+
+import (
+	"fmt"
+	"io"
+)
+
+// ErrInvalidUTF8 is returned by UTF8 reader on invalid utf8 sequence.
+var ErrInvalidUTF8 = fmt.Errorf("invalid utf8")
+
+// UTF8Reader implements io.Reader that calculates utf8 validity state after
+// every read byte from Source.
+//
+// Note that in some cases client must call r.Valid() after all bytes are read
+// to ensure that all of them are valid utf8 sequences. That is, some io helper
+// functions such io.ReadAtLeast or io.ReadFull could discard the error
+// information returned by the reader when they receive all of requested bytes.
+// For example, the last read sequence is invalid and UTF8Reader returns number
+// of bytes read and an error. But helper function decides to discard received
+// error due to all requested bytes are completely read from the source.
+//
+// Another possible case is when some valid sequence become split by the read
+// bound. Then UTF8Reader can not make decision about validity of the last
+// sequence cause it is not fully read yet. And if the read stops, Valid() will
+// return false, even if Read() by itself dit not.
+type UTF8Reader struct {
+	Source io.Reader
+
+	accepted int
+
+	state uint32
+	codep uint32
+}
+
+// NewUTF8Reader creates utf8 reader that reads from r.
+func NewUTF8Reader(r io.Reader) *UTF8Reader {
+	return &UTF8Reader{
+		Source: r,
+	}
+}
+
+// Reset resets utf8 reader to read from r.
+func (u *UTF8Reader) Reset(r io.Reader) {
+	u.Source = r
+	u.state = 0
+	u.codep = 0
+}
+
+// Read implements io.Reader.
+func (u *UTF8Reader) Read(p []byte) (n int, err error) {
+	n, err = u.Source.Read(p)
+
+	accepted := 0
+	s, c := u.state, u.codep
+	for i := 0; i < n; i++ {
+		c, s = decode(s, c, p[i])
+		if s == utf8Reject {
+			u.state = s
+			return accepted, ErrInvalidUTF8
+		}
+		if s == utf8Accept {
+			accepted = i + 1
+		}
+	}
+	u.state, u.codep = s, c
+	u.accepted = accepted
+
+	return
+}
+
+// Valid checks current reader state. It returns true if all read bytes are
+// valid UTF-8 sequences, and false if not.
+func (u *UTF8Reader) Valid() bool {
+	return u.state == utf8Accept
+}
+
+// Accepted returns number of valid bytes in last Read().
+func (u *UTF8Reader) Accepted() int {
+	return u.accepted
+}
+
+// Below is port of UTF-8 decoder from http://bjoern.hoehrmann.de/utf-8/decoder/dfa/
+//
+// Copyright (c) 2008-2009 Bjoern Hoehrmann <bjoern@hoehrmann.de>
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to
+// deal in the Software without restriction, including without limitation the
+// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+// sell copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+// IN THE SOFTWARE.
+
+const (
+	utf8Accept = 0
+	utf8Reject = 12
+)
+
+var utf8d = [...]byte{
+	// The first part of the table maps bytes to character classes that
+	// to reduce the size of the transition table and create bitmasks.
+	0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+	0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+	0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+	0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+	1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9,
+	7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7,
+	8, 8, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
+	10, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 3, 3, 11, 6, 6, 6, 5, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
+
+	// The second part is a transition table that maps a combination
+	// of a state of the automaton and a character class to a state.
+	0, 12, 24, 36, 60, 96, 84, 12, 12, 12, 48, 72, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12,
+	12, 0, 12, 12, 12, 12, 12, 0, 12, 0, 12, 12, 12, 24, 12, 12, 12, 12, 12, 24, 12, 24, 12, 12,
+	12, 12, 12, 12, 12, 12, 12, 24, 12, 12, 12, 12, 12, 24, 12, 12, 12, 12, 12, 12, 12, 24, 12, 12,
+	12, 12, 12, 12, 12, 12, 12, 36, 12, 36, 12, 12, 12, 36, 12, 12, 12, 12, 12, 36, 12, 36, 12, 12,
+	12, 36, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12,
+}
+
+func decode(state, codep uint32, b byte) (uint32, uint32) {
+	t := uint32(utf8d[b])
+
+	if state != utf8Accept {
+		codep = (uint32(b) & 0x3f) | (codep << 6)
+	} else {
+		codep = (0xff >> t) & uint32(b)
+	}
+
+	return codep, uint32(utf8d[256+state+t])
+}

vendor/github.com/gobwas/ws/wsutil/writer.go

@@ -0,0 +1,450 @@
+package wsutil
+
+import (
+	"fmt"
+	"io"
+
+	"github.com/gobwas/pool"
+	"github.com/gobwas/pool/pbytes"
+	"github.com/gobwas/ws"
+)
+
+// DefaultWriteBuffer contains size of Writer's default buffer. It used by
+// Writer constructor functions.
+var DefaultWriteBuffer = 4096
+
+var (
+	// ErrNotEmpty is returned by Writer.WriteThrough() to indicate that buffer is
+	// not empty and write through could not be done. That is, caller should call
+	// Writer.FlushFragment() to make buffer empty.
+	ErrNotEmpty = fmt.Errorf("writer not empty")
+
+	// ErrControlOverflow is returned by ControlWriter.Write() to indicate that
+	// no more data could be written to the underlying io.Writer because
+	// MaxControlFramePayloadSize limit is reached.
+	ErrControlOverflow = fmt.Errorf("control frame payload overflow")
+)
+
+// Constants which are represent frame length ranges.
+const (
+	len7  = int64(125) // 126 and 127 are reserved values
+	len16 = int64(^uint16(0))
+	len64 = int64((^uint64(0)) >> 1)
+)
+
+// ControlWriter is a wrapper around Writer that contains some guards for
+// buffered writes of control frames.
+type ControlWriter struct {
+	w     *Writer
+	limit int
+	n     int
+}
+
+// NewControlWriter contains ControlWriter with Writer inside whose buffer size
+// is at most ws.MaxControlFramePayloadSize + ws.MaxHeaderSize.
+func NewControlWriter(dest io.Writer, state ws.State, op ws.OpCode) *ControlWriter {
+	return &ControlWriter{
+		w:     NewWriterSize(dest, state, op, ws.MaxControlFramePayloadSize),
+		limit: ws.MaxControlFramePayloadSize,
+	}
+}
+
+// NewControlWriterBuffer returns a new ControlWriter with buf as a buffer.
+//
+// Note that it reserves x bytes of buf for header data, where x could be
+// ws.MinHeaderSize or ws.MinHeaderSize+4 (depending on state). At most
+// (ws.MaxControlFramePayloadSize + x) bytes of buf will be used.
+//
+// It panics if len(buf) <= ws.MinHeaderSize + x.
+func NewControlWriterBuffer(dest io.Writer, state ws.State, op ws.OpCode, buf []byte) *ControlWriter {
+	max := ws.MaxControlFramePayloadSize + headerSize(state, ws.MaxControlFramePayloadSize)
+	if len(buf) > max {
+		buf = buf[:max]
+	}
+
+	w := NewWriterBuffer(dest, state, op, buf)
+
+	return &ControlWriter{
+		w:     w,
+		limit: len(w.buf),
+	}
+}
+
+// Write implements io.Writer. It writes to the underlying Writer until it
+// returns error or until ControlWriter write limit will be exceeded.
+func (c *ControlWriter) Write(p []byte) (n int, err error) {
+	if c.n+len(p) > c.limit {
+		return 0, ErrControlOverflow
+	}
+	return c.w.Write(p)
+}
+
+// Flush flushes all buffered data to the underlying io.Writer.
+func (c *ControlWriter) Flush() error {
+	return c.w.Flush()
+}
+
+// Writer contains logic of buffering output data into a WebSocket fragments.
+// It is much the same as bufio.Writer, except the thing that it works with
+// WebSocket frames, not the raw data.
+//
+// Writer writes frames with specified OpCode.
+// It uses ws.State to decide whether the output frames must be masked.
+//
+// Note that it does not check control frame size or other RFC rules.
+// That is, it must be used with special care to write control frames without
+// violation of RFC. You could use ControlWriter that wraps Writer and contains
+// some guards for writing control frames.
+//
+// If an error occurs writing to a Writer, no more data will be accepted and
+// all subsequent writes will return the error.
+// After all data has been written, the client should call the Flush() method
+// to guarantee all data has been forwarded to the underlying io.Writer.
+type Writer struct {
+	dest io.Writer
+
+	n   int    // Buffered bytes counter.
+	raw []byte // Raw representation of buffer, including reserved header bytes.
+	buf []byte // Writeable part of buffer, without reserved header bytes.
+
+	op    ws.OpCode
+	state ws.State
+
+	dirty      bool
+	fragmented bool
+
+	err error
+}
+
+var writers = pool.New(128, 65536)
+
+// GetWriter tries to reuse Writer getting it from the pool.
+//
+// This function is intended for memory consumption optimizations, because
+// NewWriter*() functions make allocations for inner buffer.
+//
+// Note the it ceils n to the power of two.
+//
+// If you have your own bytes buffer pool you could use NewWriterBuffer to use
+// pooled bytes in writer.
+func GetWriter(dest io.Writer, state ws.State, op ws.OpCode, n int) *Writer {
+	x, m := writers.Get(n)
+	if x != nil {
+		w := x.(*Writer)
+		w.Reset(dest, state, op)
+		return w
+	}
+	// NOTE: we use m instead of n, because m is an attempt to reuse w of such
+	// size in the future.
+	return NewWriterBufferSize(dest, state, op, m)
+}
+
+// PutWriter puts w for future reuse by GetWriter().
+func PutWriter(w *Writer) {
+	w.Reset(nil, 0, 0)
+	writers.Put(w, w.Size())
+}
+
+// NewWriter returns a new Writer whose buffer has the DefaultWriteBuffer size.
+func NewWriter(dest io.Writer, state ws.State, op ws.OpCode) *Writer {
+	return NewWriterBufferSize(dest, state, op, 0)
+}
+
+// NewWriterSize returns a new Writer whose buffer size is at most n + ws.MaxHeaderSize.
+// That is, output frames payload length could be up to n, except the case when
+// Write() is called on empty Writer with len(p) > n.
+//
+// If n <= 0 then the default buffer size is used as Writer's buffer size.
+func NewWriterSize(dest io.Writer, state ws.State, op ws.OpCode, n int) *Writer {
+	if n > 0 {
+		n += headerSize(state, n)
+	}
+	return NewWriterBufferSize(dest, state, op, n)
+}
+
+// NewWriterBufferSize returns a new Writer whose buffer size is equal to n.
+// If n <= ws.MinHeaderSize then the default buffer size is used.
+//
+// Note that Writer will reserve x bytes for header data, where x is in range
+// [ws.MinHeaderSize,ws.MaxHeaderSize]. That is, frames flushed by Writer
+// will not have payload length equal to n, except the case when Write() is
+// called on empty Writer with len(p) > n.
+func NewWriterBufferSize(dest io.Writer, state ws.State, op ws.OpCode, n int) *Writer {
+	if n <= ws.MinHeaderSize {
+		n = DefaultWriteBuffer
+	}
+	return NewWriterBuffer(dest, state, op, make([]byte, n))
+}
+
+// NewWriterBuffer returns a new Writer with buf as a buffer.
+//
+// Note that it reserves x bytes of buf for header data, where x is in range
+// [ws.MinHeaderSize,ws.MaxHeaderSize] (depending on state and buf size).
+//
+// You could use ws.HeaderSize() to calculate number of bytes needed to store
+// header data.
+//
+// It panics if len(buf) is too small to fit header and payload data.
+func NewWriterBuffer(dest io.Writer, state ws.State, op ws.OpCode, buf []byte) *Writer {
+	offset := reserve(state, len(buf))
+	if len(buf) <= offset {
+		panic("buffer too small")
+	}
+
+	return &Writer{
+		dest:  dest,
+		raw:   buf,
+		buf:   buf[offset:],
+		state: state,
+		op:    op,
+	}
+}
+
+func reserve(state ws.State, n int) (offset int) {
+	var mask int
+	if state.ClientSide() {
+		mask = 4
+	}
+
+	switch {
+	case n <= int(len7)+mask+2:
+		return mask + 2
+	case n <= int(len16)+mask+4:
+		return mask + 4
+	default:
+		return mask + 10
+	}
+}
+
+// headerSize returns number of bytes needed to encode header of a frame with
+// given state and length.
+func headerSize(s ws.State, n int) int {
+	return ws.HeaderSize(ws.Header{
+		Length: int64(n),
+		Masked: s.ClientSide(),
+	})
+}
+
+// Reset discards any buffered data, clears error, and resets w to have given
+// state and write frames with given OpCode to dest.
+func (w *Writer) Reset(dest io.Writer, state ws.State, op ws.OpCode) {
+	w.n = 0
+	w.dirty = false
+	w.fragmented = false
+	w.dest = dest
+	w.state = state
+	w.op = op
+}
+
+// Size returns the size of the underlying buffer in bytes.
+func (w *Writer) Size() int {
+	return len(w.buf)
+}
+
+// Available returns how many bytes are unused in the buffer.
+func (w *Writer) Available() int {
+	return len(w.buf) - w.n
+}
+
+// Buffered returns the number of bytes that have been written into the current
+// buffer.
+func (w *Writer) Buffered() int {
+	return w.n
+}
+
+// Write implements io.Writer.
+//
+// Note that even if the Writer was created to have N-sized buffer, Write()
+// with payload of N bytes will not fit into that buffer. Writer reserves some
+// space to fit WebSocket header data.
+func (w *Writer) Write(p []byte) (n int, err error) {
+	// Even empty p may make a sense.
+	w.dirty = true
+
+	var nn int
+	for len(p) > w.Available() && w.err == nil {
+		if w.Buffered() == 0 {
+			// Large write, empty buffer. Write directly from p to avoid copy.
+			// Trade off here is that we make additional Write() to underlying
+			// io.Writer when writing frame header.
+			//
+			// On large buffers additional write is better than copying.
+			nn, _ = w.WriteThrough(p)
+		} else {
+			nn = copy(w.buf[w.n:], p)
+			w.n += nn
+			w.FlushFragment()
+		}
+		n += nn
+		p = p[nn:]
+	}
+	if w.err != nil {
+		return n, w.err
+	}
+	nn = copy(w.buf[w.n:], p)
+	w.n += nn
+	n += nn
+
+	// Even if w.Available() == 0 we will not flush buffer preventively because
+	// this could bring unwanted fragmentation. That is, user could create
+	// buffer with size that fits exactly all further Write() call, and then
+	// call Flush(), excepting that single and not fragmented frame will be
+	// sent. With preemptive flush this case will produce two frames – last one
+	// will be empty and just to set fin = true.
+
+	return n, w.err
+}
+
+// WriteThrough writes data bypassing the buffer.
+// Note that Writer's buffer must be empty before calling WriteThrough().
+func (w *Writer) WriteThrough(p []byte) (n int, err error) {
+	if w.err != nil {
+		return 0, w.err
+	}
+	if w.Buffered() != 0 {
+		return 0, ErrNotEmpty
+	}
+
+	w.err = writeFrame(w.dest, w.state, w.opCode(), false, p)
+	if w.err == nil {
+		n = len(p)
+	}
+
+	w.dirty = true
+	w.fragmented = true
+
+	return n, w.err
+}
+
+// ReadFrom implements io.ReaderFrom.
+func (w *Writer) ReadFrom(src io.Reader) (n int64, err error) {
+	var nn int
+	for err == nil {
+		if w.Available() == 0 {
+			err = w.FlushFragment()
+			continue
+		}
+
+		// We copy the behavior of bufio.Writer here.
+		// Also, from the docs on io.ReaderFrom:
+		//   ReadFrom reads data from r until EOF or error.
+		//
+		// See https://codereview.appspot.com/76400048/#ps1
+		const maxEmptyReads = 100
+		var nr int
+		for nr < maxEmptyReads {
+			nn, err = src.Read(w.buf[w.n:])
+			if nn != 0 || err != nil {
+				break
+			}
+			nr++
+		}
+		if nr == maxEmptyReads {
+			return n, io.ErrNoProgress
+		}
+
+		w.n += nn
+		n += int64(nn)
+	}
+	if err == io.EOF {
+		// NOTE: Do not flush preemptively.
+		// See the Write() sources for more info.
+		err = nil
+		w.dirty = true
+	}
+	return n, err
+}
+
+// Flush writes any buffered data to the underlying io.Writer.
+// It sends the frame with "fin" flag set to true.
+//
+// If no Write() or ReadFrom() was made, then Flush() does nothing.
+func (w *Writer) Flush() error {
+	if (!w.dirty && w.Buffered() == 0) || w.err != nil {
+		return w.err
+	}
+
+	w.err = w.flushFragment(true)
+	w.n = 0
+	w.dirty = false
+	w.fragmented = false
+
+	return w.err
+}
+
+// FlushFragment writes any buffered data to the underlying io.Writer.
+// It sends the frame with "fin" flag set to false.
+func (w *Writer) FlushFragment() error {
+	if w.Buffered() == 0 || w.err != nil {
+		return w.err
+	}
+
+	w.err = w.flushFragment(false)
+	w.n = 0
+	w.fragmented = true
+
+	return w.err
+}
+
+func (w *Writer) flushFragment(fin bool) error {
+	frame := ws.NewFrame(w.opCode(), fin, w.buf[:w.n])
+	if w.state.ClientSide() {
+		frame = ws.MaskFrameInPlace(frame)
+	}
+
+	// Write header to the header segment of the raw buffer.
+	head := len(w.raw) - len(w.buf)
+	offset := head - ws.HeaderSize(frame.Header)
+	buf := bytesWriter{
+		buf: w.raw[offset:head],
+	}
+	if err := ws.WriteHeader(&buf, frame.Header); err != nil {
+		// Must never be reached.
+		panic("dump header error: " + err.Error())
+	}
+
+	_, err := w.dest.Write(w.raw[offset : head+w.n])
+
+	return err
+}
+
+func (w *Writer) opCode() ws.OpCode {
+	if w.fragmented {
+		return ws.OpContinuation
+	}
+	return w.op
+}
+
+var errNoSpace = fmt.Errorf("not enough buffer space")
+
+type bytesWriter struct {
+	buf []byte
+	pos int
+}
+
+func (w *bytesWriter) Write(p []byte) (int, error) {
+	n := copy(w.buf[w.pos:], p)
+	w.pos += n
+	if n != len(p) {
+		return n, errNoSpace
+	}
+	return n, nil
+}
+
+func writeFrame(w io.Writer, s ws.State, op ws.OpCode, fin bool, p []byte) error {
+	var frame ws.Frame
+	if s.ClientSide() {
+		// Should copy bytes to prevent corruption of caller data.
+		payload := pbytes.GetLen(len(p))
+		defer pbytes.Put(payload)
+
+		copy(payload, p)
+
+		frame = ws.NewFrame(op, fin, payload)
+		frame = ws.MaskFrameInPlace(frame)
+	} else {
+		frame = ws.NewFrame(op, fin, p)
+	}
+
+	return ws.WriteFrame(w, frame)
+}

vendor/github.com/gobwas/ws/wsutil/wsutil.go

@@ -0,0 +1,57 @@
+/*
+Package wsutil provides utilities for working with WebSocket protocol.
+
+Overview:
+
+  // Read masked text message from peer and check utf8 encoding.
+  header, err := ws.ReadHeader(conn)
+  if err != nil {
+	  // handle err
+  }
+
+  // Prepare to read payload.
+  r := io.LimitReader(conn, header.Length)
+  r = wsutil.NewCipherReader(r, header.Mask)
+  r = wsutil.NewUTF8Reader(r)
+
+  payload, err := ioutil.ReadAll(r)
+  if err != nil {
+	  // handle err
+  }
+
+You could get the same behavior using just `wsutil.Reader`:
+
+  r := wsutil.Reader{
+	  Source:    conn,
+	  CheckUTF8: true,
+  }
+
+  payload, err := ioutil.ReadAll(r)
+  if err != nil {
+	  // handle err
+  }
+
+Or even simplest:
+
+  payload, err := wsutil.ReadClientText(conn)
+  if err != nil {
+	  // handle err
+  }
+
+Package is also exports tools for buffered writing:
+
+  // Create buffered writer, that will buffer output bytes and send them as
+  // 128-length fragments (with exception on large writes, see the doc).
+  writer := wsutil.NewWriterSize(conn, ws.StateServerSide, ws.OpText, 128)
+
+  _, err := io.CopyN(writer, rand.Reader, 100)
+  if err == nil {
+	  err = writer.Flush()
+  }
+  if err != nil {
+	  // handle error
+  }
+
+For more utils and helpers see the documentation.
+*/
+package wsutil