couchbase 3.8.2 → 3.9.0

data/ext/couchbase/core/io/mcbp_output_queue.hxx

@@ -0,0 +1,150 @@
+/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
+/*
+ *   Copyright 2024. 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
+ *
+ *       http://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.
+ */
+
+#pragma once
+
+#include <cstddef>
+#include <mutex>
+#include <utility>
+#include <vector>
+
+namespace couchbase::core::io
+{
+/**
+ * Send-side buffering for an mcbp connection.
+ *
+ * A KV operation appends its encoded frame from the caller's thread while the connection's write
+ * loop (`do_write` / the async_write completion) runs on an IO thread; this queue owns both buffers
+ * and the single flag that coordinates them, so all of that state is mutated under one lock.
+ *
+ * The flag exists to eliminate a per-operation `asio::post`. `enqueue()`/`mark_for_dispatch()` only
+ * ask the caller to schedule a write on the idle -> scheduled transition; while a write is already
+ * scheduled or in flight they return false, because the in-flight write's completion is what
+ * re-drives the loop. Draining the queue to empty returns it to idle so the next enqueue re-arms.
+ */
+class mcbp_output_queue
+{
+public:
+  using buffer = std::vector<std::byte>;
+
+  /**
+   * Append a buffer and stage it for the next write.
+   *
+   * @return true if the caller must schedule a write (the queue transitioned from idle to
+   * scheduled); false if a write is already scheduled or in flight.
+   */
+  [[nodiscard]] auto enqueue(buffer&& buf) -> bool
+  {
+    const std::scoped_lock lock(mutex_);
+    output_.emplace_back(std::move(buf));
+    return arm();
+  }
+
+  /**
+   * Append a buffer without requesting a write. Used to refill the queue before a single
+   * mark_for_dispatch(), e.g. when draining the bootstrap pending buffer.
+   */
+  void stage(buffer&& buf)
+  {
+    const std::scoped_lock lock(mutex_);
+    output_.emplace_back(std::move(buf));
+  }
+
+  /**
+   * Request a write for data staged earlier.
+   *
+   * @return true if the caller must schedule a write; false if nothing is queued or a write is
+   * already scheduled or in flight.
+   */
+  [[nodiscard]] auto mark_for_dispatch() -> bool
+  {
+    const std::scoped_lock lock(mutex_);
+    if (output_.empty()) {
+      return false;
+    }
+    return arm();
+  }
+
+  /**
+   * Move all queued buffers into the writing batch so it can be handed to async_write.
+   *
+   * @return true if there is a batch to send (see writing()); false if nothing is queued (the queue
+   * is returned to idle) or a write is already in flight (left untouched).
+   */
+  [[nodiscard]] auto begin_writing() -> bool
+  {
+    const std::scoped_lock lock(mutex_);
+    if (!writing_.empty()) {
+      // A batch is still in flight; issuing another async_write concurrently would corrupt the
+      // stream. Leave it scheduled — the in-flight completion will re-drive us.
+      return false;
+    }
+    if (output_.empty()) {
+      scheduled_ = false;
+      return false;
+    }
+    std::swap(writing_, output_);
+    return true;
+  }
+
+  /**
+   * The batch currently being written. Only valid between a begin_writing() that returned true and
+   * the matching finish_writing(); accessed solely on the IO thread during that window.
+   */
+  [[nodiscard]] auto writing() const -> const std::vector<buffer>&
+  {
+    return writing_;
+  }
+
+  /**
+   * Release the batch handed out by begin_writing() once its async_write has completed.
+   */
+  void finish_writing()
+  {
+    const std::scoped_lock lock(mutex_);
+    writing_.clear();
+  }
+
+  /**
+   * Discard everything and return to idle (used when the connection is reset or stopped).
+   */
+  void reset()
+  {
+    const std::scoped_lock lock(mutex_);
+    output_.clear();
+    writing_.clear();
+    scheduled_ = false;
+  }
+
+private:
+  // Precondition: mutex_ held. Transition idle -> scheduled, reporting whether the caller owns the
+  // resulting write post.
+  auto arm() -> bool
+  {
+    if (scheduled_) {
+      return false;
+    }
+    scheduled_ = true;
+    return true;
+  }
+
+  std::vector<buffer> output_{};
+  std::vector<buffer> writing_{};
+  bool scheduled_{ false };
+  std::mutex mutex_{};
+};
+} // namespace couchbase::core::io

data/ext/couchbase/core/io/mcbp_parser.cxx

@@ -38,25 +38,42 @@ mcbp_parser::next(mcbp_message& msg) -> mcbp_parser::result
     return result::need_data;
   }
   std::memcpy(&msg.header, buf.data(), header_size);
-  std::uint32_t body_size = utils::byte_swap(msg.header.bodylen);
+  std::uint32_t body_size = utils::network_to_host(msg.header.bodylen);
   if (body_size > 0 && buf.size() - header_size < body_size) {
     return result::need_data;
   }
   msg.body.clear();
   msg.body.reserve(body_size);
-  std::uint32_t key_size = utils::byte_swap(msg.header.keylen);
+  std::uint32_t key_size = utils::network_to_host(msg.header.keylen);
   std::uint32_t prefix_size = static_cast<std::uint32_t>(msg.header.extlen) + key_size;
   if (msg.header.magic == static_cast<std::uint8_t>(protocol::magic::alt_client_response)) {
-    const std::uint8_t framing_extras_size = msg.header.keylen & 0xffU;
-    key_size = static_cast<std::uint32_t>(msg.header.keylen >> 8U);
-    prefix_size = static_cast<std::uint32_t>(framing_extras_size) +
+    const auto sizes = msg.header.alt_sizes();
+    key_size = sizes.key;
+    prefix_size = static_cast<std::uint32_t>(sizes.framing_extras) +
                   static_cast<std::uint32_t>(msg.header.extlen) + key_size;
   }
+  // The prefix (framing extras + extras + key) must lie within the body that
+  // bodylen advertised. extlen and keylen are separate header fields from
+  // bodylen, so a frame can claim prefix_size > body_size. Only body_size was
+  // checked against the buffer above; without this guard the inserts below read
+  // past the buffer and "body_size - prefix_size" underflows on the snappy path.
+  if (prefix_size > body_size) {
+    CB_LOG_WARNING("rejecting malformed frame: prefix_size ({}) exceeds body_size ({}), "
+                   "magic={:x}, opcode={:x}, extlen={}, keylen={}",
+                   prefix_size,
+                   body_size,
+                   msg.header.magic,
+                   msg.header.opcode,
+                   msg.header.extlen,
+                   key_size);
+    reset();
+    return result::failure;
+  }
   msg.body.insert(
     msg.body.end(), buf.begin() + header_size, buf.begin() + header_size + prefix_size);
 
   const bool is_compressed =
-    (msg.header.datatype & static_cast<std::uint8_t>(protocol::datatype::snappy)) != 0;
+    protocol::has_flag(static_cast<std::byte>(msg.header.datatype), protocol::datatype::snappy);
   bool use_raw_value = true;
   if (is_compressed) {
     std::string uncompressed;
@@ -70,7 +87,7 @@ mcbp_parser::next(mcbp_message& msg) -> mcbp_parser::result
       use_raw_value = false;
       // patch header with new body size
       msg.header.bodylen =
-        utils::byte_swap(static_cast<std::uint32_t>(prefix_size + uncompressed.size()));
+        utils::host_to_network(static_cast<std::uint32_t>(prefix_size + uncompressed.size()));
     }
   }
   if (use_raw_value) {