FIX: Fix potential unaligned records

Author: threecgreen

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## 0.34.2 - TBD
+
+#### Bug fixes
+- Fixed potential for unaligned records in live and historical streaming requests
+
 ## 0.34.1 - 2025-04-29
 
 ### Enhancements

include/databento/dbn_decoder.hpp

@@ -6,6 +6,7 @@
 #include <string>
 
 #include "databento/dbn.hpp"
+#include "databento/detail/buffer.hpp"
 #include "databento/enums.hpp"  // Upgrade Policy
 #include "databento/file_stream.hpp"
 #include "databento/ireadable.hpp"
@@ -54,16 +55,14 @@ class DbnDecoder {
                                            const std::byte* buffer_end);
   bool DetectCompression();
   std::size_t FillBuffer();
-  std::size_t GetReadBufferSize() const;
   RecordHeader* BufferRecordHeader();
 
   ILogReceiver* log_receiver_;
   std::uint8_t version_{};
   VersionUpgradePolicy upgrade_policy_;
   bool ts_out_{};
   std::unique_ptr<IReadable> input_;
-  std::vector<std::byte> read_buffer_;
-  std::size_t buffer_idx_{};
+  detail::Buffer buffer_{};
   // Must be 8-byte aligned for records
   alignas(RecordHeader) std::array<std::byte, kMaxRecordLen> compat_buffer_{};
   Record current_record_{nullptr};

include/databento/detail/buffer.hpp

@@ -2,6 +2,7 @@
 
 #include <cstddef>
 #include <memory>
+#include <new>
 
 #include "databento/ireadable.hpp"
 #include "databento/iwritable.hpp"
@@ -11,7 +12,7 @@ class Buffer : public IReadable, public IWritable {
  public:
   Buffer() : Buffer(64 * std::size_t{1 << 10}) {}
   explicit Buffer(std::size_t init_capacity)
-      : buf_{std::make_unique<std::byte[]>(init_capacity)},
+      : buf_{AlignedNew(init_capacity), AlignedDelete},
         end_{buf_.get() + init_capacity},
         read_pos_{buf_.get()},
         write_pos_{buf_.get()} {}
@@ -22,7 +23,9 @@ class Buffer : public IReadable, public IWritable {
   void WriteAll(const std::byte* data, std::size_t length) override;
 
   std::byte*& WriteBegin() { return write_pos_; }
-  std::byte* WriteEnd() const { return end_; }
+  std::byte* WriteEnd() { return end_; }
+  const std::byte* WriteBegin() const { return write_pos_; }
+  const std::byte* WriteEnd() const { return end_; }
   std::size_t WriteCapacity() const {
     return static_cast<std::size_t>(end_ - write_pos_);
   }
@@ -32,7 +35,9 @@ class Buffer : public IReadable, public IWritable {
   std::size_t ReadSome(std::byte* buffer, std::size_t max_length) override;
 
   std::byte*& ReadBegin() { return read_pos_; }
-  std::byte* ReadEnd() const { return write_pos_; }
+  std::byte* ReadEnd() { return write_pos_; }
+  const std::byte* ReadBegin() const { return read_pos_; }
+  const std::byte* ReadEnd() const { return write_pos_; }
   std::size_t ReadCapacity() const {
     return static_cast<std::size_t>(write_pos_ - read_pos_);
   }
@@ -48,7 +53,16 @@ class Buffer : public IReadable, public IWritable {
   void Shift();
 
  private:
-  std::unique_ptr<std::byte[]> buf_;
+  static constexpr std::align_val_t kAlignment{8};
+
+  using UniqueBufPtr = std::unique_ptr<std::byte[], void (*)(std::byte*)>;
+
+  std::byte* AlignedNew(std::size_t capacity) {
+    return new (kAlignment) std::byte[capacity];
+  }
+  static void AlignedDelete(std::byte* p) { operator delete[](p, kAlignment); }
+
+  UniqueBufPtr buf_;
   std::byte* end_;
   std::byte* read_pos_{};
   std::byte* write_pos_{};

include/databento/detail/zstd_stream.hpp

@@ -6,6 +6,7 @@
 #include <memory>   // unique_ptr
 #include <vector>
 
+#include "databento/detail/buffer.hpp"
 #include "databento/ireadable.hpp"
 #include "databento/iwritable.hpp"
 #include "databento/log.hpp"
@@ -14,8 +15,7 @@ namespace databento::detail {
 class ZstdDecodeStream : public IReadable {
  public:
   explicit ZstdDecodeStream(std::unique_ptr<IReadable> input);
-  ZstdDecodeStream(std::unique_ptr<IReadable> input,
-                   std::vector<std::byte>&& in_buffer);
+  ZstdDecodeStream(std::unique_ptr<IReadable> input, detail::Buffer& in_buffer);
 
   // Read exactly `length` bytes into `buffer`.
   void ReadExact(std::byte* buffer, std::size_t length) override;

src/dbn_constants.hpp

@@ -13,7 +13,6 @@ constexpr std::size_t kFixedMetadataLen = 100;
 constexpr std::size_t kDatasetCstrLen = 16;
 constexpr std::size_t kMetadataReservedLen = 53;
 constexpr std::size_t kMetadataReservedLenV1 = 47;
-constexpr std::size_t kBufferCapacity = 8UL * 1024;
 constexpr std::uint16_t kNullSchema = std::numeric_limits<std::uint16_t>::max();
 constexpr std::uint8_t kNullSType = std::numeric_limits<std::uint8_t>::max();
 constexpr std::uint64_t kNullRecordCount =

src/dbn_decoder.cpp

@@ -9,6 +9,7 @@
 #include "databento/compat.hpp"
 #include "databento/constants.hpp"
 #include "databento/datetime.hpp"
+#include "databento/detail/buffer.hpp"
 #include "databento/detail/zstd_stream.hpp"
 #include "databento/enums.hpp"
 #include "databento/exceptions.hpp"
@@ -76,16 +77,12 @@ DbnDecoder::DbnDecoder(ILogReceiver* log_receiver,
     : log_receiver_{log_receiver},
       upgrade_policy_{upgrade_policy},
       input_{std::move(input)} {
-  read_buffer_.reserve(kBufferCapacity);
   if (DetectCompression()) {
-    input_ = std::make_unique<detail::ZstdDecodeStream>(
-        std::move(input_), std::move(read_buffer_));
-    // Reinitialize buffer and get it into the same state as uncompressed input
-    read_buffer_ = std::vector<std::byte>();
-    read_buffer_.reserve(kBufferCapacity);
-    read_buffer_.resize(kMagicSize);
-    input_->ReadExact(read_buffer_.data(), kMagicSize);
-    const auto* buf_ptr = read_buffer_.data();
+    input_ =
+        std::make_unique<detail::ZstdDecodeStream>(std::move(input_), buffer_);
+    input_->ReadExact(buffer_.WriteBegin(), kMagicSize);
+    buffer_.WriteBegin() += kMagicSize;
+    const auto* buf_ptr = buffer_.ReadBegin();
     if (std::strncmp(Consume(buf_ptr, 3), kDbnPrefix, 3) != 0) {
       throw DbnResponseError{"Found Zstd input, but not DBN prefix"};
     }
@@ -181,16 +178,22 @@ databento::Metadata DbnDecoder::DecodeMetadataFields(
 
 databento::Metadata DbnDecoder::DecodeMetadata() {
   // already read first 4 bytes detecting compression
-  read_buffer_.resize(kMetadataPreludeSize);
-  input_->ReadExact(&read_buffer_[4], 4);
+  const auto read_size = kMetadataPreludeSize - kMagicSize;
+  input_->ReadExact(buffer_.WriteBegin(), read_size);
+  buffer_.WriteBegin() += read_size;
   const auto [version, size] = DbnDecoder::DecodeMetadataVersionAndSize(
-      read_buffer_.data(), kMetadataPreludeSize);
+      buffer_.ReadBegin(), kMetadataPreludeSize);
+  buffer_.ReadBegin() += kMetadataPreludeSize;
   version_ = version;
-  read_buffer_.resize(size);
-  input_->ReadExact(read_buffer_.data(), read_buffer_.size());
-  buffer_idx_ = read_buffer_.size();
+  buffer_.Reserve(size);
+  input_->ReadExact(buffer_.WriteBegin(), size);
+  buffer_.WriteBegin() += size;
   auto metadata = DbnDecoder::DecodeMetadataFields(
-      version_, read_buffer_.data(), read_buffer_.data() + read_buffer_.size());
+      version_, buffer_.ReadBegin(), buffer_.ReadEnd());
+  buffer_.ReadBegin() += size;
+  // Metadata may leave buffer misaligned. Shift records to ensure 8-byte
+  // alignment
+  buffer_.Shift();
   ts_out_ = metadata.ts_out;
   metadata.Upgrade(upgrade_policy_);
   return metadata;
@@ -239,60 +242,53 @@ databento::Record DbnDecoder::DecodeRecordCompat(
 // assumes DecodeMetadata has been called
 const databento::Record* DbnDecoder::DecodeRecord() {
   // need some unread unread_bytes
-  if (GetReadBufferSize() == 0) {
+  if (buffer_.ReadCapacity() == 0) {
     if (FillBuffer() == 0) {
       return nullptr;
     }
   }
   // check length
-  while (GetReadBufferSize() < BufferRecordHeader()->Size()) {
+  while (buffer_.ReadCapacity() < BufferRecordHeader()->Size()) {
     if (FillBuffer() == 0) {
-      if (GetReadBufferSize() > 0) {
+      if (buffer_.ReadCapacity() > 0) {
         log_receiver_->Receive(
             LogLevel::Warning,
             "Unexpected partial record remaining in stream: " +
-                std::to_string(GetReadBufferSize()) + " bytes");
+                std::to_string(buffer_.ReadCapacity()) + " bytes");
       }
       return nullptr;
     }
   }
   current_record_ = Record{BufferRecordHeader()};
-  buffer_idx_ += current_record_.Size();
+  buffer_.ReadBegin() += current_record_.Size();
   current_record_ = DbnDecoder::DecodeRecordCompat(
       version_, upgrade_policy_, ts_out_, &compat_buffer_, current_record_);
   return &current_record_;
 }
 
 size_t DbnDecoder::FillBuffer() {
-  // Shift data forward
-  std::copy(read_buffer_.cbegin() + static_cast<std::ptrdiff_t>(buffer_idx_),
-            read_buffer_.cend(), read_buffer_.begin());
-  const auto unread_size = read_buffer_.size() - buffer_idx_;
-  buffer_idx_ = 0;
-  read_buffer_.resize(kBufferCapacity);
-  const auto fill_size = input_->ReadSome(&read_buffer_[unread_size],
-                                          kBufferCapacity - unread_size);
-  read_buffer_.resize(unread_size + fill_size);
+  if (buffer_.WriteCapacity() < kMaxRecordLen) {
+    buffer_.Shift();
+  }
+  const auto fill_size =
+      input_->ReadSome(buffer_.WriteBegin(), buffer_.WriteCapacity());
+  buffer_.WriteBegin() += fill_size;
   return fill_size;
 }
 
-std::size_t DbnDecoder::GetReadBufferSize() const {
-  return read_buffer_.size() - buffer_idx_;
-}
-
 databento::RecordHeader* DbnDecoder::BufferRecordHeader() {
-  return reinterpret_cast<RecordHeader*>(&read_buffer_[buffer_idx_]);
+  return reinterpret_cast<RecordHeader*>(buffer_.ReadBegin());
 }
 
 bool DbnDecoder::DetectCompression() {
-  read_buffer_.resize(kMagicSize);
-  input_->ReadExact(read_buffer_.data(), kMagicSize);
-  const auto* read_buffer_it = read_buffer_.data();
-  if (std::strncmp(Consume(read_buffer_it, 3), kDbnPrefix, 3) == 0) {
+  input_->ReadExact(buffer_.WriteBegin(), kMagicSize);
+  buffer_.WriteBegin() += kMagicSize;
+  const auto* buffer_it = buffer_.ReadBegin();
+  if (std::strncmp(Consume(buffer_it, 3), kDbnPrefix, 3) == 0) {
     return false;
   }
-  read_buffer_it = read_buffer_.data();
-  auto x = Consume<std::uint32_t>(read_buffer_it);
+  buffer_it = buffer_.ReadBegin();
+  auto x = Consume<std::uint32_t>(buffer_it);
   if (x == kZstdMagicNumber) {
     return true;
   }
@@ -302,12 +298,10 @@ bool DbnDecoder::DetectCompression() {
   if ((x & kZstdSkippableFrame) == kZstdSkippableFrame) {
     throw DbnResponseError{
         "Legacy DBZ encoding is not supported. Please use the dbn CLI tool "
-        "to "
-        "convert it to DBN."};
+        "to convert it to DBN."};
   }
   throw DbnResponseError{
-      "Couldn't detect input type. It doesn't appear to be Zstd or "
-      "DBN."};
+      "Couldn't detect input type. It doesn't appear to be Zstd or DBN."};
 }
 
 std::string DbnDecoder::DecodeSymbol(std::size_t symbol_cstr_len,

src/detail/buffer.cpp

@@ -54,7 +54,7 @@ void Buffer::Reserve(std::size_t capacity) {
   if (capacity <= Capacity()) {
     return;
   }
-  auto new_buf = std::make_unique<std::byte[]>(capacity);
+  UniqueBufPtr new_buf{AlignedNew(capacity), AlignedDelete};
   const auto unread_bytes = ReadCapacity();
   std::copy(ReadBegin(), ReadEnd(), new_buf.get());
   buf_ = std::move(new_buf);

src/detail/dbn_buffer_decoder.cpp

@@ -35,6 +35,9 @@ databento::KeepGoing DbnBufferDecoder::Process(const char* data,
       auto metadata = DbnDecoder::DecodeMetadataFields(
           input_version_, dbn_buffer_.ReadBegin(), dbn_buffer_.ReadEnd());
       dbn_buffer_.ReadBegin() += bytes_needed_;
+      // Metadata may leave buffer misaligned. Shift records to ensure 8-byte
+      // alignment
+      dbn_buffer_.Shift();
       ts_out_ = metadata.ts_out;
       metadata.Upgrade(VersionUpgradePolicy::UpgradeToV2);
       if (metadata_callback_) {

src/detail/zstd_stream.cpp

@@ -4,21 +4,28 @@
 #include <sstream>
 #include <utility>  // move
 
+#include "databento/detail/buffer.hpp"
 #include "databento/exceptions.hpp"
 #include "databento/log.hpp"
 
 using databento::detail::ZstdDecodeStream;
 
 ZstdDecodeStream::ZstdDecodeStream(std::unique_ptr<IReadable> input)
-    : ZstdDecodeStream{std::move(input), {}} {}
+    : input_{std::move(input)},
+      z_dstream_{::ZSTD_createDStream(), ::ZSTD_freeDStream},
+      read_suggestion_{::ZSTD_initDStream(z_dstream_.get())},
+      in_buffer_{},
+      z_in_buffer_{in_buffer_.data(), 0, 0} {}
 
 ZstdDecodeStream::ZstdDecodeStream(std::unique_ptr<IReadable> input,
-                                   std::vector<std::byte>&& in_buffer)
+                                   detail::Buffer& in_buffer)
     : input_{std::move(input)},
       z_dstream_{::ZSTD_createDStream(), ::ZSTD_freeDStream},
       read_suggestion_{::ZSTD_initDStream(z_dstream_.get())},
-      in_buffer_{std::move(in_buffer)},
-      z_in_buffer_{in_buffer_.data(), in_buffer_.size(), 0} {}
+      in_buffer_{in_buffer.ReadBegin(), in_buffer.ReadEnd()},
+      z_in_buffer_{in_buffer_.data(), in_buffer_.size(), 0} {
+  in_buffer.ReadBegin() += in_buffer.ReadCapacity();
+}
 
 void ZstdDecodeStream::ReadExact(std::byte* buffer, std::size_t length) {
   std::size_t size{};

src/live_blocking.cpp

@@ -148,6 +148,9 @@ databento::Metadata LiveBlocking::Start() {
   auto metadata = DbnDecoder::DecodeMetadataFields(version, buffer_.ReadBegin(),
                                                    buffer_.ReadEnd());
   buffer_.ReadBegin() += size;
+  // Metadata may leave buffer misaligned. Shift records to ensure 8-byte
+  // alignment
+  buffer_.Shift();
   version_ = metadata.version;
   metadata.Upgrade(upgrade_policy_);
   return metadata;