/* SPDX-License-Identifier: MPL-2.0 */

#include "precompiled.hpp"

#include <stdlib.h>

#include <string.h>

#include <cmath>

#include "ws_protocol.hpp"

#include "ws_decoder.hpp"

#include "likely.hpp"

#include "wire.hpp"

#include "err.hpp"

zmq::ws_decoder_t::ws_decoder_t (size_t bufsize_,

                                 int64_t maxmsgsize_,

                                 bool zero_copy_,

                                 bool must_mask_) :

    decoder_base_t<ws_decoder_t, shared_message_memory_allocator> (bufsize_),

    _msg_flags (0),

    _zero_copy (zero_copy_),

    _max_msg_size (maxmsgsize_),

    _must_mask (must_mask_),

    _size (0)

{

    memset (_tmpbuf, 0, sizeof (_tmpbuf));

    int rc = _in_progress.init ();

    errno_assert (rc == 0);

    //  At the beginning, read one byte and go to opcode_ready state.

    next_step (_tmpbuf, 1, &ws_decoder_t::opcode_ready);

}

zmq::ws_decoder_t::~ws_decoder_t ()

{

    const int rc = _in_progress.close ();

    errno_assert (rc == 0);

}

int zmq::ws_decoder_t::opcode_ready (unsigned char const *)

{

    const bool final = (_tmpbuf[0] & 0x80) != 0; // final bit

    if (!final)

        return -1; // non final messages are not supported

    _opcode = static_cast<zmq::ws_protocol_t::opcode_t> (_tmpbuf[0] & 0xF);

    _msg_flags = 0;

    switch (_opcode) {

        case zmq::ws_protocol_t::opcode_binary:

            break;

        case zmq::ws_protocol_t::opcode_close:

            _msg_flags = msg_t::command | msg_t::close_cmd;

            break;

        case zmq::ws_protocol_t::opcode_ping:

            _msg_flags = msg_t::ping | msg_t::command;

            break;

        case zmq::ws_protocol_t::opcode_pong:

            _msg_flags = msg_t::pong | msg_t::command;

            break;

        default:

            return -1;

    }

    next_step (_tmpbuf, 1, &ws_decoder_t::size_first_byte_ready);

    return 0;

}

int zmq::ws_decoder_t::size_first_byte_ready (unsigned char const *read_from_)

{

    const bool is_masked = (_tmpbuf[0] & 0x80) != 0;

    if (is_masked != _must_mask) // wrong mask value

        return -1;

    _size = static_cast<uint64_t> (_tmpbuf[0] & 0x7F);

    if (_size < 126) {

        if (_must_mask)

            next_step (_tmpbuf, 4, &ws_decoder_t::mask_ready);

        else if (_opcode == ws_protocol_t::opcode_binary) {

            if (_size == 0)

                return -1;

            next_step (_tmpbuf, 1, &ws_decoder_t::flags_ready);

        } else

            return size_ready (read_from_);

    } else if (_size == 126)

        next_step (_tmpbuf, 2, &ws_decoder_t::short_size_ready);

    else

        next_step (_tmpbuf, 8, &ws_decoder_t::long_size_ready);

    return 0;

}

int zmq::ws_decoder_t::short_size_ready (unsigned char const *read_from_)

{

    _size = (_tmpbuf[0] << 8) | _tmpbuf[1];

    if (_must_mask)

        next_step (_tmpbuf, 4, &ws_decoder_t::mask_ready);

    else if (_opcode == ws_protocol_t::opcode_binary) {

        if (_size == 0)

            return -1;

        next_step (_tmpbuf, 1, &ws_decoder_t::flags_ready);

    } else

        return size_ready (read_from_);

    return 0;

}

int zmq::ws_decoder_t::long_size_ready (unsigned char const *read_from_)

{

    //  The payload size is encoded as 64-bit unsigned integer.

    //  The most significant byte comes first.

    _size = get_uint64 (_tmpbuf);

    if (_must_mask)

        next_step (_tmpbuf, 4, &ws_decoder_t::mask_ready);

    else if (_opcode == ws_protocol_t::opcode_binary) {

        if (_size == 0)

            return -1;

        next_step (_tmpbuf, 1, &ws_decoder_t::flags_ready);

    } else

        return size_ready (read_from_);

    return 0;

}

int zmq::ws_decoder_t::mask_ready (unsigned char const *read_from_)

{

    memcpy (_mask, _tmpbuf, 4);

    if (_opcode == ws_protocol_t::opcode_binary) {

        if (_size == 0)

            return -1;

        next_step (_tmpbuf, 1, &ws_decoder_t::flags_ready);

    } else

        return size_ready (read_from_);

    return 0;

}

int zmq::ws_decoder_t::flags_ready (unsigned char const *read_from_)

{

    unsigned char flags;

    if (_must_mask)

        flags = _tmpbuf[0] ^ _mask[0];

    else

        flags = _tmpbuf[0];

    if (flags & ws_protocol_t::more_flag)

        _msg_flags |= msg_t::more;

    if (flags & ws_protocol_t::command_flag)

        _msg_flags |= msg_t::command;

    _size--;

    return size_ready (read_from_);

}

int zmq::ws_decoder_t::size_ready (unsigned char const *read_pos_)

{

    //  Message size must not exceed the maximum allowed size.

    if (_max_msg_size >= 0)

        if (unlikely (_size > static_cast<uint64_t> (_max_msg_size))) {

            errno = EMSGSIZE;

            return -1;

        }

    //  Message size must fit into size_t data type.

    if (unlikely (_size != static_cast<size_t> (_size))) {

        errno = EMSGSIZE;

        return -1;

    }

    int rc = _in_progress.close ();

    assert (rc == 0);

    // the current message can exceed the current buffer. We have to copy the buffer

    // data into a new message and complete it in the next receive.

    shared_message_memory_allocator &allocator = get_allocator ();

    if (unlikely (!_zero_copy || allocator.data () > read_pos_

                  || static_cast<size_t> (read_pos_ - allocator.data ())

                       > allocator.size ()

                  || _size > static_cast<size_t> (

                       allocator.data () + allocator.size () - read_pos_))) {

        // a new message has started, but the size would exceed the pre-allocated arena

        // (or read_pos_ is in the initial handshake buffer)

        // this happens every time when a message does not fit completely into the buffer

        rc = _in_progress.init_size (static_cast<size_t> (_size));

    } else {

        // construct message using n bytes from the buffer as storage

        // increase buffer ref count

        // if the message will be a large message, pass a valid refcnt memory location as well

        rc = _in_progress.init (

          const_cast<unsigned char *> (read_pos_), static_cast<size_t> (_size),

          shared_message_memory_allocator::call_dec_ref, allocator.buffer (),

          allocator.provide_content ());

        // For small messages, data has been copied and refcount does not have to be increased

        if (_in_progress.is_zcmsg ()) {

            allocator.advance_content ();

            allocator.inc_ref ();

        }

    }

    if (unlikely (rc)) {

        errno_assert (errno == ENOMEM);

        rc = _in_progress.init ();

        errno_assert (rc == 0);

        errno = ENOMEM;

        return -1;

    }

    _in_progress.set_flags (_msg_flags);

    // this sets read_pos to

    // the message data address if the data needs to be copied

    // for small message / messages exceeding the current buffer

    // or

    // to the current start address in the buffer because the message

    // was constructed to use n bytes from the address passed as argument

    next_step (_in_progress.data (), _in_progress.size (),

               &ws_decoder_t::message_ready);

    return 0;

}

int zmq::ws_decoder_t::message_ready (unsigned char const *)

{

    if (_must_mask) {

        int mask_index = _opcode == ws_protocol_t::opcode_binary ? 1 : 0;

        unsigned char *data =

          static_cast<unsigned char *> (_in_progress.data ());

        for (size_t i = 0; i < _size; ++i, mask_index++)

            data[i] = data[i] ^ _mask[mask_index % 4];

    }

    //  Message is completely read. Signal this to the caller

    //  and prepare to decode next message.

    next_step (_tmpbuf, 1, &ws_decoder_t::opcode_ready);

    return 1;

}