WebSocket.cpp

//
// Created by rahul on 3/21/16.
//

#include "WebSocket.h"

#include <array>
#include <bitset>
#include <boost/asio/buffer.hpp>
#include <boost/asio/read.hpp>
#include <boost/asio/write.hpp>
#include <boost/endian/conversion.hpp>

#include "spdlog/spdlog.h"

using namespace std;
using namespace boost;
using namespace boost::asio;
namespace spd = spdlog;

static const uint8_t continuation_frame = 0x0;
static const uint8_t text_frame = 0x1;
static const uint8_t binary_frame = 0x2;
static const uint8_t connection_close = 0x8;
static const uint8_t ping = 0x9;
static const uint8_t pong = 0xA;

static void read_start_frame(Handler<FrameData> &handler,
                             FrameData &frame_data,
                             const boost::system::error_code &ec,
                             std::size_t);

static void apply_mask(vector<uint8_t> &payload, const uint8_t (&masking_key)[4]) {
  const size_t size = payload.size();
  const size_t size4 = size / 4;
  size_t i = 0;
  uint8_t *p = &payload[0];
  uint32_t *p32 = reinterpret_cast<uint32_t *>(p);
  const uint32_t m = *reinterpret_cast<const uint32_t *>(&masking_key[0]);

#pragma clang loop vectorize(enable) interleave(enable)
  for (i = 0; i < size4; i++) {
    p32[i] = p32[i] ^ m;
  }

  for (i = (size4*4); i < size; i++) {
    p[i] = p[i] ^ masking_key[i % 4];
  }
}

static void read_payload(Handler<FrameData> &handler,
                         FrameData &frame_data,
                         const boost::system::error_code &ec,
                         std::size_t) {
  if (ec) {
    spd::get("console")->error("read_payload Error:{}",ec.message());
    frame_data.websocket.close();
    return;
  }

  auto console = spd::get("console");

  uint8_t opcode = frame_data.frame_start[0] & 0b00001111;
  bool final = (frame_data.frame_start[0] & 0b10000000) == 0b10000000;

  const bool is_text_binary = opcode == text_frame || opcode == binary_frame;

  if (final && frame_data.is_fragmented && is_text_binary) {
    frame_data.is_fragmented = false;
  }

  apply_mask(frame_data.payload, frame_data.masking_key);

  auto pc = frame_data.websocket.acceptor.payload_callback;

  switch (opcode) {
    case continuation_frame:
      pc(frame_data.websocket, frame_data.payload, Opcode::continuation, final);
      break;
    case text_frame:
      pc(frame_data.websocket, frame_data.payload, Opcode::text, final);
      break;
    case binary_frame:
      pc(frame_data.websocket, frame_data.payload, Opcode::binary, final);
      break;
    case connection_close: {
      uint16_t close_code = 0;
      if (frame_data.payload.size() >= 2) {
        close_code = *reinterpret_cast<uint16_t *>(&frame_data.payload[0]);
        endian::big_to_native_inplace(close_code);
        console->trace("Close Code: {}", close_code);
      }
      frame_data.websocket.send_close(true, frame_data.payload);
      return;
    }
    case ping:
      frame_data.websocket.send(frame_data.payload, Opcode::pong, true);
      break;
    case pong: {
      //TODO: Implement Ping with Pong timeout
      break;
    }
    default: {
      console->error("Invalid State");
      frame_data.websocket.close();
      return;
    }
  }

  frame_data.cb = read_start_frame;
  async_read(frame_data.websocket.socket, buffer(frame_data.frame_start, 2), std::move(handler));
}

static void read_length_64(Handler<FrameData> &handler,
                           FrameData &frame_data,
                           const boost::system::error_code &ec,
                           std::size_t) {
  if (ec) {
    spd::get("console")->error("read_length_64 Error:{}",ec.message());
    frame_data.websocket.close();
    return;
  }

  endian::big_to_native_inplace(frame_data.length);
  spd::get("console")->trace("read_length_64: {}", frame_data.length);
  frame_data.payload.resize(frame_data.length);

  frame_data.cb = read_payload;
  async_read(frame_data.websocket.socket, buffer(frame_data.payload), std::move(handler));
}


static void read_length_16(Handler<FrameData> &handler,
                           FrameData &frame_data,
                           const boost::system::error_code &ec,
                           std::size_t) {
  if (ec) {
    spd::get("console")->error("read_length_16 Error:{}",ec.message());
    frame_data.websocket.close();
    return;
  }

  uint16_t &length = *(reinterpret_cast<uint16_t *>(&frame_data.length));
  spd::get("console")->trace("read_length_16: {}", length);
  endian::big_to_native_inplace(length);
  frame_data.payload.resize(length);

  frame_data.cb = read_payload;
  async_read(frame_data.websocket.socket, buffer(frame_data.payload), std::move(handler));
}


static void read_start_frame(Handler<FrameData> &handler,
                             FrameData &frame_data,
                             const boost::system::error_code &ec,
                             std::size_t) {
  if (ec) {
    spd::get("console")->error("read_start_frame Error:{}",ec.message());
    frame_data.websocket.close();
    return;
  }

  auto console = spd::get("console");

  frame_data.length = static_cast<uint64_t>(frame_data.frame_start[1] & 127);
  bool final = (frame_data.frame_start[0] & 0b10000000) == 0b10000000;
  uint8_t opcode = frame_data.frame_start[0] & 0b00001111;
  uint8_t rsv = frame_data.frame_start[0] & 0b01110000;

  const bool is_control_frame = opcode == ping || opcode == pong || opcode == connection_close;
  const bool is_text_binary = opcode == text_frame || opcode == binary_frame;
  bool is_first_fragment = false;

  bitset<8> b1(frame_data.frame_start[0]);
  bitset<8> b2(frame_data.frame_start[1]);
  console->trace("nreceived frame_bits: {} : {}, Final:{}, op: {}", b1, b2, final, static_cast<unsigned int>(opcode));

  if (!final && is_text_binary) {
    frame_data.is_fragmented = true;
    is_first_fragment = true;
  }

  //Sanity Checks
  if (
      rsv != 0  //RSV is always 0
          || (is_control_frame && (frame_data.length > 125 || !final))  //Control frames are always final and <= 125 in length
          || (!frame_data.is_fragmented && opcode == continuation_frame) //Continuation opcode on non fragmented package
          || (!is_first_fragment && is_text_binary && frame_data.is_fragmented && opcode != continuation_frame) //Non Continuation opcode on fragmented package
      ) {
    //5.5.  Control Frames - All control frames MUST have a payload length of 125 bytes or less
    //and MUST NOT be fragmented.
    console->error("Failed Sanity Checks");
    frame_data.websocket.send_close(false, 1002, "Protocol Error");
    return;
  }

  console->trace("frame_start length: {}", frame_data.length);
  if (frame_data.length == 126) {
    frame_data.length = 0;

    frame_data.cb = read_length_16;
    async_read(
        frame_data.websocket.socket,
        std::array<mutable_buffer, 2>{
            buffer(&frame_data.length, 2),
            buffer(&frame_data.masking_key, 4)
        },
        std::move(handler));
  } else if (frame_data.length == 127) {
    frame_data.cb = read_length_64;
    async_read(
        frame_data.websocket.socket,
        std::array<mutable_buffer, 2>{
            buffer(&frame_data.length, 8),
            buffer(&frame_data.masking_key, 4)
        },
        std::move(handler));
  } else {
    frame_data.cb = read_payload;

    if (frame_data.length > 0) {
      frame_data.payload.resize(frame_data.length);
      async_read(
          frame_data.websocket.socket,
          std::array<mutable_buffer, 2> {
              buffer(&frame_data.masking_key, 4),
              buffer(frame_data.payload)
          },
          std::move(handler));
    } else {
      // No Payload data
      async_read(frame_data.websocket.socket, buffer(&frame_data.masking_key, 4),
                 std::move(handler));
    }
  }
}

static void header_write_handler(Handler<HeaderData> &,
                                 HeaderData &header_data,
                                 const boost::system::error_code &ec,
                                 std::size_t bytes_transferred) {
  if (ec) {
    spd::get("console")->error("HeaderWriteHandler::operator Error: {}",ec.message());
    header_data.websocket.close();
    return;
  }

  spd::get("console")->trace("Header written: {}", bytes_transferred);
}

static void header_read_handler(Handler<HeaderData> &handler,
                                HeaderData &header_data,
                                const boost::system::error_code &ec,
                                std::size_t bytes_transferred) {
  if (ec) {
    spd::get("console")->error("HeaderReceiveHandler::operator() Error: {}",ec.message());
    header_data.websocket.close();
    return;
  }

  auto console = spd::get("console");

  if (bytes_transferred > 0) {
    HTTPParser &par = header_data.par;

    try {
      par.received_data(header_data.cbuffer, bytes_transferred);
    } catch (const ParserException &pex) {
      console->error("ParserException& pex Error: {}", pex.what());
      header_data.websocket.close();
      return;
    } catch (const InvalidData &idex) {
      console->error("InvalidData& idex Error: {}", idex.what());
      header_data.websocket.close();
      return;
    }

    if (par.message_complete) {
      Handler<FrameData> frame_handler(make_unique<FrameData>(header_data.websocket, read_start_frame));
      async_read(header_data.websocket.socket,
                 buffer(frame_handler.state->frame_start, 2),
                 std::move(frame_handler));

      derive_sec_websocket_accept(par.sec_websocket_key.c_str(),
                                  par.sec_websocket_key.length(),
                                  header_data.sec_websocket_accept);

      std::array<const_buffer, 3> response_buffer = {
          buffer(header_data.response_header_start, 97),
          buffer(header_data.sec_websocket_accept, 28),
          buffer(header_data.response_header_end, 4),
      };
      header_data.cb = header_write_handler;
      async_write(header_data.websocket.socket, response_buffer,
                  std::move(handler));
      return;
    }
  }

  header_data.websocket.socket.async_receive(buffer(header_data.cbuffer, 64),
                                             std::move(handler));
}

void WebSocket::start_receive() {
  Handler<HeaderData> handler(make_unique<HeaderData>(*this, header_read_handler));
  socket.async_receive(buffer(handler.state->cbuffer, 64), std::move(handler));
}

static void sent_handler(Handler<BufferedSendData> &,
                         BufferedSendData &sending_data,
                         const boost::system::error_code &ec,
                         std::size_t) {
  if (ec) {
    spd::get("console")->error("SendPacketHandler::operator() Error:{}",ec.message());
    sending_data.websocket.close();
    return;
  }
  sending_data.websocket.is_sending = false;

  size_t package_size = sending_data.payload_list.size();
  spd::get("console")->trace("Sent: {}", package_size);

  if(sending_data.has_close) {
    sending_data.websocket.close();
    return;
  }

  sending_data.websocket.send_buffered();
}

void WebSocket::send(FrameData::Payload &payload, Opcode opcode, bool is_final) {
  auto sending_data = make_unique<SendingData>(payload);
  if (is_final) {
    sending_data->frame_start[0] = 0b10000000;
  } else {
    sending_data->frame_start[0] = 0;
  }
  sending_data->frame_start[0] = sending_data->frame_start[0] | static_cast<uint8_t>(opcode);

  size_t package_size = sending_data->payload.size();

  bitset<8> b1(sending_data->frame_start[0]);
  bitset<8> b2(sending_data->frame_start[1]);
  spd::get("console")->trace("sending package size:{} ,send frame_bits:{} : {}", package_size, b1, b2);

  if (package_size < 126) {
    sending_data->frame_start[1] = static_cast<uint8_t>(package_size);

    buffer_list.push_back(buffer(sending_data->frame_start, 2));
    buffer_list.push_back(buffer(sending_data->payload));

  } else if (package_size >= 126 && package_size <= 65535) {
    sending_data->frame_start[1] = 126;
    uint16_t &ex_length = *(reinterpret_cast<uint16_t *>(&sending_data->extended_length));
    ex_length = static_cast<uint16_t>(package_size);
    endian::native_to_big_inplace(ex_length);

    buffer_list.push_back(buffer(sending_data->frame_start, 2));
    buffer_list.push_back(buffer(&ex_length, 2));
    buffer_list.push_back(buffer(sending_data->payload));

  } else if (package_size > 65535) {
    sending_data->frame_start[1] = 127;
    sending_data->extended_length = endian::native_to_big(package_size);

    buffer_list.push_back(buffer(sending_data->frame_start, 2));
    buffer_list.push_back(buffer(&sending_data->extended_length, 8));
    buffer_list.push_back(buffer(sending_data->payload));
  }

  payload_list.emplace_back(std::move(sending_data));

  send_buffered();
}

void Acceptor::start_accept() {
  socket_list.emplace_front(io_service, *this);
  auto socket_it = socket_list.begin();
  socket_it->it = socket_it;

  spd::get("console")->error("Total connections: {}", socket_list.size());

  acceptor.async_accept(
      socket_it->socket,
      [this, socket_it](const boost::system::error_code &error) {
        if (error) {
          this->socket_list.erase(socket_it);
          return;
        }

        socket_it->start_receive();
        this->start_accept();
      });
}

const char *HeaderData::response_header_start = "HTTP/1.1 101 Switching Protocols\r\n"
    "Upgrade: websocket\r\n"
    "Connection: Upgrade\r\n"
    "Sec-WebSocket-Accept: ";
const char *HeaderData::response_header_end = "\r\n\r\n";

void WebSocket::send_close(bool /*is_close_response*/, FrameData::Payload &payload) {
  //TODO: see if we need is_close_response
  spd::get("console")->trace("send_close");

  size_t payload_size = payload.size();
  if (payload_size > 125) {
    //TODO: Throw Exception
    payload.resize(123);
  }

  auto sending_data = make_unique<SendingData>(payload);
  SendingData &send_data = *sending_data;

  send_data.frame_start[0] = 0b10000000 | connection_close;
  send_data.frame_start[1] = static_cast<uint8_t>(payload_size);

  if (payload_size == 0) {
    buffer_list.push_back(buffer(send_data.frame_start, 2));
  } else {
    buffer_list.push_back(buffer(send_data.frame_start, 2));
    buffer_list.push_back(buffer(send_data.payload));
  }

  payload_list.emplace_back(std::move(sending_data));
  close_requested = true;
  send_buffered();
}

void WebSocket::send_close(bool is_close_response, std::uint16_t code, std::string reason) {
  FrameData::Payload pl;
  size_t reason_size = reason.size();
  pl.reserve(2 + reason_size);
  endian::native_to_big_inplace(code);
  uint8_t *code8 = reinterpret_cast<uint8_t *>(&code);
  std::copy(code8, code8 + 2, std::back_inserter(pl));
  std::copy(reason.begin(), reason.end(), std::back_inserter(pl));
  send_close(is_close_response, pl);
}

void WebSocket::send_buffered() {
  if (is_sending || payload_list.empty()) {
    return;
  }
  is_sending = true;

  Handler<BufferedSendData> handler(make_unique<BufferedSendData>(payload_list, buffer_list, *this, sent_handler, close_requested));
  BufferRef<BufferedSendData::BufferList> cref(handler.state->buffer_list);
  async_write(socket, cref, std::move(handler));
}

void WebSocket::close()  {
  acceptor.socket_list.erase(it);
}