| /*----------------------------------------------------------------------------*/ |
| /* 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); |
| } |
| } |