/* SPDX-License-Identifier: MPL-2.0 */
 
#include "precompiled.hpp"
#include "ws_protocol.hpp"
#include "ws_encoder.hpp"
#include "msg.hpp"
#include "likely.hpp"
#include "wire.hpp"
#include "random.hpp"
 
#include <limits.h>
 
zmq::ws_encoder_t::ws_encoder_t (size_t bufsize_, bool must_mask_) :
    encoder_base_t<ws_encoder_t> (bufsize_), _must_mask (must_mask_)
{
    //  Write 0 bytes to the batch and go to message_ready state.
    next_step (NULL, 0, &ws_encoder_t::message_ready, true);
    _masked_msg.init ();
}
 
zmq::ws_encoder_t::~ws_encoder_t ()
{
    _masked_msg.close ();
}
 
void zmq::ws_encoder_t::message_ready ()
{
    int offset = 0;
 
    _is_binary = false;
 
    if (in_progress ()->is_ping ())
        _tmp_buf[offset++] = 0x80 | zmq::ws_protocol_t::opcode_ping;
    else if (in_progress ()->is_pong ())
        _tmp_buf[offset++] = 0x80 | zmq::ws_protocol_t::opcode_pong;
    else if (in_progress ()->is_close_cmd ())
        _tmp_buf[offset++] = 0x80 | zmq::ws_protocol_t::opcode_close;
    else {
        _tmp_buf[offset++] = 0x82; // Final | binary
        _is_binary = true;
    }
 
    _tmp_buf[offset] = _must_mask ? 0x80 : 0x00;
 
    size_t size = in_progress ()->size ();
    if (_is_binary)
        size++;
    //  TODO: create an opcode for subscribe/cancel
    if (in_progress ()->is_subscribe () || in_progress ()->is_cancel ())
        size++;
 
    if (size <= 125)
        _tmp_buf[offset++] |= static_cast<unsigned char> (size & 127);
    else if (size <= 0xFFFF) {
        _tmp_buf[offset++] |= 126;
        _tmp_buf[offset++] = static_cast<unsigned char> ((size >> 8) & 0xFF);
        _tmp_buf[offset++] = static_cast<unsigned char> (size & 0xFF);
    } else {
        _tmp_buf[offset++] |= 127;
        put_uint64 (_tmp_buf + offset, size);
        offset += 8;
    }
 
    if (_must_mask) {
        const uint32_t random = generate_random ();
        put_uint32 (_tmp_buf + offset, random);
        put_uint32 (_mask, random);
        offset += 4;
    }
 
    int mask_index = 0;
    if (_is_binary) {
        //  Encode flags.
        unsigned char protocol_flags = 0;
        if (in_progress ()->flags () & msg_t::more)
            protocol_flags |= ws_protocol_t::more_flag;
        if (in_progress ()->flags () & msg_t::command)
            protocol_flags |= ws_protocol_t::command_flag;
 
        _tmp_buf[offset++] =
          _must_mask ? protocol_flags ^ _mask[mask_index++] : protocol_flags;
    }
 
    //  Encode the subscribe/cancel byte.
    //  TODO: remove once there is an opcode for subscribe/cancel
    if (in_progress ()->is_subscribe ())
        _tmp_buf[offset++] = _must_mask ? 1 ^ _mask[mask_index++] : 1;
    else if (in_progress ()->is_cancel ())
        _tmp_buf[offset++] = _must_mask ? 0 ^ _mask[mask_index++] : 0;
 
    next_step (_tmp_buf, offset, &ws_encoder_t::size_ready, false);
}
 
void zmq::ws_encoder_t::size_ready ()
{
    if (_must_mask) {
        assert (in_progress () != &_masked_msg);
        const size_t size = in_progress ()->size ();
 
        unsigned char *src =
          static_cast<unsigned char *> (in_progress ()->data ());
        unsigned char *dest = src;
 
        //  If msg is shared or data is constant we cannot mask in-place, allocate a new msg for it
        if (in_progress ()->flags () & msg_t::shared
            || in_progress ()->is_cmsg ()) {
            _masked_msg.close ();
            _masked_msg.init_size (size);
            dest = static_cast<unsigned char *> (_masked_msg.data ());
        }
 
        int mask_index = 0;
        if (_is_binary)
            ++mask_index;
        //  TODO: remove once there is an opcode for subscribe/cancel
        if (in_progress ()->is_subscribe () || in_progress ()->is_cancel ())
            ++mask_index;
        for (size_t i = 0; i < size; ++i, mask_index++)
            dest[i] = src[i] ^ _mask[mask_index % 4];
 
        next_step (dest, size, &ws_encoder_t::message_ready, true);
    } else {
        next_step (in_progress ()->data (), in_progress ()->size (),
                   &ws_encoder_t::message_ready, true);
    }
}