websocket_mock_server.cc

Go to the documentation of this file.

/*
 * Copyright (C) Max Dymond, <cmeister2@gmail.com>, et al.
 *
 * SPDX-License-Identifier: curl
 */

 
#include "proto_fuzzer/websocket_mock_server.h"
 
#include <curl/websockets.h>
 
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <string>
#include <utility>
 
#include "proto_fuzzer/option_apply.h"
#include "proto_fuzzer/scenario_limits.h"
#include "proto_fuzzer/ws_accept_key.h"
#include "proto_fuzzer/ws_frame.h"
 
namespace proto_fuzzer {
 
namespace {
 
// Iteration caps for the post-handshake manual WS drive. Kept small because
// every iteration makes a curl_ws_recv / curl_ws_send call.
constexpr std::size_t kMaxWsRecvIterations = 128;
constexpr std::size_t kMaxWsSendIterations = 32;

constexpr unsigned int kWsSendFlagMatrix[] = {
    CURLWS_CONT,
    CURLWS_TEXT,
    CURLWS_BINARY,
    CURLWS_TEXT | CURLWS_OFFSET,
    CURLWS_BINARY | CURLWS_OFFSET,
    CURLWS_TEXT | CURLWS_CONT,
    CURLWS_CONT,
    0,
    CURLWS_TEXT,
    CURLWS_BINARY | CURLWS_CONT,
    CURLWS_BINARY,
    CURLWS_PING,
    CURLWS_PONG,
    CURLWS_CLOSE,
    0,
    CURLWS_CLOSE | CURLWS_CONT,
    CURLWS_PING | CURLWS_CONT,
    CURLWS_PONG | CURLWS_CONT,
    CURLWS_TEXT | CURLWS_BINARY,
};

std::string ExtractWebSocketKey(const std::string& request) {
  static const char kHeader[] = "Sec-WebSocket-Key:";
  std::size_t pos = request.find(kHeader);
  if (pos == std::string::npos) {
    return {};
  }
  pos += sizeof(kHeader) - 1;
  while (pos < request.size() && (request[pos] == ' ' || request[pos] == '\t')) {
    ++pos;
  }
  std::size_t end = request.find("\r\n", pos);
  if (end == std::string::npos) {
    return {};
  }
  while (end > pos && (request[end - 1] == ' ' || request[end - 1] == '\t')) {
    --end;
  }
  return request.substr(pos, end - pos);
}

std::string ComputeWebSocketAccept(const std::string& key) { return proto_fuzzer::ComputeWebSocketAcceptKey(key); }

std::vector<std::string> BuildFrameChunks(const curl::fuzzer::proto::Connection& conn) {
  std::vector<std::string> chunks;
  chunks.reserve(scenario_limits::kMaxResponseChunks);
  const std::size_t raw_budget = std::min<std::size_t>(scenario_limits::kMaxResponseChunks, conn.on_readable_size());
  for (std::size_t i = 0; i < raw_budget; ++i) {
    chunks.emplace_back(conn.on_readable(i));
  }
  const std::size_t frame_budget = scenario_limits::kMaxResponseChunks - chunks.size();
  const std::size_t frame_count = std::min<std::size_t>(frame_budget, conn.server_frames_size());
  for (std::size_t i = 0; i < frame_count; ++i) {
    chunks.emplace_back(SerializeWebSocketFrame(conn.server_frames(static_cast<int>(i))));
  }
  return chunks;
}

size_t WebSocketWriteCallback(void* /*contents*/, size_t size, size_t nmemb, void* userdata) {
  auto* server = static_cast<WebSocketMockServer*>(userdata);
  if (server != nullptr) {
    CURL* easy = server->easy_handle();
    if (easy != nullptr) {
      (void)curl_ws_meta(easy);
      if (!server->ws_probe_fired()) {
        server->MarkWsProbeFired();
        static unsigned char kProbe[16384];
        std::fill(kProbe, kProbe + sizeof(kProbe), 'P');
        std::size_t sent = 0;
        (void)curl_ws_send(easy, kProbe, sizeof(kProbe), &sent, 0, 0);
      }
    }
  }
  return size * nmemb;
}

std::size_t DrainWsRecv(CURL* easy) {
  std::size_t total_received = 0;
  for (std::size_t i = 0; i < kMaxWsRecvIterations; ++i) {
    unsigned char buffer[4096];
    std::size_t nread = 0;
    const struct curl_ws_frame* meta = nullptr;
    CURLcode rr = curl_ws_recv(easy, buffer, sizeof(buffer), &nread, &meta);
    if (rr == CURLE_AGAIN) {
      break;
    }
    if (rr != CURLE_OK && rr != CURLE_GOT_NOTHING) {
      break;
    }
    if (nread == 0 && meta == nullptr) {
      break;
    }
    total_received += nread;
  }
  return total_received;
}
 
}  // namespace

bool ScenarioRequestsManualWsDrive(const curl::fuzzer::proto::Scenario& scenario) {
  bool manual_delivery = false;
  const std::size_t option_count = RuntimeOptionCount(scenario);
  for (std::size_t index = 0; index < option_count; ++index) {
    const auto& option = scenario.options(static_cast<int>(index));
    if (option.option_id() != curl::fuzzer::proto::CURLOPT_CONNECT_ONLY) {
      continue;
    }
    const std::uint64_t value = DecodeIntegralOptionValue(option);
    if (value <= 2) {
      manual_delivery = value == 2;
    }
  }
  return manual_delivery;
}

WebSocketMockServer::WebSocketMockServer()
    : next_chunk_(0),
      manual_delivery_(false),
      handshake_sent_(false),
      ws_probe_fired_(false),
      easy_handle_(nullptr),
      manual_received_bytes_(0) {}

WebSocketMockServer::~WebSocketMockServer() = default;

void WebSocketMockServer::Install(CURL* easy) {
  MockServerBase::Install(easy);
  easy_handle_ = easy;
  ws_probe_fired_ = false;
  curl_easy_setopt(easy, CURLOPT_WRITEFUNCTION, &WebSocketWriteCallback);
  curl_easy_setopt(easy, CURLOPT_WRITEDATA, this);
  curl_easy_setopt(easy, CURLOPT_HEADERFUNCTION, &WebSocketWriteCallback);
}

bool WebSocketMockServer::ws_probe_fired() const { return ws_probe_fired_; }

void WebSocketMockServer::MarkWsProbeFired() { ws_probe_fired_ = true; }

CURL* WebSocketMockServer::easy_handle() const { return easy_handle_; }

std::size_t WebSocketMockServer::manual_received_bytes() const { return manual_received_bytes_; }

void WebSocketMockServer::SetFrames(std::vector<std::string> frames) {
  frames_ = std::move(frames);
  next_chunk_ = 0;
}

void WebSocketMockServer::SetManualDelivery(bool manual) { manual_delivery_ = manual; }

bool WebSocketMockServer::manual_delivery() const { return manual_delivery_; }

bool WebSocketMockServer::handshake_sent() const { return handshake_sent_; }

bool WebSocketMockServer::has_more_chunks() const { return next_chunk_ < frames_.size(); }

std::size_t WebSocketMockServer::remaining_chunks() const {
  return next_chunk_ >= frames_.size() ? 0 : frames_.size() - next_chunk_;
}

const std::string& WebSocketMockServer::PeekChunk(std::size_t index) const { return frames_[next_chunk_ + index]; }

void WebSocketMockServer::ConsumeChunk() {
  if (next_chunk_ < frames_.size()) {
    ++next_chunk_;
  }
}

std::unique_ptr<MockConnection> WebSocketMockServer::CreateConnection() { return std::make_unique<MockConnection>(); }

std::size_t WebSocketMockServer::DrainHandshakeData() {
  if (!connection_) {
    return 0;
  }
  const std::size_t size_before = ws_request_buffer_.size();
  connection_->ReadAvailable(&ws_request_buffer_);
  return ws_request_buffer_.size() - size_before;
}

void WebSocketMockServer::FlushTransport() {}

void WebSocketMockServer::ObserveIncomingData(const unsigned char* data, std::size_t size) {
  if (!handshake_sent_ && data != nullptr && size != 0) {
    ws_request_buffer_.append(reinterpret_cast<const char*>(data), size);
  }
}

curl_socket_t WebSocketMockServer::HandleOpenSocket(curlsocktype purpose, struct curl_sockaddr* address) {
  (void)purpose;
  (void)address;
  if (connection_) {
    return CURL_SOCKET_BAD;
  }
  connection_ = CreateConnection();
  if (!connection_ || !connection_->ok()) {
    connection_.reset();
    return CURL_SOCKET_BAD;
  }
  connection_->SetIncomingDataObserver(this);
  ApplyPendingBackpressure();
  // Wait for curl's Upgrade request before we write anything — the drive
  // loop calls TryAdvanceHandshake() to drive that exchange.
  return connection_->take_client_fd();
}

bool WebSocketMockServer::PushRawBytes(const unsigned char* data, std::size_t size) {
  if (!connection_) {
    return false;
  }
  connection_->DrainIncoming();
  const bool written = connection_->WriteAll(data, size);
  if (written) {
    FlushTransport();
  }
  return written;
}

bool WebSocketMockServer::DeliverNextChunk() {
  if (!connection_ || next_chunk_ >= frames_.size()) {
    return false;
  }
  connection_->DrainIncoming();
  const std::string& chunk = frames_[next_chunk_++];
  if (!chunk.empty()) {
    if (connection_->WriteAll(reinterpret_cast<const unsigned char*>(chunk.data()), chunk.size())) {
      FlushTransport();
    }
  }
  if (next_chunk_ >= frames_.size()) {
    connection_->ShutdownWrite();
  }
  return true;
}

bool WebSocketMockServer::TryAdvanceHandshake(bool* made_progress) {
  if (made_progress != nullptr) {
    *made_progress = false;
  }
  if (handshake_sent_ || !connection_) {
    return handshake_sent_;
  }
  if (DrainHandshakeData() != 0 && made_progress != nullptr) {
    *made_progress = true;
  }
  if (ws_request_buffer_.find("\r\n\r\n") == std::string::npos) {
    return false;
  }
  std::string key = ExtractWebSocketKey(ws_request_buffer_);
  // Even if parsing failed, reply with *something* so curl doesn't wedge.
  // A bad Accept exercises curl's handshake-error path.
  std::string accept = key.empty() ? std::string("AAAAAAAAAAAAAAAAAAAAAAAAAAA=") : ComputeWebSocketAccept(key);
  std::string response =
      "HTTP/1.1 101 Switching Protocols\r\n"
      "Upgrade: websocket\r\n"
      "Connection: Upgrade\r\n"
      "Sec-WebSocket-Accept: " +
      accept + "\r\n\r\n";
  // CONNECT_ONLY may leave already-read bytes after the 101 headers. Put one
  // bounded frame in the same peer write so curl has to preserve those bytes
  // for the later curl_ws_recv call instead of discarding them when the
  // opening transfer finishes. Remaining chunks still exercise the explicit
  // manual-delivery tail below.
  if (manual_delivery() && has_more_chunks()) {
    response.append(PeekChunk(0));
    ConsumeChunk();
  }
  connection_->WriteAll(reinterpret_cast<const unsigned char*>(response.data()), response.size());
  FlushTransport();
  handshake_sent_ = true;
  if (made_progress != nullptr) {
    *made_progress = true;
  }
  return true;
}

void WebSocketMockServer::RunLoop(CURLM* multi, CURL* easy, const curl::fuzzer::proto::Scenario& scenario) {
  SetManualDelivery(ScenarioRequestsManualWsDrive(scenario));
  manual_received_bytes_ = 0;
  // initial_response is unused in WS mode; we synthesise the 101 dynamically
  // from curl's Upgrade request.
  SetFrames(BuildFrameChunks(scenario.connection()));
 
  int still_running = 1;
  int idle_iterations = 0;
  int drive_iterations = 0;
  bool pollset_probed = false;
  const bool timed_drive = UsesTimedDrive(scenario);
  const int idle_limit = timed_drive ? kMaxTimedIdleIterations : kMaxIdleIterations;
  CURLMcode rc = CURLM_OK;
 
  while (still_running && idle_iterations < idle_limit && drive_iterations++ < kMaxDriveIterations) {
    bool made_progress = false;
    const int running_before = still_running;
    rc = curl_multi_perform(multi, &still_running);
    if (rc != CURLM_OK) {
      break;
    }
    made_progress = still_running != running_before;
    if (timed_drive && still_running && !pollset_probed) {
      // At this point curl has created the Upgrade connection but the mock has
      // not necessarily replied, which gives curl_multi_poll a meaningful WS
      // filter state without adding a timeout to fixed fast-lane inputs.
      ProbeMultiPollset(multi);
      pollset_probed = true;
    }
    // Drive the 101 handshake on every iteration; no-op once sent.
    if (!handshake_sent()) {
      bool handshake_progress = false;
      (void)TryAdvanceHandshake(&handshake_progress);
      made_progress = handshake_progress || made_progress;
    }
    if (!still_running) {
      break;
    }
 
    // Only push frame chunks in streaming mode — in manual mode the caller
    // will push them below after the handshake.
    if (!manual_delivery() && handshake_sent() && has_more_chunks()) {
      made_progress = DeliverNextChunk() || made_progress;
    } else if (handshake_sent() && connection() != nullptr) {
      // Uploading WS scenarios can have no server frames at all. Drain their
      // encoded client data here so they make deterministic progress and do
      // not have to wait for CURLOPT_TIMEOUT_MS just to exercise cr_ws_read.
      made_progress = connection()->DrainIncoming() != 0 || made_progress;
    }
 
    if (made_progress) {
      idle_iterations = 0;
      continue;
    }
 
    if (timed_drive) {
      (void)WaitOnMultiFdset(multi, &rc);
      if (rc != CURLM_OK) {
        break;
      }
    }
    ++idle_iterations;
  }
 
  if (!manual_delivery() || !handshake_sent()) {
    return;
  }
 
  // Manual-drive tail: feed every remaining scripted chunk straight onto the
  // server fd as raw frame bytes, draining curl_ws_recv between each push.
  while (remaining_chunks() > 0) {
    const std::string& chunk = PeekChunk(0);
    if (!chunk.empty()) {
      PushRawBytes(reinterpret_cast<const unsigned char*>(chunk.data()), chunk.size());
    }
    ConsumeChunk();
    manual_received_bytes_ += DrainWsRecv(easy);
  }
  // Final drain in case frame parsing produced more work after the last push.
  manual_received_bytes_ += DrainWsRecv(easy);
 
  const auto& probes = scenario.connection().manual_probes();
  static const unsigned char kPayload[] = "hello-from-proto-fuzzer";
  const std::size_t payload_len = sizeof(kPayload) - 1;
 
  if (probes.flag_matrix()) {
    // Exercise curl_ws_send with a fixed matrix of flags. We don't care whether
    // the send actually lands on the wire — the point is to reach the encode
    // paths in ws_enc_add_frame / ws_enc_write_head.
    std::size_t iteration = 0;
    for (unsigned int flags : kWsSendFlagMatrix) {
      if (iteration++ >= kMaxWsSendIterations) {
        break;
      }
      // Announce the frame via the public curl_ws_start_frame entrypoint before
      // the actual send. In non-raw mode this writes the frame head into the
      // sendbuf; the follow-up curl_ws_send with the same flags finishes the
      // exchange or fails cleanly on invalid flag combos.
      (void)curl_ws_start_frame(easy, flags, static_cast<curl_off_t>(payload_len));
      std::size_t sent = 0;
      curl_off_t fragsize = (flags & CURLWS_OFFSET) ? static_cast<curl_off_t>(payload_len) : 0;
      (void)curl_ws_send(easy, kPayload, payload_len, &sent, fragsize, flags);
      if (connection() != nullptr) {
        connection()->DrainIncoming();
      }
    }
  }
 
  if (probes.unaligned_send()) {
    // Multi-call mis-sized send probe: declare a fragsize=200 frame, send 23
    // bytes, then call curl_ws_send again with a buflen much bigger than the
    // remaining payload. Hits the "unaligned frame size" failf in ws_enc_send.
    constexpr curl_off_t kBigFrag = 200;
    (void)curl_ws_start_frame(easy, CURLWS_TEXT | CURLWS_OFFSET, kBigFrag);
    std::size_t sent = 0;
    (void)curl_ws_send(easy, kPayload, payload_len, &sent, kBigFrag, CURLWS_TEXT | CURLWS_OFFSET);
    // Now enc.payload_remain ≈ kBigFrag - payload_len. A follow-up send with
    // buflen > remaining trips the guard at ws_enc_send:~1050.
    std::size_t sent2 = 0;
    constexpr std::size_t kOverrun = 500;
    unsigned char overrun[kOverrun];
    std::fill(overrun, overrun + kOverrun, 'X');
    (void)curl_ws_send(easy, overrun, kOverrun, &sent2, kBigFrag, CURLWS_TEXT | CURLWS_OFFSET);
    if (connection() != nullptr) {
      connection()->DrainIncoming();
    }
  }
 
  if (probes.raw_send()) {
    // Raw-mode send path: reachable only when CURLOPT_WS_OPTIONS has
    // CURLWS_RAW_MODE set. curl_ws_send with flags=0, fragsize=0 takes the
    // data->set.ws_raw_mode branch → ws_send_raw. No-op for non-raw scenarios
    // (hits the "no flags given" failure path instead).
    std::size_t sent = 0;
    (void)curl_ws_send(easy, kPayload, payload_len, &sent, 0, 0);
    if (connection() != nullptr) {
      connection()->DrainIncoming();
    }
  }
}

void WebSocketMockServer::ObserveState() {
#if defined(PROTO_FUZZER_OBSERVE_STATE)
  MockServerBase::ObserveState();
  const volatile std::size_t websocket_observation = manual_received_bytes();
  (void)websocket_observation;
#endif
}
 
}  // namespace proto_fuzzer