/* SPDX-License-Identifier: MPL-2.0 */
 
#include "precompiled.hpp"
#include "macros.hpp"
 
#include <limits.h>
#include <string.h>
 
#ifndef ZMQ_HAVE_WINDOWS
#include <unistd.h>
#endif
 
#include <new>
#include <sstream>
 
#include "zmtp_engine.hpp"
#include "io_thread.hpp"
#include "session_base.hpp"
#include "v1_encoder.hpp"
#include "v1_decoder.hpp"
#include "v2_encoder.hpp"
#include "v2_decoder.hpp"
#include "v3_1_encoder.hpp"
#include "null_mechanism.hpp"
#include "plain_client.hpp"
#include "plain_server.hpp"
#include "gssapi_client.hpp"
#include "gssapi_server.hpp"
#include "curve_client.hpp"
#include "curve_server.hpp"
#include "raw_decoder.hpp"
#include "raw_encoder.hpp"
#include "config.hpp"
#include "err.hpp"
#include "ip.hpp"
#include "likely.hpp"
#include "wire.hpp"
 
zmq::zmtp_engine_t::zmtp_engine_t (
  fd_t fd_,
  const options_t &options_,
  const endpoint_uri_pair_t &endpoint_uri_pair_) :
    stream_engine_base_t (fd_, options_, endpoint_uri_pair_, true),
    _greeting_size (v2_greeting_size),
    _greeting_bytes_read (0),
    _subscription_required (false),
    _heartbeat_timeout (0)
{
    _next_msg = static_cast<int (stream_engine_base_t::*) (msg_t *)> (
      &zmtp_engine_t::routing_id_msg);
    _process_msg = static_cast<int (stream_engine_base_t::*) (msg_t *)> (
      &zmtp_engine_t::process_routing_id_msg);
 
    int rc = _pong_msg.init ();
    errno_assert (rc == 0);
 
    rc = _routing_id_msg.init ();
    errno_assert (rc == 0);
 
    if (_options.heartbeat_interval > 0) {
        _heartbeat_timeout = _options.heartbeat_timeout;
        if (_heartbeat_timeout == -1)
            _heartbeat_timeout = _options.heartbeat_interval;
    }
}
 
zmq::zmtp_engine_t::~zmtp_engine_t ()
{
    const int rc = _routing_id_msg.close ();
    errno_assert (rc == 0);
}
 
void zmq::zmtp_engine_t::plug_internal ()
{
    // start optional timer, to prevent handshake hanging on no input
    set_handshake_timer ();
 
    //  Send the 'length' and 'flags' fields of the routing id message.
    //  The 'length' field is encoded in the long format.
    _outpos = _greeting_send;
    _outpos[_outsize++] = UCHAR_MAX;
    put_uint64 (&_outpos[_outsize], _options.routing_id_size + 1);
    _outsize += 8;
    _outpos[_outsize++] = 0x7f;
 
    set_pollin ();
    set_pollout ();
    //  Flush all the data that may have been already received downstream.
    in_event ();
}
 
//  Position of the revision and minor fields in the greeting.
const size_t revision_pos = 10;
const size_t minor_pos = 11;
 
bool zmq::zmtp_engine_t::handshake ()
{
    zmq_assert (_greeting_bytes_read < _greeting_size);
    //  Receive the greeting.
    const int rc = receive_greeting ();
    if (rc == -1)
        return false;
    const bool unversioned = rc != 0;
 
    if (!(this
            ->*select_handshake_fun (unversioned, _greeting_recv[revision_pos],
                                     _greeting_recv[minor_pos])) ())
        return false;
 
    // Start polling for output if necessary.
    if (_outsize == 0)
        set_pollout ();
 
    return true;
}
 
int zmq::zmtp_engine_t::receive_greeting ()
{
    bool unversioned = false;
    while (_greeting_bytes_read < _greeting_size) {
        const int n = read (_greeting_recv + _greeting_bytes_read,
                            _greeting_size - _greeting_bytes_read);
        if (n == -1) {
            if (errno != EAGAIN)
                error (connection_error);
            return -1;
        }
 
        _greeting_bytes_read += n;
 
        //  We have received at least one byte from the peer.
        //  If the first byte is not 0xff, we know that the
        //  peer is using unversioned protocol.
        if (_greeting_recv[0] != 0xff) {
            unversioned = true;
            break;
        }
 
        if (_greeting_bytes_read < signature_size)
            continue;
 
        //  Inspect the right-most bit of the 10th byte (which coincides
        //  with the 'flags' field if a regular message was sent).
        //  Zero indicates this is a header of a routing id message
        //  (i.e. the peer is using the unversioned protocol).
        if (!(_greeting_recv[9] & 0x01)) {
            unversioned = true;
            break;
        }
 
        //  The peer is using versioned protocol.
        receive_greeting_versioned ();
    }
    return unversioned ? 1 : 0;
}
 
void zmq::zmtp_engine_t::receive_greeting_versioned ()
{
    //  Send the major version number.
    if (_outpos + _outsize == _greeting_send + signature_size) {
        if (_outsize == 0)
            set_pollout ();
        _outpos[_outsize++] = 3; //  Major version number
    }
 
    if (_greeting_bytes_read > signature_size) {
        if (_outpos + _outsize == _greeting_send + signature_size + 1) {
            if (_outsize == 0)
                set_pollout ();
 
            //  Use ZMTP/2.0 to talk to older peers.
            if (_greeting_recv[revision_pos] == ZMTP_1_0
                || _greeting_recv[revision_pos] == ZMTP_2_0)
                _outpos[_outsize++] = _options.type;
            else {
                _outpos[_outsize++] = 1; //  Minor version number
                memset (_outpos + _outsize, 0, 20);
 
                zmq_assert (_options.mechanism == ZMQ_NULL
                            || _options.mechanism == ZMQ_PLAIN
                            || _options.mechanism == ZMQ_CURVE
                            || _options.mechanism == ZMQ_GSSAPI);
 
                if (_options.mechanism == ZMQ_NULL)
                    memcpy (_outpos + _outsize, "NULL", 4);
                else if (_options.mechanism == ZMQ_PLAIN)
                    memcpy (_outpos + _outsize, "PLAIN", 5);
                else if (_options.mechanism == ZMQ_GSSAPI)
                    memcpy (_outpos + _outsize, "GSSAPI", 6);
                else if (_options.mechanism == ZMQ_CURVE)
                    memcpy (_outpos + _outsize, "CURVE", 5);
                _outsize += 20;
                memset (_outpos + _outsize, 0, 32);
                _outsize += 32;
                _greeting_size = v3_greeting_size;
            }
        }
    }
}
 
zmq::zmtp_engine_t::handshake_fun_t zmq::zmtp_engine_t::select_handshake_fun (
  bool unversioned_, unsigned char revision_, unsigned char minor_)
{
    //  Is the peer using ZMTP/1.0 with no revision number?
    if (unversioned_) {
        return &zmtp_engine_t::handshake_v1_0_unversioned;
    }
    switch (revision_) {
        case ZMTP_1_0:
            return &zmtp_engine_t::handshake_v1_0;
        case ZMTP_2_0:
            return &zmtp_engine_t::handshake_v2_0;
        case ZMTP_3_x:
            switch (minor_) {
                case 0:
                    return &zmtp_engine_t::handshake_v3_0;
                default:
                    return &zmtp_engine_t::handshake_v3_1;
            }
        default:
            return &zmtp_engine_t::handshake_v3_1;
    }
}
 
bool zmq::zmtp_engine_t::handshake_v1_0_unversioned ()
{
    //  We send and receive rest of routing id message
    if (session ()->zap_enabled ()) {
        // reject ZMTP 1.0 connections if ZAP is enabled
        error (protocol_error);
        return false;
    }
 
    _encoder = new (std::nothrow) v1_encoder_t (_options.out_batch_size);
    alloc_assert (_encoder);
 
    _decoder = new (std::nothrow)
      v1_decoder_t (_options.in_batch_size, _options.maxmsgsize);
    alloc_assert (_decoder);
 
    //  We have already sent the message header.
    //  Since there is no way to tell the encoder to
    //  skip the message header, we simply throw that
    //  header data away.
    const size_t header_size =
      _options.routing_id_size + 1 >= UCHAR_MAX ? 10 : 2;
    unsigned char tmp[10], *bufferp = tmp;
 
    //  Prepare the routing id message and load it into encoder.
    //  Then consume bytes we have already sent to the peer.
    int rc = _routing_id_msg.close ();
    zmq_assert (rc == 0);
    rc = _routing_id_msg.init_size (_options.routing_id_size);
    zmq_assert (rc == 0);
    memcpy (_routing_id_msg.data (), _options.routing_id,
            _options.routing_id_size);
    _encoder->load_msg (&_routing_id_msg);
    const size_t buffer_size = _encoder->encode (&bufferp, header_size);
    zmq_assert (buffer_size == header_size);
 
    //  Make sure the decoder sees the data we have already received.
    _inpos = _greeting_recv;
    _insize = _greeting_bytes_read;
 
    //  To allow for interoperability with peers that do not forward
    //  their subscriptions, we inject a phantom subscription message
    //  message into the incoming message stream.
    if (_options.type == ZMQ_PUB || _options.type == ZMQ_XPUB)
        _subscription_required = true;
 
    //  We are sending our routing id now and the next message
    //  will come from the socket.
    _next_msg = &zmtp_engine_t::pull_msg_from_session;
 
    //  We are expecting routing id message.
    _process_msg = static_cast<int (stream_engine_base_t::*) (msg_t *)> (
      &zmtp_engine_t::process_routing_id_msg);
 
    return true;
}
 
bool zmq::zmtp_engine_t::handshake_v1_0 ()
{
    if (session ()->zap_enabled ()) {
        // reject ZMTP 1.0 connections if ZAP is enabled
        error (protocol_error);
        return false;
    }
 
    _encoder = new (std::nothrow) v1_encoder_t (_options.out_batch_size);
    alloc_assert (_encoder);
 
    _decoder = new (std::nothrow)
      v1_decoder_t (_options.in_batch_size, _options.maxmsgsize);
    alloc_assert (_decoder);
 
    return true;
}
 
bool zmq::zmtp_engine_t::handshake_v2_0 ()
{
    if (session ()->zap_enabled ()) {
        // reject ZMTP 2.0 connections if ZAP is enabled
        error (protocol_error);
        return false;
    }
 
    _encoder = new (std::nothrow) v2_encoder_t (_options.out_batch_size);
    alloc_assert (_encoder);
 
    _decoder = new (std::nothrow) v2_decoder_t (
      _options.in_batch_size, _options.maxmsgsize, _options.zero_copy);
    alloc_assert (_decoder);
 
    return true;
}
 
bool zmq::zmtp_engine_t::handshake_v3_x (const bool downgrade_sub_)
{
    if (_options.mechanism == ZMQ_NULL
        && memcmp (_greeting_recv + 12, "NULL\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",
                   20)
             == 0) {
        _mechanism = new (std::nothrow)
          null_mechanism_t (session (), _peer_address, _options);
        alloc_assert (_mechanism);
    } else if (_options.mechanism == ZMQ_PLAIN
               && memcmp (_greeting_recv + 12,
                          "PLAIN\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0", 20)
                    == 0) {
        if (_options.as_server)
            _mechanism = new (std::nothrow)
              plain_server_t (session (), _peer_address, _options);
        else
            _mechanism =
              new (std::nothrow) plain_client_t (session (), _options);
        alloc_assert (_mechanism);
    }
#ifdef ZMQ_HAVE_CURVE
    else if (_options.mechanism == ZMQ_CURVE
             && memcmp (_greeting_recv + 12,
                        "CURVE\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0", 20)
                  == 0) {
        if (_options.as_server)
            _mechanism = new (std::nothrow) curve_server_t (
              session (), _peer_address, _options, downgrade_sub_);
        else
            _mechanism = new (std::nothrow)
              curve_client_t (session (), _options, downgrade_sub_);
        alloc_assert (_mechanism);
    }
#endif
#ifdef HAVE_LIBGSSAPI_KRB5
    else if (_options.mechanism == ZMQ_GSSAPI
             && memcmp (_greeting_recv + 12,
                        "GSSAPI\0\0\0\0\0\0\0\0\0\0\0\0\0\0", 20)
                  == 0) {
        if (_options.as_server)
            _mechanism = new (std::nothrow)
              gssapi_server_t (session (), _peer_address, _options);
        else
            _mechanism =
              new (std::nothrow) gssapi_client_t (session (), _options);
        alloc_assert (_mechanism);
    }
#endif
    else {
        socket ()->event_handshake_failed_protocol (
          session ()->get_endpoint (),
          ZMQ_PROTOCOL_ERROR_ZMTP_MECHANISM_MISMATCH);
        error (protocol_error);
        return false;
    }
#ifndef ZMQ_HAVE_CURVE
    LIBZMQ_UNUSED (downgrade_sub_);
#endif
    _next_msg = &zmtp_engine_t::next_handshake_command;
    _process_msg = &zmtp_engine_t::process_handshake_command;
 
    return true;
}
 
bool zmq::zmtp_engine_t::handshake_v3_0 ()
{
    _encoder = new (std::nothrow) v2_encoder_t (_options.out_batch_size);
    alloc_assert (_encoder);
 
    _decoder = new (std::nothrow) v2_decoder_t (
      _options.in_batch_size, _options.maxmsgsize, _options.zero_copy);
    alloc_assert (_decoder);
 
    return zmq::zmtp_engine_t::handshake_v3_x (true);
}
 
bool zmq::zmtp_engine_t::handshake_v3_1 ()
{
    _encoder = new (std::nothrow) v3_1_encoder_t (_options.out_batch_size);
    alloc_assert (_encoder);
 
    _decoder = new (std::nothrow) v2_decoder_t (
      _options.in_batch_size, _options.maxmsgsize, _options.zero_copy);
    alloc_assert (_decoder);
 
    return zmq::zmtp_engine_t::handshake_v3_x (false);
}
 
int zmq::zmtp_engine_t::routing_id_msg (msg_t *msg_)
{
    const int rc = msg_->init_size (_options.routing_id_size);
    errno_assert (rc == 0);
    if (_options.routing_id_size > 0)
        memcpy (msg_->data (), _options.routing_id, _options.routing_id_size);
    _next_msg = &zmtp_engine_t::pull_msg_from_session;
    return 0;
}
 
int zmq::zmtp_engine_t::process_routing_id_msg (msg_t *msg_)
{
    if (_options.recv_routing_id) {
        msg_->set_flags (msg_t::routing_id);
        const int rc = session ()->push_msg (msg_);
        errno_assert (rc == 0);
    } else {
        int rc = msg_->close ();
        errno_assert (rc == 0);
        rc = msg_->init ();
        errno_assert (rc == 0);
    }
 
    if (_subscription_required) {
        msg_t subscription;
 
        //  Inject the subscription message, so that also
        //  ZMQ 2.x peers receive published messages.
        int rc = subscription.init_size (1);
        errno_assert (rc == 0);
        *static_cast<unsigned char *> (subscription.data ()) = 1;
        rc = session ()->push_msg (&subscription);
        errno_assert (rc == 0);
    }
 
    _process_msg = &zmtp_engine_t::push_msg_to_session;
 
    return 0;
}
 
int zmq::zmtp_engine_t::produce_ping_message (msg_t *msg_)
{
    // 16-bit TTL + \4PING == 7
    const size_t ping_ttl_len = msg_t::ping_cmd_name_size + 2;
    zmq_assert (_mechanism != NULL);
 
    int rc = msg_->init_size (ping_ttl_len);
    errno_assert (rc == 0);
    msg_->set_flags (msg_t::command);
    // Copy in the command message
    memcpy (msg_->data (), "\4PING", msg_t::ping_cmd_name_size);
 
    uint16_t ttl_val = htons (_options.heartbeat_ttl);
    memcpy (static_cast<uint8_t *> (msg_->data ()) + msg_t::ping_cmd_name_size,
            &ttl_val, sizeof (ttl_val));
 
    rc = _mechanism->encode (msg_);
    _next_msg = &zmtp_engine_t::pull_and_encode;
    if (!_has_timeout_timer && _heartbeat_timeout > 0) {
        add_timer (_heartbeat_timeout, heartbeat_timeout_timer_id);
        _has_timeout_timer = true;
    }
    return rc;
}
 
int zmq::zmtp_engine_t::produce_pong_message (msg_t *msg_)
{
    zmq_assert (_mechanism != NULL);
 
    int rc = msg_->move (_pong_msg);
    errno_assert (rc == 0);
 
    rc = _mechanism->encode (msg_);
    _next_msg = &zmtp_engine_t::pull_and_encode;
    return rc;
}
 
int zmq::zmtp_engine_t::process_heartbeat_message (msg_t *msg_)
{
    if (msg_->is_ping ()) {
        // 16-bit TTL + \4PING == 7
        const size_t ping_ttl_len = msg_t::ping_cmd_name_size + 2;
        const size_t ping_max_ctx_len = 16;
        uint16_t remote_heartbeat_ttl;
 
        // Get the remote heartbeat TTL to setup the timer
        memcpy (&remote_heartbeat_ttl,
                static_cast<uint8_t *> (msg_->data ())
                  + msg_t::ping_cmd_name_size,
                ping_ttl_len - msg_t::ping_cmd_name_size);
        remote_heartbeat_ttl = ntohs (remote_heartbeat_ttl);
        // The remote heartbeat is in 10ths of a second
        // so we multiply it by 100 to get the timer interval in ms.
        remote_heartbeat_ttl *= 100;
 
        if (!_has_ttl_timer && remote_heartbeat_ttl > 0) {
            add_timer (remote_heartbeat_ttl, heartbeat_ttl_timer_id);
            _has_ttl_timer = true;
        }
 
        //  As per ZMTP 3.1 the PING command might contain an up to 16 bytes
        //  context which needs to be PONGed back, so build the pong message
        //  here and store it. Truncate it if it's too long.
        //  Given the engine goes straight to out_event, sequential PINGs will
        //  not be a problem.
        const size_t context_len =
          std::min (msg_->size () - ping_ttl_len, ping_max_ctx_len);
        const int rc =
          _pong_msg.init_size (msg_t::ping_cmd_name_size + context_len);
        errno_assert (rc == 0);
        _pong_msg.set_flags (msg_t::command);
        memcpy (_pong_msg.data (), "\4PONG", msg_t::ping_cmd_name_size);
        if (context_len > 0)
            memcpy (static_cast<uint8_t *> (_pong_msg.data ())
                      + msg_t::ping_cmd_name_size,
                    static_cast<uint8_t *> (msg_->data ()) + ping_ttl_len,
                    context_len);
 
        _next_msg = static_cast<int (stream_engine_base_t::*) (msg_t *)> (
          &zmtp_engine_t::produce_pong_message);
        out_event ();
    }
 
    return 0;
}
 
int zmq::zmtp_engine_t::process_command_message (msg_t *msg_)
{
    const uint8_t cmd_name_size =
      *(static_cast<const uint8_t *> (msg_->data ()));
    const size_t ping_name_size = msg_t::ping_cmd_name_size - 1;
    const size_t sub_name_size = msg_t::sub_cmd_name_size - 1;
    const size_t cancel_name_size = msg_t::cancel_cmd_name_size - 1;
    //  Malformed command
    if (unlikely (msg_->size () < cmd_name_size + sizeof (cmd_name_size)))
        return -1;
 
    const uint8_t *const cmd_name =
      static_cast<const uint8_t *> (msg_->data ()) + 1;
    if (cmd_name_size == ping_name_size
        && memcmp (cmd_name, "PING", cmd_name_size) == 0)
        msg_->set_flags (zmq::msg_t::ping);
    if (cmd_name_size == ping_name_size
        && memcmp (cmd_name, "PONG", cmd_name_size) == 0)
        msg_->set_flags (zmq::msg_t::pong);
    if (cmd_name_size == sub_name_size
        && memcmp (cmd_name, "SUBSCRIBE", cmd_name_size) == 0)
        msg_->set_flags (zmq::msg_t::subscribe);
    if (cmd_name_size == cancel_name_size
        && memcmp (cmd_name, "CANCEL", cmd_name_size) == 0)
        msg_->set_flags (zmq::msg_t::cancel);
 
    if (msg_->is_ping () || msg_->is_pong ())
        return process_heartbeat_message (msg_);
 
    return 0;
}