
/*---------------------------------------------------------------------------*/
/* websocketframe.cpp                                                        */
/* copyright (c) innovaphone 2015                                            */
/*                                                                           */
/*---------------------------------------------------------------------------*/

#include <platform.h>

#include "sdk/common/interface/websocket_client.h"
#include "sdk/common/interface/random.h"
#include "websocketframe.h"

#ifdef OVERRIDE_WEBSOCKET_MAX_PAYLOAD_SIZE
#ifdef WEBSOCKET_MAX_PAYLOAD_SIZE
#undef WEBSOCKET_MAX_PAYLOAD_SIZE
#endif
#define WEBSOCKET_MAX_PAYLOAD_SIZE OVERRIDE_WEBSOCKET_MAX_PAYLOAD_SIZE
#endif

#define BIT_SET(i, b) ((i & b) == b)

WebsocketFrame::WebsocketFrame()
    : _headerSize(0),
      _extLenAndMaskSize(0),
      _dataLength(0),
      _type(WSF_SINGLE),
      _serverMode(false)
{
    _buffer = (byte *)malloc(WEBSOCKET_MAX_PAYLOAD_SIZE + 14);
    memset(_buffer, 0, WEBSOCKET_MAX_PAYLOAD_SIZE + 14);
    _dataBuffer = nullptr;
}


WebsocketFrame::~WebsocketFrame()
{
    free(_buffer);
}


void WebsocketFrame::SetServerMode(bool serverMode)
{
    _serverMode = serverMode;
}


//   0                   1                   2                   3
//   0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
//  +-+-+-+-+-------+-+-------------+-------------------------------+
//  |F|R|R|R| opcode|M| Payload len |    Extended payload length    |
//  |I|S|S|S|  (4)  |A|     (7)     |             (16/64)           |
//  |N|V|V|V|       |S|             |   (if payload len==126/127)   |
//  | |1|2|3|       |K|             |                               |
//  +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
//  |     Extended payload length continued, if payload len == 127  |
//  + - - - - - - - - - - - - - - - +-------------------------------+
//  |                               |Masking-key, if MASK set to 1  |
//  +-------------------------------+-------------------------------+
//  | Masking-key (continued)       |          Payload Data         |
//  +-------------------------------- - - - - - - - - - - - - - - - +
//  :                     Payload Data continued ...                :
//  + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - +
//  |                     Payload Data continued ...                |
//  +---------------------------------------------------------------+
void WebsocketFrame::DecodeHeaderBase()
{
    _headerSize = 2;
    _extLenAndMaskSize = 0;
    _maskKeyStart = 0;

    _dataLength = _buffer[1] & 0x7F;

    if (_dataLength == 126)
        _extLenAndMaskSize += 2;
    else if (_dataLength == 127)
        _extLenAndMaskSize += 8;
    
    if (this->maskBit())
        _extLenAndMaskSize += 4;

    opcode_t oc = this->opcode();
    if (!this->finBit()) {
        if (oc != WSOC_CONTINOUS_FRAME)
            _type = WSF_FRAGMENTED_START;
        if (oc == WSOC_CONTINOUS_FRAME)
            _type = WSF_FRAGMENTED_CONTINOUS;
    }
    else {
        if (oc == WSOC_CONTINOUS_FRAME)
            _type = WSF_FRAGMENTED_END;
        else
            _type = WSF_SINGLE;
    }

    if (oc == WSOC_BINARY_FRAME)
        _type = (frame_t)(_type | WSF_BINARY);
}


void WebsocketFrame::DecodeHeaderExt()
{
    byte * b = &_buffer[2];

    if (_dataLength == 126) {
        _dataLength = READ_WORD_NET(b);
        _headerSize += 2;
    }
    else if (_dataLength == 127) {
        _dataLength = READ_QUAD_NET(b);
        _headerSize += 8;
    }

    if (this->maskBit()) {
        _maskKeyStart = (int)_headerSize;
        _headerSize += 4;
    }
}


void WebsocketFrame::EncodeFrame(frame_t ftype, opcode_t opcode, byte * data, size_t len)
{
    _headerSize = 2;
    _extLenAndMaskSize = 0;
    _maskKeyStart = 0;

    // The bit to encode or decode depends on the mode...
    if (BIT_SET(ftype, WSF_SINGLE) || BIT_SET(ftype, WSF_FRAGMENTED_END)) {
        _buffer[0] = opcode | 0x80;
    }
    else
        _buffer[0] = opcode;

    byte plen = 0;
    _dataLength = len;
    if (_dataLength < 126) {
        plen = (byte)_dataLength;
    }
    else if (_dataLength > 0xFFFF) {
        plen = 127;
        WRITE_QUAD_NET((dword *)&_buffer[2], _dataLength);
        _headerSize += 8;
        _extLenAndMaskSize += 8;
    }
    else {
        plen = 126;
        WRITE_WORD_NET((word *)&_buffer[2], (word)_dataLength);
        _headerSize += 2;
        _extLenAndMaskSize += 2;
    }
    
    if (_serverMode) {
        _buffer[1] = plen;
        if (data)
            _dataBuffer = data;
            //memcpy(&_buffer[_headerSize], data, len);
    }
    else {
        _buffer[1] = plen | 0x80;
        dword random = IRandom::GetRandom();
        memcpy(&_buffer[_headerSize], &random, sizeof(dword));
        //dword * tmp = (dword *)&_buffer[_headerSize];
        //*tmp = IRandom::GetRandom();
        _maskKeyStart = (int)_headerSize;
        _headerSize += 4;
        _extLenAndMaskSize += 4;

        if (data)
            this->MaskData(data, &_buffer[_headerSize], len);
    }
}


void WebsocketFrame::MaskData(byte * in, byte * out, size_t len)
{
    for (size_t i = 0; i < len; i++)
        out[i] = in[i] ^ _buffer[_maskKeyStart + (i & 3)];
}


void WebsocketFrame::EncodeDataFrame(frame_t ftype, void * data, size_t len)
{
    opcode_t oc;
    if (BIT_SET(ftype, WSF_FRAGMENTED_CONTINOUS) || BIT_SET(ftype, WSF_FRAGMENTED_END))
        oc = WSOC_CONTINOUS_FRAME;
    else if (BIT_SET(ftype, WSF_BINARY))
        oc = WSOC_BINARY_FRAME;
    else
        oc = WSOC_TEXT_FRAME;
    this->EncodeFrame(ftype, oc, (byte *)data, len);
}


void WebsocketFrame::EncodeCloseFrame(word status, const char * reason)
{
    frame_t ftype = WSF_SINGLE;

    size_t l = strlen(reason);
    ASSERT(l < WEBSOCKET_MAX_PAYLOAD_SIZE, "WebsocketFrame::EncodeCloseFrame() - reason longer than WEBSOCKET_MAX_PAYLOAD_SIZE");

    this->EncodeFrame(ftype, WSOC_CLOSE, NULL, l + 2);

    // Copy the reason manually....
    _dataBuffer = &_buffer[_headerSize];
    byte * data = _dataBuffer;
    WRITE_WORD_NET(data, status);
    if (l > 0) {
        data += 2;
        if (_serverMode)
            memcpy(data, reason, l);
        else
            this->MaskData((byte *)reason, data, l);
    }
}


void WebsocketFrame::EncodePingFrame()
{
    frame_t ftype = WSF_SINGLE;
    this->EncodeFrame(ftype, WSOC_PING, NULL, 0);
}


void WebsocketFrame::EncodePongFrame()
{
    frame_t ftype = WSF_SINGLE;
    this->EncodeFrame(ftype, WSOC_PONG, NULL, 0);
}
