340 lines
6.9 KiB
Go
340 lines
6.9 KiB
Go
package websocket
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"fmt"
|
|
"io"
|
|
"unicode/utf8"
|
|
)
|
|
|
|
var ErrUnmaskedFrame = fmt.Errorf("Received unmasked frame")
|
|
var ErrTooLongControlFrame = fmt.Errorf("Received a control frame that is fragmented or too long")
|
|
var ErrInvalidFragment = fmt.Errorf("Received invalid fragmentation")
|
|
var ErrUnexpectedContinuation = fmt.Errorf("Received unexpected continuation frame")
|
|
var ErrInvalidSize = fmt.Errorf("Received invalid payload size")
|
|
var ErrInvalidPayload = fmt.Errorf("Received invalid utf8 payload")
|
|
var ErrInvalidCloseStatus = fmt.Errorf("Received invalid close status")
|
|
var ErrInvalidOpCode = fmt.Errorf("Received invalid OpCode")
|
|
var ErrReservedBits = fmt.Errorf("Received reserved bits")
|
|
var CloseFrame = fmt.Errorf("Received close Frame")
|
|
|
|
// Maximum Header Size = Final/OpCode + isMask/Length + Length + Mask
|
|
const maximumHeaderSize = 1 + 1 + 8 + 4
|
|
const maxWriteChunk = 0x7fff
|
|
|
|
// Lists websocket close status
|
|
type MessageError uint16
|
|
|
|
const (
|
|
NONE MessageError = 0
|
|
NORMAL MessageError = 1000
|
|
GOING_AWAY MessageError = 1001
|
|
PROTOCOL_ERR MessageError = 1002
|
|
UNACCEPTABLE_OPCODE MessageError = 1003
|
|
INVALID_PAYLOAD MessageError = 1007 // utf8
|
|
MESSAGE_TOO_LARGE MessageError = 1009
|
|
)
|
|
|
|
// Lists websocket message types
|
|
type MessageType byte
|
|
|
|
const (
|
|
CONTINUATION MessageType = 0x00
|
|
TEXT MessageType = 0x01
|
|
BINARY MessageType = 0x02
|
|
CLOSE MessageType = 0x08
|
|
PING MessageType = 0x09
|
|
PONG MessageType = 0x0a
|
|
)
|
|
|
|
// Represents a websocket message
|
|
type Message struct {
|
|
Final bool
|
|
Type MessageType
|
|
Size uint
|
|
Data []byte
|
|
}
|
|
|
|
// receive reads a message form reader
|
|
func readMessage(reader io.Reader) (*Message, error) {
|
|
|
|
var err error
|
|
var tmpBuf []byte
|
|
var mask []byte
|
|
var cursor int
|
|
|
|
m := new(Message)
|
|
|
|
/* (2) Byte 1: FIN and OpCode */
|
|
tmpBuf = make([]byte, 1)
|
|
err = readBytes(reader, tmpBuf)
|
|
if err != nil {
|
|
return m, err
|
|
}
|
|
|
|
// check reserved bits
|
|
if tmpBuf[0]&0x70 != 0 {
|
|
return m, ErrReservedBits
|
|
}
|
|
|
|
m.Final = bool(tmpBuf[0]&0x80 == 0x80)
|
|
m.Type = MessageType(tmpBuf[0] & 0x0f)
|
|
|
|
/* (3) Byte 2: Mask and Length[0] */
|
|
tmpBuf = make([]byte, 1)
|
|
err = readBytes(reader, tmpBuf)
|
|
if err != nil {
|
|
return m, err
|
|
}
|
|
|
|
// if mask, byte array not nil
|
|
if tmpBuf[0]&0x80 == 0x80 {
|
|
mask = make([]byte, 0)
|
|
}
|
|
|
|
// payload length
|
|
m.Size = uint(tmpBuf[0] & 0x7f)
|
|
|
|
/* (4) Extended payload */
|
|
if m.Size == 127 {
|
|
|
|
tmpBuf = make([]byte, 8)
|
|
err := readBytes(reader, tmpBuf)
|
|
if err != nil {
|
|
return m, err
|
|
}
|
|
|
|
m.Size = uint(binary.BigEndian.Uint64(tmpBuf))
|
|
|
|
} else if m.Size == 126 {
|
|
|
|
tmpBuf = make([]byte, 2)
|
|
err := readBytes(reader, tmpBuf)
|
|
if err != nil {
|
|
return m, err
|
|
}
|
|
|
|
m.Size = uint(binary.BigEndian.Uint16(tmpBuf))
|
|
|
|
}
|
|
|
|
/* (5) Masking key */
|
|
if mask != nil {
|
|
|
|
tmpBuf = make([]byte, 4)
|
|
err := readBytes(reader, tmpBuf)
|
|
if err != nil {
|
|
return m, err
|
|
}
|
|
|
|
mask = make([]byte, 4)
|
|
copy(mask, tmpBuf)
|
|
|
|
}
|
|
|
|
/* (6) Read payload by chunks */
|
|
m.Data = make([]byte, int(m.Size))
|
|
|
|
cursor = 0
|
|
|
|
// {1} While we have data to read //
|
|
for uint(cursor) < m.Size {
|
|
|
|
// {2} Try to read (at least 1 byte) //
|
|
nbread, err := io.ReadAtLeast(reader, m.Data[cursor:m.Size], 1)
|
|
if err != nil {
|
|
return m, err
|
|
}
|
|
|
|
// {3} Unmask data //
|
|
if mask != nil {
|
|
for i, l := cursor, cursor+nbread; i < l; i++ {
|
|
|
|
mi := i % 4 // mask index
|
|
m.Data[i] = m.Data[i] ^ mask[mi]
|
|
|
|
}
|
|
}
|
|
|
|
// {4} Update cursor //
|
|
cursor += nbread
|
|
|
|
}
|
|
|
|
// return error if unmasked frame
|
|
// we have to fully read it for read buffer to be clean
|
|
err = nil
|
|
if mask == nil {
|
|
err = ErrUnmaskedFrame
|
|
}
|
|
|
|
return m, err
|
|
|
|
}
|
|
|
|
// Send sends a frame over a socket
|
|
func (m Message) Send(writer io.Writer) error {
|
|
|
|
header := make([]byte, 0, maximumHeaderSize)
|
|
|
|
// fix size
|
|
if uint(len(m.Data)) <= m.Size {
|
|
m.Size = uint(len(m.Data))
|
|
}
|
|
|
|
/* (1) Byte 0 : FIN + opcode */
|
|
var final byte = 0x80
|
|
if !m.Final {
|
|
final = 0
|
|
}
|
|
header = append(header, final|byte(m.Type))
|
|
|
|
/* (2) Get payload length */
|
|
if m.Size < 126 { // simple
|
|
|
|
header = append(header, byte(m.Size))
|
|
|
|
} else if m.Size <= 0xffff { // extended: 16 bits
|
|
|
|
header = append(header, 126)
|
|
|
|
buf := make([]byte, 2)
|
|
binary.BigEndian.PutUint16(buf, uint16(m.Size))
|
|
header = append(header, buf...)
|
|
|
|
} else if m.Size <= 0xffffffffffffffff { // extended: 64 bits
|
|
|
|
header = append(header, 127)
|
|
|
|
buf := make([]byte, 8)
|
|
binary.BigEndian.PutUint64(buf, uint64(m.Size))
|
|
header = append(header, buf...)
|
|
|
|
}
|
|
|
|
/* (3) Build write buffer */
|
|
writeBuf := make([]byte, 0, len(header)+int(m.Size))
|
|
writeBuf = append(writeBuf, header...)
|
|
writeBuf = append(writeBuf, m.Data[0:m.Size]...)
|
|
|
|
/* (4) Send over socket by chunks */
|
|
toWrite := len(header) + int(m.Size)
|
|
cursor := 0
|
|
for cursor < toWrite {
|
|
|
|
maxBoundary := cursor + maxWriteChunk
|
|
if maxBoundary > toWrite {
|
|
maxBoundary = toWrite
|
|
}
|
|
|
|
// Try to wrote (at max 1024 bytes) //
|
|
nbwritten, err := writer.Write(writeBuf[cursor:maxBoundary])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Update cursor //
|
|
cursor += nbwritten
|
|
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// Check for message errors with:
|
|
// (m) the current message
|
|
// (fragment) whether there is a fragment in construction
|
|
// returns the message error
|
|
func (m *Message) check(fragment bool) error {
|
|
|
|
/* (1) Invalid first fragment (not TEXT nor BINARY) */
|
|
if !m.Final && !fragment && m.Type != TEXT && m.Type != BINARY {
|
|
return ErrInvalidFragment
|
|
}
|
|
|
|
/* (2) Waiting fragment but received standalone frame */
|
|
if fragment && m.Type != CONTINUATION && m.Type != CLOSE && m.Type != PING && m.Type != PONG {
|
|
return ErrInvalidFragment
|
|
}
|
|
|
|
/* (3) Control frame too long */
|
|
if (m.Type == CLOSE || m.Type == PING || m.Type == PONG) && (m.Size > 125 || !m.Final) {
|
|
return ErrTooLongControlFrame
|
|
}
|
|
|
|
switch m.Type {
|
|
case CONTINUATION:
|
|
// unexpected continuation
|
|
if !fragment {
|
|
return ErrUnexpectedContinuation
|
|
}
|
|
return nil
|
|
|
|
case TEXT:
|
|
if m.Final && !utf8.Valid(m.Data) {
|
|
return ErrInvalidPayload
|
|
}
|
|
return nil
|
|
|
|
case BINARY:
|
|
return nil
|
|
|
|
case CLOSE:
|
|
// incomplete code
|
|
if m.Size == 1 {
|
|
return ErrInvalidCloseStatus
|
|
}
|
|
|
|
// invalid utf-8 reason
|
|
if m.Size > 2 && !utf8.Valid(m.Data[2:]) {
|
|
return ErrInvalidPayload
|
|
}
|
|
|
|
// invalid code
|
|
if m.Size >= 2 {
|
|
c := binary.BigEndian.Uint16(m.Data[0:2])
|
|
|
|
if c < 1000 || c >= 1004 && c <= 1006 || c >= 1012 && c <= 1016 || c == 1100 || c == 2000 || c == 2999 {
|
|
return ErrInvalidCloseStatus
|
|
}
|
|
}
|
|
return CloseFrame
|
|
|
|
case PING:
|
|
return nil
|
|
|
|
case PONG:
|
|
return nil
|
|
|
|
default:
|
|
return ErrInvalidOpCode
|
|
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// readBytes reads from a reader into a byte array
|
|
// until the byte length is fully filled with data
|
|
// loops while there is no error
|
|
//
|
|
// It manages connections which chunks data
|
|
func readBytes(reader io.Reader, buffer []byte) error {
|
|
|
|
var cur, len int = 0, len(buffer)
|
|
|
|
// try to read until the full size is read
|
|
for cur < len {
|
|
|
|
nbread, err := reader.Read(buffer[cur:])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
cur += nbread
|
|
}
|
|
|
|
return nil
|
|
|
|
}
|