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

280 lines
9 KiB
Go
Raw Permalink Normal View History

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