flipper/xplat/FlipperWebSocket/WebSocketTLSClient.cpp

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#ifdef FB_SONARKIT_ENABLED

#include "WebSocketTLSClient.h"
#include <Flipper/ConnectionContextStore.h>
#include <Flipper/FlipperExceptions.h>
#include <Flipper/FlipperTransportTypes.h>
#include <Flipper/FlipperURLSerializer.h>
#include <Flipper/Log.h>
#include <folly/String.h>
#include <folly/futures/Future.h>
#include <folly/io/async/SSLContext.h>
#include <folly/json.h>
#include <websocketpp/common/memory.hpp>
#include <websocketpp/common/thread.hpp>
#include <websocketpp/config/asio.hpp>
#include <cctype>
#include <iomanip>
#include <sstream>
#include <stdexcept>
#include <string>
#include <thread>

namespace facebook {
namespace flipper {

WebSocketTLSClient::WebSocketTLSClient(
    FlipperConnectionEndpoint endpoint,
    std::unique_ptr<FlipperSocketBasePayload> payload,
    folly::EventBase* eventBase)
    : WebSocketTLSClient(
          std::move(endpoint),
          std::move(payload),
          eventBase,
          nullptr) {}

WebSocketTLSClient::WebSocketTLSClient(
    FlipperConnectionEndpoint endpoint,
    std::unique_ptr<FlipperSocketBasePayload> payload,
    folly::EventBase* eventBase,
    ConnectionContextStore* connectionContextStore)
    : BaseClient(
          std::move(endpoint),
          std::move(payload),
          eventBase,
          connectionContextStore) {
  status_ = Status::Unconnected;

  socket_.clear_access_channels(websocketpp::log::alevel::all);
  socket_.clear_error_channels(websocketpp::log::elevel::all);
  socket_.init_asio();
  socket_.start_perpetual();

  thread_ = websocketpp::lib::make_shared<websocketpp::lib::thread>(
      &SocketTLSClient::run, &socket_);
}

WebSocketTLSClient::~WebSocketTLSClient() {
  disconnect();
}

bool WebSocketTLSClient::connect(FlipperConnectionManager* manager) {
  if (status_ != Status::Unconnected) {
    return false;
  }

  status_ = Status::Connecting;

  std::string connectionURL = endpoint_.secure ? "wss://" : "ws://";
  connectionURL += endpoint_.host;
  connectionURL += ":";
  connectionURL += std::to_string(endpoint_.port);

  auto serializer = URLSerializer{};
  payload_->serialize(serializer);
  auto payload = serializer.serialize();

  if (payload.size()) {
    connectionURL += "/?";
    connectionURL += payload;
  }

  socket_.set_tls_init_handler(bind(
      &WebSocketTLSClient::onTLSInit,
      this,
      endpoint_.host.c_str(),
      websocketpp::lib::placeholders::_1));

  auto uri = websocketpp::lib::make_shared<websocketpp::uri>(connectionURL);
  websocketpp::lib::error_code ec;
  connection_ = socket_.get_connection(uri, ec);

  if (ec) {
    status_ = Status::Failed;
    return false;
  }

  handle_ = connection_->get_handle();

  connection_->set_open_handler(websocketpp::lib::bind(
      &WebSocketTLSClient::onOpen,
      this,
      &socket_,
      websocketpp::lib::placeholders::_1));

  connection_->set_message_handler(websocketpp::lib::bind(
      &WebSocketTLSClient::onMessage,
      this,
      &socket_,
      websocketpp::lib::placeholders::_1,
      websocketpp::lib::placeholders::_2));

  connection_->set_fail_handler(websocketpp::lib::bind(
      &WebSocketTLSClient::onFail,
      this,
      &socket_,
      websocketpp::lib::placeholders::_1));

  connection_->set_close_handler(websocketpp::lib::bind(
      &WebSocketTLSClient::onClose,
      this,
      &socket_,
      websocketpp::lib::placeholders::_1));

  auto connected = connected_.get_future();

  socket_.connect(connection_);

  auto state = connected.wait_for(std::chrono::seconds(10));
  if (state == std::future_status::ready) {
    return connected.get();
  }

  disconnect();

  return false;
}

void WebSocketTLSClient::disconnect() {
  socket_.stop_perpetual();

  if (status_ == Status::Connecting || status_ == Status::Open ||
      status_ == Status::Failed) {
    websocketpp::lib::error_code ec;
    socket_.close(handle_, websocketpp::close::status::going_away, "", ec);
  }
  socket_.stop();
  status_ = Status::Closed;
  if (thread_ && thread_->joinable()) {
    thread_->join();
  }

  thread_ = nullptr;
  eventBase_->add(
      [eventHandler = eventHandler_]() { eventHandler(SocketEvent::CLOSE); });
}

void WebSocketTLSClient::send(
    const folly::dynamic& message,
    SocketSendHandler completion) {
  std::string json = folly::toJson(message);
  send(json, std::move(completion));
}

void WebSocketTLSClient::send(
    const std::string& message,
    SocketSendHandler completion) {
  websocketpp::lib::error_code ec;
  socket_.send(
      handle_,
      &message[0],
      message.size(),
      websocketpp::frame::opcode::text,
      ec);
  completion();
}

/**
    Only ever used for insecure connections to receive the device_id from a
    signCertificate request. If the intended usage ever changes, then a better
    approach needs to be put in place.
 */
void WebSocketTLSClient::sendExpectResponse(
    const std::string& message,
    SocketSendExpectResponseHandler completion) {
  connection_->set_message_handler(
      [completion, eventBase = eventBase_](
          websocketpp::connection_hdl hdl, SocketTLSClient::message_ptr msg) {
        const std::string& payload = msg->get_payload();
        eventBase->add([completion, payload] { completion(payload, false); });
      });
  websocketpp::lib::error_code ec;
  socket_.send(
      handle_,
      &message[0],
      message.size(),
      websocketpp::frame::opcode::text,
      ec);
  if (ec) {
    auto reason = ec.message();
    completion(reason, true);
  }
}

void WebSocketTLSClient::onOpen(
    SocketTLSClient* c,
    websocketpp::connection_hdl hdl) {
  if (status_ == Status::Connecting) {
    connected_.set_value(true);
  }

  status_ = Status::Initializing;
  eventBase_->add(
      [eventHandler = eventHandler_]() { eventHandler(SocketEvent::OPEN); });
}

void WebSocketTLSClient::onMessage(
    SocketTLSClient* c,
    websocketpp::connection_hdl hdl,
    SocketTLSClient::message_ptr msg) {
  const std::string& payload = msg->get_payload();
  if (messageHandler_) {
    eventBase_->add([payload, messageHandler = messageHandler_]() {
      messageHandler(payload);
    });
  }
}

void WebSocketTLSClient::onFail(
    SocketTLSClient* c,
    websocketpp::connection_hdl hdl) {
  SocketTLSClient::connection_ptr con = c->get_con_from_hdl(hdl);
  auto server = con->get_response_header("Server");
  auto reason = con->get_ec().message();
  auto sslError =
      (reason.find("TLS handshake failed") != std::string::npos ||
       reason.find("Generic TLS related error") != std::string::npos);

  if (status_ == Status::Connecting) {
    if (sslError) {
      try {
        connected_.set_exception(
            std::make_exception_ptr(SSLException("SSL handshake failed")));
      } catch (...) {
        // set_exception() may throw an exception
        // In that case, just set the value to false.
        connected_.set_value(false);
      }
    } else {
      connected_.set_value(false);
    }
  }
  status_ = Status::Failed;
  eventBase_->add([eventHandler = eventHandler_, sslError]() {
    if (sslError) {
      eventHandler(SocketEvent::SSL_ERROR);
    } else {
      eventHandler(SocketEvent::ERROR);
    }
  });
}

void WebSocketTLSClient::onClose(
    SocketTLSClient* c,
    websocketpp::connection_hdl hdl) {
  status_ = Status::Closed;
  eventBase_->add(
      [eventHandler = eventHandler_]() { eventHandler(SocketEvent::CLOSE); });
}

SocketTLSContext WebSocketTLSClient::onTLSInit(
    const char* hostname,
    websocketpp::connection_hdl) {
  namespace asio = websocketpp::lib::asio;
  SocketTLSContext ctx = websocketpp::lib::make_shared<asio::ssl::context>(
      asio::ssl::context::sslv23);

  ctx->set_options(
      asio::ssl::context::default_workarounds | asio::ssl::context::no_sslv2 |
      asio::ssl::context::no_sslv3 | asio::ssl::context::single_dh_use);

  ctx->set_verify_mode(asio::ssl::verify_peer);
  ctx->load_verify_file(connectionContextStore_->getPath(
      ConnectionContextStore::StoreItem::FLIPPER_CA));

  asio::error_code error;
  ctx->use_certificate_file(
      connectionContextStore_->getPath(
          ConnectionContextStore::StoreItem::CLIENT_CERT),
      asio::ssl::context::pem,
      error);
  ctx->use_private_key_file(
      connectionContextStore_->getPath(
          ConnectionContextStore::StoreItem::PRIVATE_KEY),
      asio::ssl::context::pem,
      error);

  return ctx;
}

} // namespace flipper
} // namespace facebook

#endif