text: /**
   *
   *  WebSocketConnectionImpl.cc
   *  An Tao
   *
   *  Copyright 2018, An Tao.  All rights reserved.
   *  https://github.com/an-tao/drogon
   *  Use of this source code is governed by a MIT license
   *  that can be found in the License file.
   *
   *  Drogon
   *
   */
  
  #include "WebSocketConnectionImpl.h"
  #include "HttpAppFrameworkImpl.h"
  #include <thread>
  
  using namespace drogon;
  WebSocketConnectionImpl::WebSocketConnectionImpl(
      const trantor::TcpConnectionPtr &conn,
      HttpAppFrameworkImpl *app,
      bool isServer)
      : tcpConnectionPtr_(conn),
        localAddr_(conn->localAddr()),
        peerAddr_(conn->peerAddr()),
        isServer_(isServer),
        parser_(app)
  {
      app_ = app;
  }
  
  void WebSocketConnectionImpl::send(const char *msg,
                                     uint64_t len,
                                     const WebSocketMessageType type)
  {
      unsigned char opcode;
      if (type == WebSocketMessageType::Text)
          opcode = 1;
      else if (type == WebSocketMessageType::Binary)
          opcode = 2;
      else if (type == WebSocketMessageType::Close)
      {
          assert(len <= 125);
          opcode = 8;
      }
      else if (type == WebSocketMessageType::Ping)
      {
          assert(len <= 125);
          opcode = 9;
      }
      else if (type == WebSocketMessageType::Pong)
      {
          assert(len <= 125);
          opcode = 10;
      }
      else
      {
          opcode = 0;
          assert(0);
      }
      sendWsData(msg, len, opcode);
  }
  
  void WebSocketConnectionImpl::sendWsData(const char *msg,
                                           uint64_t len,
                                           unsigned char opcode)
  {
      LOG_TRACE << "send " << len << " bytes";
  
      // Format the frame
      std::string bytesFormatted;
      bytesFormatted.resize(len + 10);
      bytesFormatted[0] = char(0x80 | (opcode & 0x0f));
  
      int indexStartRawData = -1;
  
      if (len <= 125)
      {
          bytesFormatted[1] = static_cast<char>(len);
          indexStartRawData = 2;
      }
      else if (len <= 65535)
      {
          bytesFormatted[1] = 126;
          bytesFormatted[2] = ((len >> 8) & 255);
          bytesFormatted[3] = ((len)&255);
          LOG_TRACE << "bytes[2]=" << (size_t)bytesFormatted[2];
          LOG_TRACE << "bytes[3]=" << (size_t)bytesFormatted[3];
          indexStartRawData = 4;
      }
      else
      {
          bytesFormatted[1] = 127;
          bytesFormatted[2] = ((len >> 56) & 255);
          bytesFormatted[3] = ((len >> 48) & 255);
          bytesFormatted[4] = ((len >> 40) & 255);
          bytesFormatted[5] = ((len >> 32) & 255);
          bytesFormatted[6] = ((len >> 24) & 255);
          bytesFormatted[7] = ((len >> 16) & 255);
          bytesFormatted[8] = ((len >> 8) & 255);
          bytesFormatted[9] = ((len)&255);
  
          indexStartRawData = 10;
      }
      if (!isServer_)
      {
          // Add masking key;
          static std::once_flag once;
          std::call_once(once, []() {
              std::srand(static_cast<unsigned int>(time(nullptr)));
          });
          int random = std::rand();
  
          bytesFormatted[1] = (bytesFormatted[1] | 0x80);
          bytesFormatted.resize(indexStartRawData + 4 + len);
          *((int *)&bytesFormatted[indexStartRawData]) = random;
          for (size_t i = 0; i < len; ++i)
          {
              bytesFormatted[indexStartRawData + 4 + i] =
                  (msg[i] ^ bytesFormatted[indexStartRawData + (i % 4)]);
          }
      }
      else
      {
          bytesFormatted.resize(indexStartRawData);
          bytesFormatted.append(msg, len);
      }
      tcpConnectionPtr_->send(std::move(bytesFormatted));
  }
  void WebSocketConnectionImpl::send(const std::string &msg,
                                     const WebSocketMessageType type)
  {
      send(msg.data(), msg.length(), type);
  }
  const trantor::InetAddress &WebSocketConnectionImpl::localAddr() const
  {
      return localAddr_;
  }
  const trantor::InetAddress &WebSocketConnectionImpl::peerAddr() const
  {
      return peerAddr_;
  }
  
  bool WebSocketConnectionImpl::connected() const
  {
      return tcpConnectionPtr_->connected();
  }
  bool WebSocketConnectionImpl::disconnected() const
  {
      return tcpConnectionPtr_->disconnected();
  }
  void WebSocketConnectionImpl::WebSocketConnectionImpl::shutdown(
      const CloseCode code,
      const std::string &reason)
  {
      tcpConnectionPtr_->getLoop()->invalidateTimer(pingTimerId_);
      std::string message;
      message.resize(reason.length() + 2);
      auto c = htons(static_cast<unsigned short>(code));
      memcpy(&message[0], &c, 2);
      if (!reason.empty())
          memcpy(&message[2], reason.data(), reason.length());
      send(message, WebSocketMessageType::Close);
      tcpConnectionPtr_->shutdown();
  }
  void WebSocketConnectionImpl::WebSocketConnectionImpl::forceClose()
  {
      tcpConnectionPtr_->forceClose();
  }
  
  void WebSocketConnectionImpl::setPingMessage(
      const std::string &message,
      const std::chrono::duration<long double> &interval)
  {
      std::weak_ptr<WebSocketConnectionImpl> weakPtr = shared_from_this();
      pingTimerId_ = tcpConnectionPtr_->getLoop()->runEvery(
          interval.count(), [weakPtr, message]() {
              auto thisPtr = weakPtr.lock();
              if (thisPtr)
              {
                  thisPtr->send(message, WebSocketMessageType::Ping);
              }
          });
  }
  
  bool WebSocketMessageParser::parse(trantor::MsgBuffer *buffer)
  {
      // According to the rfc6455
      gotAll_ = false;
      if (buffer->readableBytes() >= 2)
      {
          unsigned char opcode = (*buffer)[0] & 0x0f;
          bool isControlFrame = false;
          switch (opcode)
          {
              case 0:
                  // continuation frame
                  break;
              case 1:
                  type_ = WebSocketMessageType::Text;
                  break;
              case 2:
                  type_ = WebSocketMessageType::Binary;
                  break;
              case 8:
                  type_ = WebSocketMessageType::Close;
                  isControlFrame = true;
                  break;
              case 9:
                  type_ = WebSocketMessageType::Ping;
                  isControlFrame = true;
                  break;
              case 10:
                  type_ = WebSocketMessageType::Pong;
                  isControlFrame = true;
                  break;
              default:
                  LOG_ERROR << "Unknown frame type";
                  return false;
                  break;
          }
  
          bool isFin = (((*buffer)[0] & 0x80) == 0x80);
          if (!isFin && isControlFrame)
          {
              // rfc6455-5.5
              LOG_ERROR << "Bad frame: all control frames MUST NOT be fragmented";
              return false;
          }
          auto secondByte = (*buffer)[1];
          size_t length = secondByte & 127;
          int isMasked = (secondByte & 0x80);
          if (isMasked != 0)
          {
              LOG_TRACE << "data encoded!";
          }
          else
              LOG_TRACE << "plain data";
          size_t indexFirstMask = 2;
  
          if (length == 126)
          {
              indexFirstMask = 4;
          }
          else if (length == 127)
          {
              indexFirstMask = 10;
          }
          if (indexFirstMask > 2 && buffer->readableBytes() >= indexFirstMask)
          {
              if (isControlFrame)
              {
                  // rfc6455-5.5
                  LOG_ERROR << "Bad frame: all control frames MUST have a "
                               "payload length "
                               "of 125 bytes or less";
                  return false;
              }
              if (indexFirstMask == 4)
              {
                  length = (unsigned char)(*buffer)[2];
                  length = (length << 8) + (unsigned char)(*buffer)[3];
              }
              else if (indexFirstMask == 10)
              {
                  length = (unsigned char)(*buffer)[2];
                  length = (length << 8) + (unsigned char)(*buffer)[3];
                  length = (length << 8) + (unsigned char)(*buffer)[4];
                  length = (length << 8) + (unsigned char)(*buffer)[5];
                  length = (length << 8) + (unsigned char)(*buffer)[6];
                  length = (length << 8) + (unsigned char)(*buffer)[7];
                  length = (length << 8) + (unsigned char)(*buffer)[8];
                  length = (length << 8) + (unsigned char)(*buffer)[9];
              }
              else
              {
                  LOG_ERROR << "Websock parsing failed!";
                  return false;
              }
          }
          if (isMasked != 0)
          {
              // The message is sent by the client, check the length
              if (length > app_->getClientMaxWebSocketMessageSize())
              {
                  LOG_ERROR << "The size of the WebSocket message is too large!";
                  buffer->retrieveAll();
                  return false;
              }
              if (buffer->readableBytes() >= (indexFirstMask + 4 + length))
              {
                  auto masks = buffer->peek() + indexFirstMask;
                  auto indexFirstDataByte = indexFirstMask + 4;
                  auto rawData = buffer->peek() + indexFirstDataByte;
                  auto oldLen = message_.length();
                  message_.resize(oldLen + length);
                  for (size_t i = 0; i < length; ++i)
                  {
                      message_[oldLen + i] = (rawData[i] ^ masks[i % 4]);
                  }
                  if (isFin)
                      gotAll_ = true;
                  buffer->retrieve(indexFirstMask + 4 + length);
                  return true;
              }
          }
          else
          {
              if (buffer->readableBytes() >= (indexFirstMask + length))
              {
                  auto rawData = buffer->peek() + indexFirstMask;
                  message_.append(rawData, length);
                  if (isFin)
                      gotAll_ = true;
                  buffer->retrieve(indexFirstMask + length);
                  return true;
              }
          }
      }
      return true;
  }
  
  void WebSocketConnectionImpl::onNewMessage(
      const trantor::TcpConnectionPtr &connPtr,
      trantor::MsgBuffer *buffer)
  {
      while (buffer->readableBytes() > 0)
      {
          auto success = parser_.parse(buffer);
          if (success)
          {
              std::string message;
              WebSocketMessageType type;
              if (parser_.gotAll(message, type))
              {
                  if (type == WebSocketMessageType::Ping)
                  {
                      // ping
                      send(message, WebSocketMessageType::Pong);
                  }
                  else if (type == WebSocketMessageType::Close)
                  {
                      // close
                      connPtr->shutdown();
                  }
                  else if (type == WebSocketMessageType::Unknown)
                  {
                      return;
                  }
                  messageCallback_(std::move(message), shared_from_this(), type);
              }
              else
              {
                  return;
              }
          }
          else
          {
              // Websock error!
              connPtr->shutdown();
              return;
          }
      }
      return;
  }