ws/ws/message.go

package ws

import (
	"fmt"
	"io"
	"encoding/binary"
)

var UnmaskedFrameErr = fmt.Errorf("Received unmasked 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 = reader.Read(tmpBuf)
	if err != nil { return nil, err }


	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 = reader.Read(tmpBuf)
	if err != nil { return nil, 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 := reader.Read(tmpBuf)
		if err != nil { return nil, err }

		m.Size = uint( binary.BigEndian.Uint64(tmpBuf) )

	} else if m.Size == 126 {

		tmpBuf = make([]byte, 2)
		_, err := reader.Read(tmpBuf)
		if err != nil { return nil, err }

		m.Size = uint( binary.BigEndian.Uint16(tmpBuf) )

	}

	/* (5) Masking key */
	if mask != nil {

		tmpBuf = make([]byte, 4)
		_, err := reader.Read(tmpBuf)
		if err != nil { return nil, err }

		mask = make([]byte, 4)
		copy(mask, tmpBuf)

	}

	/* (6) Read payload by chunks */
	m.Data = make([]byte, int(m.Size))

	// If empty payload
	if m.Size <= 0 {
		return m, nil
	}

	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 nil, 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
	if mask == nil {
		return nil, UnmaskedFrameErr
	}

	return m, nil

}




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