couchbase 3.5.6 → 3.6.0

data/ext/couchbase/core/websocket_codec.cxx

@@ -0,0 +1,647 @@
+/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
+/*
+ * Copyright 2024-Present Couchbase, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under
+ * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+ * ANY KIND, either express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+#include "websocket_codec.hxx"
+
+#include "core/crypto/cbcrypto.h"
+#include "core/platform/base64.h"
+#include "core/platform/random.h"
+
+#include <llhttp.h>
+#include <spdlog/fmt/bundled/core.h>
+
+#include <algorithm>
+#include <array>
+#include <cstdint>
+#include <cstring>
+#include <map>
+#include <optional>
+#include <random>
+#include <variant>
+
+namespace couchbase::core
+{
+namespace
+{
+/*
+     0                   1                   2                   3
+     0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+    +-+-+-+-+-------+-+-------------+-------------------------------+
+    |F|R|R|R| opcode|M| Payload len | Extended payload length       |
+    |I|S|S|S|   (4) |A|     (7)     |          (16/64)              |
+    |N|V|V|V|       |S|             | (if payload len==126/127)     |
+    | |1|2|3|       |K|             |                               |
+    +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
+    |     Extended payload length continued, if payload len == 127  |
+    + - - - - - - - - - - - - - - - +-------------------------------+
+    |                               | Masking-key, if MASK set to 1 |
+    +-------------------------------+-------------------------------+
+    | Masking-key (continued)       | Payload Data                  |
+    +-------------------------------- - - - - - - - - - - - - - - - +
+    :                        Payload Data continued ...             :
+    + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - +
+    |                        Payload Data continued ...             |
+    +---------------------------------------------------------------+
+ */
+constexpr std::uint8_t flag_fin{ 0b1000'0000 };
+constexpr std::uint8_t flag_mask{ 0b1000'0000 };
+
+constexpr std::uint8_t reserved_bit_mask{ 0b0111'0000 };
+constexpr std::uint8_t opcode_mask{ 0b0000'1111 };
+constexpr std::uint8_t payload_length_7_mask{ 0b0111'1111 };
+
+constexpr std::uint8_t opcode_continuation{ 0x00 };
+constexpr std::uint8_t opcode_text{ 0x01 };
+constexpr std::uint8_t opcode_binary{ 0x2 };
+constexpr std::uint8_t opcode_close{ 0x08 };
+constexpr std::uint8_t opcode_ping{ 0x09 };
+constexpr std::uint8_t opcode_pong{ 0x0a };
+
+auto
+generate_masking_key() -> std::array<std::byte, 4>
+{
+  thread_local std::random_device rd;
+  thread_local std::uniform_int_distribution<std::uint16_t> dist(0, 0xff);
+
+  return {
+    static_cast<std::byte>(dist(rd)),
+    static_cast<std::byte>(dist(rd)),
+    static_cast<std::byte>(dist(rd)),
+    static_cast<std::byte>(dist(rd)),
+  };
+}
+
+auto
+generate_session_key() -> std::string
+{
+  const couchbase::core::RandomGenerator randomGenerator;
+  std::array<std::byte, 16> key{};
+  if (!core::RandomGenerator::getBytes(key.data(), key.size())) {
+    throw std::bad_alloc();
+  }
+  return core::base64::encode(key, false);
+}
+
+auto
+signature_is_valid(const std::string& session_key, const std::string& signature) -> bool
+{
+  // RFC 6455, Section 1.3
+  static const std::string websocket_guid{ "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" };
+  const std::string salted_key{ fmt::format("{}{}", session_key, websocket_guid) };
+  auto hash = core::crypto::digest(core::crypto::Algorithm::ALG_SHA1, salted_key);
+  return core::base64::encode(hash, false) == signature;
+}
+
+auto
+case_insensitive_equals(const std::string& str1, const std::string& str2) -> bool
+{
+  if (str1.size() != str2.size()) {
+    return false;
+  }
+
+  return std::equal(str1.begin(), str1.end(), str2.begin(), [](unsigned char c1, unsigned char c2) {
+    return std::tolower(c1) == std::tolower(c2);
+  });
+}
+
+struct context {
+  websocket_callbacks& callbacks;
+  const websocket_codec& ws;
+};
+} // namespace
+
+class websocket_handler
+{
+public:
+  websocket_handler() = default;
+  websocket_handler(const websocket_handler&) = delete;
+  websocket_handler(websocket_handler&&) noexcept = delete;
+  auto operator=(const websocket_handler&) -> websocket_handler& = delete;
+  auto operator=(websocket_handler&&) noexcept -> websocket_handler& = delete;
+  virtual ~websocket_handler() = default;
+
+  virtual auto feed(gsl::span<std::byte> data, const context& ctx)
+    -> std::unique_ptr<websocket_handler> = 0;
+};
+
+namespace
+{
+
+void
+mask_payload_data(gsl::span<std::byte> masking_key, gsl::span<std::byte> payload)
+{
+  for (std::size_t i = 0; i < payload.size(); ++i) {
+    payload[i] ^= masking_key[i % masking_key.size()];
+  }
+}
+
+struct decoded_frame {
+  std::uint8_t type;
+  gsl::span<std::byte> payload;
+  std::size_t consumed_bytes;
+  bool expected_continuation;
+};
+
+struct partial_frame {
+  std::uint8_t type;
+  std::vector<std::byte> payload;
+};
+
+struct decoding_error {
+  std::string message;
+};
+
+struct need_more_data {
+};
+
+using decode_status = std::variant<decoded_frame, decoding_error, need_more_data>;
+
+auto
+decode_frame(gsl::span<std::byte> data, bool expected_continuation) -> decode_status;
+
+class error_handler : public websocket_handler
+{
+private:
+  std::string message_;
+
+public:
+  error_handler(std::string message, const context& ctx)
+    : message_{ std::move(message) }
+  {
+    ctx.callbacks.on_error(ctx.ws, message_);
+  }
+
+  auto feed(gsl::span<std::byte> /* data */, const context& ctx)
+    -> std::unique_ptr<websocket_handler> override
+  {
+    ctx.callbacks.on_error(ctx.ws, message_);
+    return nullptr;
+  }
+};
+
+class data_handler : public websocket_handler
+{
+private:
+  std::vector<std::byte> buffer_{};
+  std::optional<partial_frame> partial_response_{};
+
+public:
+  explicit data_handler(const context& ctx, gsl::span<std::byte> remaining = {})
+    : buffer_{ remaining.begin(), remaining.end() }
+  {
+    ctx.callbacks.on_ready(ctx.ws);
+  }
+
+  auto feed(gsl::span<std::byte> data, const context& ctx)
+    -> std::unique_ptr<websocket_handler> override
+  {
+    std::vector<std::byte> buffer{};
+    if (!buffer_.empty()) {
+      std::swap(buffer_, buffer);
+      std::copy(data.begin(), data.end(), std::back_insert_iterator(buffer));
+      data = buffer;
+    }
+    while (!data.empty()) {
+      auto status = decode_frame(data, partial_response_.has_value());
+      if (std::holds_alternative<decoding_error>(status)) {
+        return std::make_unique<error_handler>(
+          fmt::format("Decoding error: {}", std::get<decoding_error>(status).message), ctx);
+      }
+      if (std::holds_alternative<need_more_data>(status)) {
+        std::copy(data.begin(), data.end(), std::back_insert_iterator(buffer_));
+        return nullptr;
+      }
+      if (std::holds_alternative<decoded_frame>(status)) {
+        auto frame = std::get<decoded_frame>(status);
+        switch (frame.type) {
+          case opcode_text:
+            if (frame.expected_continuation) {
+              partial_response_ = partial_frame{
+                frame.type,
+                { frame.payload.begin(), frame.payload.end() },
+              };
+            } else {
+              ctx.callbacks.on_text(ctx.ws, frame.payload);
+            }
+            break;
+          case opcode_binary:
+            if (frame.expected_continuation) {
+              partial_response_ = partial_frame{
+                frame.type,
+                { frame.payload.begin(), frame.payload.end() },
+              };
+            } else {
+              ctx.callbacks.on_binary(ctx.ws, frame.payload);
+            }
+            break;
+          case opcode_close:
+            ctx.callbacks.on_close(ctx.ws, frame.payload);
+            break;
+          case opcode_ping:
+            ctx.callbacks.on_ping(ctx.ws, frame.payload);
+            break;
+          case opcode_pong:
+            ctx.callbacks.on_pong(ctx.ws, frame.payload);
+            break;
+          case opcode_continuation:
+            if (auto& response = partial_response_; response.has_value()) {
+              if (frame.expected_continuation) {
+                std::copy(frame.payload.begin(),
+                          frame.payload.end(),
+                          std::back_insert_iterator(response->payload));
+              } else {
+                if (response->type == opcode_text) {
+                  ctx.callbacks.on_text(ctx.ws, response->payload);
+                } else {
+                  ctx.callbacks.on_binary(ctx.ws, response->payload);
+                }
+              }
+            } else {
+              return std::make_unique<error_handler>("Unexpected continuation frame", ctx);
+            }
+            break;
+          default:
+            return std::make_unique<error_handler>(
+              fmt::format("Unexpected frame.type: {}", frame.type), ctx);
+        }
+        data = data.subspan(frame.consumed_bytes);
+      }
+    }
+
+    return nullptr;
+  }
+};
+
+class open_handshake : public websocket_handler
+{
+public:
+  open_handshake()
+  {
+    llhttp_settings_init(&settings_);
+    settings_.on_status = on_status;
+    settings_.on_header_field = on_header_field;
+    settings_.on_header_value = on_header_value;
+    settings_.on_body = on_body;
+    settings_.on_message_complete = on_message_complete;
+    llhttp_init(&parser_, HTTP_RESPONSE, &settings_);
+    parser_.data = this;
+  }
+
+  auto feed(gsl::span<std::byte> data, const context& ctx)
+    -> std::unique_ptr<websocket_handler> override
+  {
+    auto error = llhttp_execute(&parser_, reinterpret_cast<const char*>(data.data()), data.size());
+    if (error != HPE_OK && error != HPE_PAUSED_UPGRADE) {
+      return std::make_unique<error_handler>(
+        fmt::format("Failed to parse HTTP response: {}",
+                    llhttp_errno_name(llhttp_get_errno(&parser_))),
+        ctx);
+    }
+    if (complete_) {
+      if (status_code_ != 101) {
+        return std::make_unique<error_handler>(
+          fmt::format("Response status must be 101. ({} {})", status_code_, status_message_), ctx);
+      }
+      if (!case_insensitive_equals(headers_["connection"], "upgrade")) {
+        return std::make_unique<error_handler>(
+          "Request has MUST contain Connection header field with value including \"Upgrade\"", ctx);
+      }
+      if (!case_insensitive_equals(headers_["upgrade"], "websocket")) {
+        return std::make_unique<error_handler>(
+          "Request has MUST contain Upgrade header field with value including \"websocket\"", ctx);
+      }
+
+      if (!signature_is_valid(ctx.ws.session_key(), headers_["sec-websocket-accept"])) {
+        return std::make_unique<error_handler>(
+          "Request has MUST contain Sec-WebSocket-Accept with valid key", ctx);
+      }
+
+      if (error == HPE_PAUSED_UPGRADE) {
+        auto bytes_parsed =
+          llhttp_get_error_pos(&parser_) - reinterpret_cast<const char*>(data.data());
+        return std::make_unique<data_handler>(ctx,
+                                              data.subspan(static_cast<std::size_t>(bytes_parsed)));
+      }
+      return std::make_unique<data_handler>(ctx);
+    }
+    return nullptr;
+  }
+
+private:
+  static auto on_status(llhttp_t* parser, const char* at, std::size_t length) -> int
+  {
+    auto* self = static_cast<open_handshake*>(parser->data);
+    self->status_message_.assign(at, length);
+    self->status_code_ = parser->status_code;
+    return 0;
+  }
+
+  static auto on_header_field(llhttp_t* parser, const char* at, std::size_t length) -> int
+  {
+    auto* self = static_cast<open_handshake*>(parser->data);
+    self->header_field_.clear();
+    self->header_field_.reserve(length);
+    for (std::size_t i = 0; i < length; ++i) {
+      self->header_field_ += static_cast<char>(std::tolower(at[i]));
+    }
+    return 0;
+  }
+
+  static auto on_header_value(llhttp_t* parser, const char* at, std::size_t length) -> int
+  {
+    auto* self = static_cast<open_handshake*>(parser->data);
+    self->headers_[self->header_field_] = std::string(at, length);
+    return 0;
+  }
+
+  static auto on_body(llhttp_t* parser, const char* at, std::size_t length) -> int
+  {
+    auto* self = static_cast<open_handshake*>(parser->data);
+    self->body_.append(std::string_view{ at, length });
+    return 0;
+  }
+
+  static auto on_body_after_upgrade(llhttp_t* parser, const char* at, std::size_t length) -> int
+  {
+    auto* self = static_cast<open_handshake*>(parser->data);
+    self->body_ = std::string_view{ at, length };
+    return 0;
+  }
+
+  static auto on_message_complete(llhttp_t* parser) -> int
+  {
+    auto* self = static_cast<open_handshake*>(parser->data);
+    self->complete_ = true;
+    return 0;
+  }
+
+  llhttp_settings_t settings_{};
+  llhttp_t parser_{};
+
+  std::string header_field_{};
+  bool complete_{};
+
+  std::uint32_t status_code_{ 0 };
+  std::string status_message_{};
+  std::map<std::string, std::string> headers_{};
+  std::string body_{};
+};
+
+constexpr auto
+is_data_frame(std::uint8_t opcode) -> bool
+{
+  switch (opcode) {
+    case opcode_text:
+    case opcode_binary:
+      return true;
+    default:
+      break;
+  }
+  return false;
+}
+
+constexpr auto
+is_control_frame(std::uint8_t opcode) -> bool
+{
+  switch (opcode) {
+    case opcode_close:
+    case opcode_ping:
+    case opcode_pong:
+      return true;
+    default:
+      break;
+  }
+  return false;
+}
+
+auto
+encode_payload_length(std::size_t length) -> std::vector<std::byte>
+{
+  if (length <= 125) { // 7 bit
+    const std::uint8_t field_7_bit = static_cast<std::uint8_t>(length) | flag_mask;
+    return { std::byte{ field_7_bit } };
+  }
+  if (length <= 0xFFFF) { // 7 + 16 bit
+    constexpr std::uint8_t field_7_bit = 126 | flag_mask;
+    return {
+      std::byte{ field_7_bit },
+      static_cast<std::byte>((length >> 8) & 0xff),
+      static_cast<std::byte>(length & 0xff),
+    };
+  }
+  // 7 + 64
+  constexpr std::uint8_t field_7_bit = 127 | flag_mask;
+  return {
+    std::byte{ field_7_bit },
+    static_cast<std::byte>((length >> 56) & 0xff),
+    static_cast<std::byte>((length >> 48) & 0xff),
+    static_cast<std::byte>((length >> 40) & 0xff),
+    static_cast<std::byte>((length >> 32) & 0xff),
+    static_cast<std::byte>((length >> 24) & 0xff),
+    static_cast<std::byte>((length >> 16) & 0xff),
+    static_cast<std::byte>((length >> 8) & 0xff),
+    static_cast<std::byte>(length & 0xff),
+  };
+}
+
+auto
+decode_uint64(gsl::span<std::byte> data) -> std::uint64_t
+{
+  std::uint64_t result{};
+  std::memcpy(&result, data.data(), sizeof(result));
+
+  return                                  //
+    (result & 0x00000000000000FF) << 56 | //
+    (result & 0x000000000000FF00) << 40 | //
+    (result & 0x0000000000FF0000) << 24 | //
+    (result & 0x00000000FF000000) << 8 |  //
+    (result & 0x000000FF00000000) >> 8 |  //
+    (result & 0x0000FF0000000000) >> 24 | //
+    (result & 0x00FF000000000000) >> 40 | //
+    (result & 0xFF00000000000000) >> 56;
+}
+
+auto
+decode_uint16(gsl::span<std::byte> data) -> std::uint16_t
+{
+  std::uint16_t result{};
+  std::memcpy(&result, data.data(), sizeof(result));
+
+  return static_cast<std::uint16_t>(result << 8 | result >> 8);
+}
+
+auto
+decode_frame(gsl::span<std::byte> data, bool expected_continuation) -> decode_status
+{
+  auto first_byte = static_cast<std::uint8_t>(data[0]);
+  if ((first_byte & reserved_bit_mask) != 0) {
+    return decoding_error{ "unsupported error: reserved bit used" };
+  }
+
+  const bool expect_more = (first_byte & flag_fin) == 0;
+  const std::uint8_t frame_type = first_byte & opcode_mask;
+
+  if (expect_more && is_control_frame(frame_type)) {
+    return decoding_error{ "unsupported error: fragmented control frame" };
+  }
+
+  if (is_data_frame(frame_type) && expected_continuation) {
+    return decoding_error{ "unsupported error: expected continuation frame" };
+  }
+
+  const auto second_byte = static_cast<std::uint8_t>(data[1]);
+
+  const bool masked = (second_byte & flag_mask) != 0;
+
+  const std::uint8_t length_7 = second_byte & payload_length_7_mask;
+  if (is_control_frame(frame_type) && length_7 > 125) {
+    return decoding_error{ "unsupported error: control frame is too long" };
+  }
+
+  std::size_t header_length{};
+  std::size_t payload_length{};
+
+  switch (length_7) {
+    case 127:
+      if (data.size() < 10) {
+        return need_more_data{};
+      }
+      header_length = 10;
+      payload_length = decode_uint64(data.subspan(2, sizeof(std::uint64_t)));
+      break;
+
+    case 126:
+      if (data.size() < 4) {
+        return need_more_data{};
+      }
+      header_length = 4;
+      payload_length = decode_uint16(data.subspan(2, sizeof(std::uint16_t)));
+      break;
+
+    default:
+      header_length = 2;
+      payload_length = length_7;
+  }
+
+  constexpr std::size_t masking_key_size{ 4 };
+  if (data.size() < header_length + payload_length + (masked ? masking_key_size : 0)) {
+    return need_more_data{};
+  }
+  auto consumed_bytes = header_length;
+
+  auto payload = data.subspan(header_length);
+  if (masked) {
+    auto masking_key = data.subspan(header_length, masking_key_size);
+    consumed_bytes += masking_key.size();
+    payload = data.subspan(header_length + masking_key.size());
+    mask_payload_data(masking_key, payload);
+  }
+  consumed_bytes += payload.size();
+
+  return decoded_frame{
+    frame_type,
+    payload,
+    consumed_bytes,
+    expect_more,
+  };
+}
+
+template<typename T>
+auto
+encode_frame(std::uint8_t opcode, T message) -> std::vector<std::byte>
+{
+  std::vector<std::byte> payload{};
+  payload.reserve(7 + message.size());
+  payload.emplace_back(static_cast<std::byte>(opcode | flag_fin));
+  auto encoded_length = encode_payload_length(message.size());
+  std::copy(encoded_length.begin(), encoded_length.end(), std::back_insert_iterator(payload));
+  auto masking_key = generate_masking_key();
+  std::copy(masking_key.begin(), masking_key.end(), std::back_insert_iterator(payload));
+  auto header_length = payload.size();
+  std::for_each(message.begin(), message.end(), [&payload](auto ch) {
+    payload.push_back(static_cast<std::byte>(ch));
+  });
+  mask_payload_data(masking_key,
+                    { payload.data() + header_length, payload.size() - header_length });
+  return payload;
+}
+} // namespace
+
+websocket_codec::websocket_codec(websocket_callbacks* callbacks)
+  : session_key_{ generate_session_key() }
+  , callbacks_{ callbacks }
+  , handler_{ std::make_unique<open_handshake>() }
+{
+}
+
+auto
+websocket_codec::session_key() const -> const std::string&
+{
+  return session_key_;
+}
+
+websocket_codec::~websocket_codec() = default;
+
+void
+websocket_codec::feed(gsl::span<std::byte> chunk)
+{
+  auto new_handler = handler_->feed(chunk,
+                                    context{
+                                      *callbacks_,
+                                      *this,
+                                    });
+  if (new_handler) {
+    std::swap(handler_, new_handler);
+  }
+}
+
+void
+websocket_codec::feed(std::string_view chunk)
+{
+  std::vector<std::byte> copy{ reinterpret_cast<const std::byte*>(chunk.data()),
+                               reinterpret_cast<const std::byte*>(chunk.data()) + chunk.size() };
+  feed(copy);
+}
+
+auto
+websocket_codec::text(std::string_view message) const -> std::vector<std::byte>
+{
+  return encode_frame(opcode_text, message);
+}
+
+auto
+websocket_codec::binary(gsl::span<std::byte> message) const -> std::vector<std::byte>
+{
+  return encode_frame(opcode_binary, message);
+}
+
+auto
+websocket_codec::ping(gsl::span<std::byte> message) const -> std::vector<std::byte>
+{
+  return encode_frame(opcode_ping, message);
+}
+
+auto
+websocket_codec::pong(gsl::span<std::byte> message) const -> std::vector<std::byte>
+{
+  return encode_frame(opcode_pong, message);
+}
+
+auto
+websocket_codec::close(gsl::span<std::byte> message) const -> std::vector<std::byte>
+{
+  return encode_frame(opcode_close, message);
+}
+} // namespace couchbase::core