text: /*----------------------------------------------------------------------------*/
  /* Copyright (c) 2018-2019 FIRST. All Rights Reserved.                        */
  /* Open Source Software - may be modified and shared by FRC teams. The code   */
  /* must be accompanied by the FIRST BSD license file in the root directory of */
  /* the project.                                                               */
  /*----------------------------------------------------------------------------*/
  
  #include "wpi/WebSocket.h"
  
  #include <random>
  
  #include "wpi/Base64.h"
  #include "wpi/HttpParser.h"
  #include "wpi/SmallString.h"
  #include "wpi/SmallVector.h"
  #include "wpi/raw_uv_ostream.h"
  #include "wpi/sha1.h"
  #include "wpi/uv/Stream.h"
  
  using namespace wpi;
  
  namespace {
  class WebSocketWriteReq : public uv::WriteReq {
   public:
    explicit WebSocketWriteReq(
        std::function<void(MutableArrayRef<uv::Buffer>, uv::Error)> callback) {
      finish.connect([=](uv::Error err) {
        MutableArrayRef<uv::Buffer> bufs{m_bufs};
        for (auto&& buf : bufs.slice(0, m_startUser)) buf.Deallocate();
        callback(bufs.slice(m_startUser), err);
      });
    }
  
    SmallVector<uv::Buffer, 4> m_bufs;
    size_t m_startUser;
  };
  }  // namespace
  
  class WebSocket::ClientHandshakeData {
   public:
    ClientHandshakeData() {
      // key is a random nonce
      static std::random_device rd;
      static std::default_random_engine gen{rd()};
      std::uniform_int_distribution<unsigned int> dist(0, 255);
      char nonce[16];  // the nonce sent to the server
      for (char& v : nonce) v = static_cast<char>(dist(gen));
      raw_svector_ostream os(key);
      Base64Encode(os, StringRef{nonce, 16});
    }
    ~ClientHandshakeData() {
      if (auto t = timer.lock()) {
        t->Stop();
        t->Close();
      }
    }
  
    SmallString<64> key;                       // the key sent to the server
    SmallVector<std::string, 2> protocols;     // valid protocols
    HttpParser parser{HttpParser::kResponse};  // server response parser
    bool hasUpgrade = false;
    bool hasConnection = false;
    bool hasAccept = false;
    bool hasProtocol = false;
  
    std::weak_ptr<uv::Timer> timer;
  };
  
  static StringRef AcceptHash(StringRef key, SmallVectorImpl<char>& buf) {
    SHA1 hash;
    hash.Update(key);
    hash.Update("258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
    SmallString<64> hashBuf;
    return Base64Encode(hash.RawFinal(hashBuf), buf);
  }
  
  WebSocket::WebSocket(uv::Stream& stream, bool server, const private_init&)
      : m_stream{stream}, m_server{server} {
    // Connect closed and error signals to ourselves
    m_stream.closed.connect([this]() { SetClosed(1006, "handle closed"); });
    m_stream.error.connect([this](uv::Error err) {
      Terminate(1006, "stream error: " + Twine(err.name()));
    });
  
    // Start reading
    m_stream.StopRead();  // we may have been reading
    m_stream.StartRead();
    m_stream.data.connect(
        [this](uv::Buffer& buf, size_t size) { HandleIncoming(buf, size); });
    m_stream.end.connect(
        [this]() { Terminate(1006, "remote end closed connection"); });
  }
  
  WebSocket::~WebSocket() {}
  
  std::shared_ptr<WebSocket> WebSocket::CreateClient(
      uv::Stream& stream, const Twine& uri, const Twine& host,
      ArrayRef<StringRef> protocols, const ClientOptions& options) {
    auto ws = std::make_shared<WebSocket>(stream, false, private_init{});
    stream.SetData(ws);
    ws->StartClient(uri, host, protocols, options);
    return ws;
  }
  
  std::shared_ptr<WebSocket> WebSocket::CreateServer(uv::Stream& stream,
                                                     StringRef key,
                                                     StringRef version,
                                                     StringRef protocol) {
    auto ws = std::make_shared<WebSocket>(stream, true, private_init{});
    stream.SetData(ws);
    ws->StartServer(key, version, protocol);
    return ws;
  }
  
  void WebSocket::Close(uint16_t code, const Twine& reason) {
    SendClose(code, reason);
    if (m_state != FAILED && m_state != CLOSED) m_state = CLOSING;
  }
  
  void WebSocket::Fail(uint16_t code, const Twine& reason) {
    if (m_state == FAILED || m_state == CLOSED) return;
    SendClose(code, reason);
    SetClosed(code, reason, true);
    Shutdown();
  }
  
  void WebSocket::Terminate(uint16_t code, const Twine& reason) {
    if (m_state == FAILED || m_state == CLOSED) return;
    SetClosed(code, reason);
    Shutdown();
  }
  
  void WebSocket::StartClient(const Twine& uri, const Twine& host,
                              ArrayRef<StringRef> protocols,
                              const ClientOptions& options) {
    // Create client handshake data
    m_clientHandshake = std::make_unique<ClientHandshakeData>();
  
    // Build client request
    SmallVector<uv::Buffer, 4> bufs;
    raw_uv_ostream os{bufs, 4096};
  
    os << "GET " << uri << " HTTP/1.1\r\n";
    os << "Host: " << host << "\r\n";
    os << "Upgrade: websocket\r\n";
    os << "Connection: Upgrade\r\n";
    os << "Sec-WebSocket-Key: " << m_clientHandshake->key << "\r\n";
    os << "Sec-WebSocket-Version: 13\r\n";
  
    // protocols (if provided)
    if (!protocols.empty()) {
      os << "Sec-WebSocket-Protocol: ";
      bool first = true;
      for (auto protocol : protocols) {
        if (!first)
          os << ", ";
        else
          first = false;
        os << protocol;
        // also save for later checking against server response
        m_clientHandshake->protocols.emplace_back(protocol);
      }
      os << "\r\n";
    }
  
    // other headers
    for (auto&& header : options.extraHeaders)
      os << header.first << ": " << header.second << "\r\n";
  
    // finish headers
    os << "\r\n";
  
    // Send client request
    m_stream.Write(bufs, [](auto bufs, uv::Error) {
      for (auto& buf : bufs) buf.Deallocate();
    });
  
    // Set up client response handling
    m_clientHandshake->parser.status.connect([this](StringRef status) {
      unsigned int code = m_clientHandshake->parser.GetStatusCode();
      if (code != 101) Terminate(code, status);
    });
    m_clientHandshake->parser.header.connect(
        [this](StringRef name, StringRef value) {
          value = value.trim();
          if (name.equals_lower("upgrade")) {
            if (!value.equals_lower("websocket"))
              return Terminate(1002, "invalid upgrade response value");
            m_clientHandshake->hasUpgrade = true;
          } else if (name.equals_lower("connection")) {
            if (!value.equals_lower("upgrade"))
              return Terminate(1002, "invalid connection response value");
            m_clientHandshake->hasConnection = true;
          } else if (name.equals_lower("sec-websocket-accept")) {
            // Check against expected response
            SmallString<64> acceptBuf;
            if (!value.equals(AcceptHash(m_clientHandshake->key, acceptBuf)))
              return Terminate(1002, "invalid accept key");
            m_clientHandshake->hasAccept = true;
          } else if (name.equals_lower("sec-websocket-extensions")) {
            // No extensions are supported
            if (!value.empty()) return Terminate(1010, "unsupported extension");
          } else if (name.equals_lower("sec-websocket-protocol")) {
            // Make sure it was one of the provided protocols
            bool match = false;
            for (auto&& protocol : m_clientHandshake->protocols) {
              if (value.equals_lower(protocol)) {
                match = true;
                break;
              }
            }
            if (!match) return Terminate(1003, "unsupported protocol");
            m_clientHandshake->hasProtocol = true;
            m_protocol = value;
          }
        });
    m_clientHandshake->parser.headersComplete.connect([this](bool) {
      if (!m_clientHandshake->hasUpgrade || !m_clientHandshake->hasConnection ||
          !m_clientHandshake->hasAccept ||
          (!m_clientHandshake->hasProtocol &&
           !m_clientHandshake->protocols.empty())) {
        return Terminate(1002, "invalid response");
      }
      if (m_state == CONNECTING) {
        m_state = OPEN;
        open(m_protocol);
      }
    });
  
    // Start handshake timer if a timeout was specified
    if (options.handshakeTimeout != (uv::Timer::Time::max)()) {
      auto timer = uv::Timer::Create(m_stream.GetLoopRef());
      timer->timeout.connect(
          [this]() { Terminate(1006, "connection timed out"); });
      timer->Start(options.handshakeTimeout);
      m_clientHandshake->timer = timer;
    }
  }
  
  void WebSocket::StartServer(StringRef key, StringRef version,
                              StringRef protocol) {
    m_protocol = protocol;
  
    // Build server response
    SmallVector<uv::Buffer, 4> bufs;
    raw_uv_ostream os{bufs, 4096};
  
    // Handle unsupported version
    if (version != "13") {
      os << "HTTP/1.1 426 Upgrade Required\r\n";
      os << "Upgrade: WebSocket\r\n";
      os << "Sec-WebSocket-Version: 13\r\n\r\n";
      m_stream.Write(bufs, [this](auto bufs, uv::Error) {
        for (auto& buf : bufs) buf.Deallocate();
        // XXX: Should we support sending a new handshake on the same connection?
        // XXX: "this->" is required by GCC 5.5 (bug)
        this->Terminate(1003, "unsupported protocol version");
      });
      return;
    }
  
    os << "HTTP/1.1 101 Switching Protocols\r\n";
    os << "Upgrade: websocket\r\n";
    os << "Connection: Upgrade\r\n";
  
    // accept hash
    SmallString<64> acceptBuf;
    os << "Sec-WebSocket-Accept: " << AcceptHash(key, acceptBuf) << "\r\n";
  
    if (!protocol.empty()) os << "Sec-WebSocket-Protocol: " << protocol << "\r\n";
  
    // end headers
    os << "\r\n";
  
    // Send server response
    m_stream.Write(bufs, [this](auto bufs, uv::Error) {
      for (auto& buf : bufs) buf.Deallocate();
      if (m_state == CONNECTING) {
        m_state = OPEN;
        open(m_protocol);
      }
    });
  }
  
  void WebSocket::SendClose(uint16_t code, const Twine& reason) {
    SmallVector<uv::Buffer, 4> bufs;
    if (code != 1005) {
      raw_uv_ostream os{bufs, 4096};
      const uint8_t codeMsb[] = {static_cast<uint8_t>((code >> 8) & 0xff),
                                 static_cast<uint8_t>(code & 0xff)};
      os << ArrayRef<uint8_t>(codeMsb);
      reason.print(os);
    }
    Send(kFlagFin | kOpClose, bufs, [](auto bufs, uv::Error) {
      for (auto&& buf : bufs) buf.Deallocate();
    });
  }
  
  void WebSocket::SetClosed(uint16_t code, const Twine& reason, bool failed) {
    if (m_state == FAILED || m_state == CLOSED) return;
    m_state = failed ? FAILED : CLOSED;
    SmallString<64> reasonBuf;
    closed(code, reason.toStringRef(reasonBuf));
  }
  
  void WebSocket::Shutdown() {
    m_stream.Shutdown([this] { m_stream.Close(); });
  }
  
  void WebSocket::HandleIncoming(uv::Buffer& buf, size_t size) {
    // ignore incoming data if we're failed or closed
    if (m_state == FAILED || m_state == CLOSED) return;
  
    StringRef data{buf.base, size};
  
    // Handle connecting state (mainly on client)
    if (m_state == CONNECTING) {
      if (m_clientHandshake) {
        data = m_clientHandshake->parser.Execute(data);
        // check for parser failure
        if (m_clientHandshake->parser.HasError())
          return Terminate(1003, "invalid response");
        if (m_state != OPEN) return;  // not done with handshake yet
  
        // we're done with the handshake, so release its memory
        m_clientHandshake.reset();
  
        // fall through to process additional data after handshake
      } else {
        return Terminate(1003, "got data on server before response");
      }
    }
  
    // Message processing
    while (!data.empty()) {
      if (m_frameSize == UINT64_MAX) {
        // Need at least two bytes to determine header length
        if (m_header.size() < 2u) {
          size_t toCopy = (std::min)(2u - m_header.size(), data.size());
          m_header.append(data.bytes_begin(), data.bytes_begin() + toCopy);
          data = data.drop_front(toCopy);
          if (m_header.size() < 2u) return;  // need more data
  
          // Validate RSV bits are zero
          if ((m_header[0] & 0x70) != 0) return Fail(1002, "nonzero RSV");
        }
  
        // Once we have first two bytes, we can calculate the header size
        if (m_headerSize == 0) {
          m_headerSize = 2;
          uint8_t len = m_header[1] & kLenMask;
          if (len == 126)
            m_headerSize += 2;
          else if (len == 127)
            m_headerSize += 8;
          bool masking = (m_header[1] & kFlagMasking) != 0;
          if (masking) m_headerSize += 4;  // masking key
          // On server side, incoming messages MUST be masked
          // On client side, incoming messages MUST NOT be masked
          if (m_server && !masking) return Fail(1002, "client data not masked");
          if (!m_server && masking) return Fail(1002, "server data masked");
        }
  
        // Need to complete header to calculate message size
        if (m_header.size() < m_headerSize) {
          size_t toCopy = (std::min)(m_headerSize - m_header.size(), data.size());
          m_header.append(data.bytes_begin(), data.bytes_begin() + toCopy);
          data = data.drop_front(toCopy);
          if (m_header.size() < m_headerSize) return;  // need more data
        }
  
        if (m_header.size() >= m_headerSize) {
          // get payload length
          uint8_t len = m_header[1] & kLenMask;
          if (len == 126)
            m_frameSize = (static_cast<uint16_t>(m_header[2]) << 8) |
                          static_cast<uint16_t>(m_header[3]);
          else if (len == 127)
            m_frameSize = (static_cast<uint64_t>(m_header[2]) << 56) |
                          (static_cast<uint64_t>(m_header[3]) << 48) |
                          (static_cast<uint64_t>(m_header[4]) << 40) |
                          (static_cast<uint64_t>(m_header[5]) << 32) |
                          (static_cast<uint64_t>(m_header[6]) << 24) |
                          (static_cast<uint64_t>(m_header[7]) << 16) |
                          (static_cast<uint64_t>(m_header[8]) << 8) |
                          static_cast<uint64_t>(m_header[9]);
          else
            m_frameSize = len;
  
          // limit maximum size
          if ((m_payload.size() + m_frameSize) > m_maxMessageSize)
            return Fail(1009, "message too large");
        }
      }
  
      if (m_frameSize != UINT64_MAX) {
        size_t need = m_frameStart + m_frameSize - m_payload.size();
        size_t toCopy = (std::min)(need, data.size());
        m_payload.append(data.bytes_begin(), data.bytes_begin() + toCopy);
        data = data.drop_front(toCopy);
        need -= toCopy;
        if (need == 0) {
          // We have a complete frame
          // If the message had masking, unmask it
          if ((m_header[1] & kFlagMasking) != 0) {
            uint8_t key[4] = {
                m_header[m_headerSize - 4], m_header[m_headerSize - 3],
                m_header[m_headerSize - 2], m_header[m_headerSize - 1]};
            int n = 0;
            for (uint8_t& ch :
                 MutableArrayRef<uint8_t>{m_payload}.slice(m_frameStart)) {
              ch ^= key[n++];
              if (n >= 4) n = 0;
            }
          }
  
          // Handle message
          bool fin = (m_header[0] & kFlagFin) != 0;
          uint8_t opcode = m_header[0] & kOpMask;
          switch (opcode) {
            case kOpCont:
              switch (m_fragmentOpcode) {
                case kOpText:
                  if (!m_combineFragments || fin)
                    text(StringRef{reinterpret_cast<char*>(m_payload.data()),
                                   m_payload.size()},
                         fin);
                  break;
                case kOpBinary:
                  if (!m_combineFragments || fin) binary(m_payload, fin);
                  break;
                default:
                  // no preceding message?
                  return Fail(1002, "invalid continuation message");
              }
              if (fin) m_fragmentOpcode = 0;
              break;
            case kOpText:
              if (m_fragmentOpcode != 0) return Fail(1002, "incomplete fragment");
              if (!m_combineFragments || fin)
                text(StringRef{reinterpret_cast<char*>(m_payload.data()),
                               m_payload.size()},
                     fin);
              if (!fin) m_fragmentOpcode = opcode;
              break;
            case kOpBinary:
              if (m_fragmentOpcode != 0) return Fail(1002, "incomplete fragment");
              if (!m_combineFragments || fin) binary(m_payload, fin);
              if (!fin) m_fragmentOpcode = opcode;
              break;
            case kOpClose: {
              uint16_t code;
              StringRef reason;
              if (!fin) {
                code = 1002;
                reason = "cannot fragment control frames";
              } else if (m_payload.size() < 2) {
                code = 1005;
              } else {
                code = (static_cast<uint16_t>(m_payload[0]) << 8) |
                       static_cast<uint16_t>(m_payload[1]);
                reason = StringRef{reinterpret_cast<char*>(m_payload.data()),
                                   m_payload.size()}
                             .drop_front(2);
              }
              // Echo the close if we didn't previously send it
              if (m_state != CLOSING) SendClose(code, reason);
              SetClosed(code, reason);
              // If we're the server, shutdown the connection.
              if (m_server) Shutdown();
              break;
            }
            case kOpPing:
              if (!fin) return Fail(1002, "cannot fragment control frames");
              ping(m_payload);
              break;
            case kOpPong:
              if (!fin) return Fail(1002, "cannot fragment control frames");
              pong(m_payload);
              break;
            default:
              return Fail(1002, "invalid message opcode");
          }
  
          // Prepare for next message
          m_header.clear();
          m_headerSize = 0;
          if (!m_combineFragments || fin) m_payload.clear();
          m_frameStart = m_payload.size();
          m_frameSize = UINT64_MAX;
        }
      }
    }
  }
  
  void WebSocket::Send(
      uint8_t opcode, ArrayRef<uv::Buffer> data,
      std::function<void(MutableArrayRef<uv::Buffer>, uv::Error)> callback) {
    // If we're not open, emit an error and don't send the data
    if (m_state != OPEN) {
      int err;
      if (m_state == CONNECTING)
        err = UV_EAGAIN;
      else
        err = UV_ESHUTDOWN;
      SmallVector<uv::Buffer, 4> bufs{data.begin(), data.end()};
      callback(bufs, uv::Error{err});
      return;
    }
  
    auto req = std::make_shared<WebSocketWriteReq>(callback);
    raw_uv_ostream os{req->m_bufs, 4096};
  
    // opcode (includes FIN bit)
    os << static_cast<unsigned char>(opcode);
  
    // payload length
    uint64_t size = 0;
    for (auto&& buf : data) size += buf.len;
    if (size < 126) {
      os << static_cast<unsigned char>((m_server ? 0x00 : kFlagMasking) | size);
    } else if (size <= 0xffff) {
      os << static_cast<unsigned char>((m_server ? 0x00 : kFlagMasking) | 126);
      const uint8_t sizeMsb[] = {static_cast<uint8_t>((size >> 8) & 0xff),
                                 static_cast<uint8_t>(size & 0xff)};
      os << ArrayRef<uint8_t>(sizeMsb);
    } else {
      os << static_cast<unsigned char>((m_server ? 0x00 : kFlagMasking) | 127);
      const uint8_t sizeMsb[] = {static_cast<uint8_t>((size >> 56) & 0xff),
                                 static_cast<uint8_t>((size >> 48) & 0xff),
                                 static_cast<uint8_t>((size >> 40) & 0xff),
                                 static_cast<uint8_t>((size >> 32) & 0xff),
                                 static_cast<uint8_t>((size >> 24) & 0xff),
                                 static_cast<uint8_t>((size >> 16) & 0xff),
                                 static_cast<uint8_t>((size >> 8) & 0xff),
                                 static_cast<uint8_t>(size & 0xff)};
      os << ArrayRef<uint8_t>(sizeMsb);
    }
  
    // clients need to mask the input data
    if (!m_server) {
      // generate masking key
      static std::random_device rd;
      static std::default_random_engine gen{rd()};
      std::uniform_int_distribution<unsigned int> dist(0, 255);
      uint8_t key[4];
      for (uint8_t& v : key) v = dist(gen);
      os << ArrayRef<uint8_t>{key, 4};
      // copy and mask data
      int n = 0;
      for (auto&& buf : data) {
        for (auto&& ch : buf.data()) {
          os << static_cast<unsigned char>(static_cast<uint8_t>(ch) ^ key[n++]);
          if (n >= 4) n = 0;
        }
      }
      req->m_startUser = req->m_bufs.size();
      req->m_bufs.append(data.begin(), data.end());
      // don't send the user bufs as we copied their data
      m_stream.Write(ArrayRef<uv::Buffer>{req->m_bufs}.slice(0, req->m_startUser),
                     req);
    } else {
      // servers can just send the buffers directly without masking
      req->m_startUser = req->m_bufs.size();
      req->m_bufs.append(data.begin(), data.end());
      m_stream.Write(req->m_bufs, req);
    }
  }
  