280 lines
9.0 KiB
Go
280 lines
9.0 KiB
Go
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
|
|
}
|
|
}
|