EXP: [C++] Make IPC message decoding stricter

Error out when the advertised message metadata length is larger than expected.

Author: pitrou

cpp/src/arrow/buffer.cc

@@ -87,7 +87,7 @@ Result<std::shared_ptr<Buffer>> SliceMutableBufferSafe(std::shared_ptr<Buffer> b
   return SliceMutableBuffer(std::move(buffer), offset, length);
 }
 
-std::string Buffer::ToHexString() {
+std::string Buffer::ToHexString() const {
   return HexEncode(data(), static_cast<size_t>(size()));
 }
 

cpp/src/arrow/buffer.h

@@ -121,7 +121,7 @@ class ARROW_EXPORT Buffer {
 
   /// \brief Construct a new std::string with a hexadecimal representation of the buffer.
   /// \return std::string
-  std::string ToHexString();
+  std::string ToHexString() const;
 
   /// Return true if both buffers are the same size and contain the same bytes
   /// up to the number of compared bytes

cpp/src/arrow/ipc/message.cc

@@ -285,17 +285,21 @@ std::string FormatMessageType(MessageType type) {
 
 namespace {
 
-Status ReadFieldsSubset(int64_t offset, int32_t metadata_length,
-                        io::RandomAccessFile* file,
+Status ReadFieldsSubset(int64_t offset, io::RandomAccessFile* file,
                         const FieldsLoaderFunction& fields_loader,
                         const std::shared_ptr<Buffer>& metadata, int64_t required_size,
                         std::shared_ptr<Buffer>& body) {
+  DCHECK_GE(static_cast<size_t>(metadata->size()), sizeof(int32_t));
+  const auto continuation = util::SafeLoadAs<int32_t>(metadata->data());
+  // Either 8 bytes (32-bit continuation indicator + 32-bit little-endian length prefix)
+  // or 4 bytes for legacy IPC without continuation indicator
+  const auto continuation_size = (continuation == internal::kIpcContinuationToken)
+                                     ? 2 * sizeof(int32_t)
+                                     : sizeof(int32_t);
+
   const flatbuf::Message* message = nullptr;
-  uint8_t continuation_metadata_size = sizeof(int32_t) + sizeof(int32_t);
-  // skip 8 bytes (32-bit continuation indicator + 32-bit little-endian length prefix)
-  RETURN_NOT_OK(internal::VerifyMessage(metadata->data() + continuation_metadata_size,
-                                        metadata->size() - continuation_metadata_size,
-                                        &message));
+  RETURN_NOT_OK(internal::VerifyMessage(metadata->data() + continuation_size,
+                                        metadata->size() - continuation_size, &message));
   auto batch = message->header_as_RecordBatch();
   if (batch == nullptr) {
     return Status::IOError(
@@ -305,8 +309,8 @@ Status ReadFieldsSubset(int64_t offset, int32_t metadata_length,
   RETURN_NOT_OK(fields_loader(batch, &io_recorded_random_access_file));
   const auto& read_ranges = io_recorded_random_access_file.GetReadRanges();
   for (const auto& range : read_ranges) {
-    auto read_result = file->ReadAt(offset + metadata_length + range.offset, range.length,
-                                    body->mutable_data() + range.offset);
+    auto read_result = file->ReadAt(offset + metadata->size() + range.offset,
+                                    range.length, body->mutable_data() + range.offset);
     if (!read_result.ok()) {
       return Status::IOError("Failed to read message body, error ",
                              read_result.status().ToString());
@@ -330,6 +334,10 @@ Result<std::unique_ptr<Message>> ReadMessage(std::shared_ptr<Buffer> metadata,
   }
 
   ARROW_RETURN_NOT_OK(decoder.Consume(metadata));
+  if (decoder.buffered_size() > 0) {
+    return Status::Invalid("Message metadata too long by ", decoder.buffered_size(),
+                           " bytes");
+  }
 
   switch (decoder.state()) {
     case MessageDecoder::State::INITIAL:
@@ -384,6 +392,10 @@ static Result<std::unique_ptr<Message>> ReadMessageInternal(
                                      /*allow_short_read=*/false));
 
   ARROW_RETURN_NOT_OK(decoder.Consume(SliceBuffer(metadata, 0, metadata_length)));
+  if (decoder.buffered_size() > 0) {
+    return Status::Invalid("Message metadata too long by ", decoder.buffered_size(),
+                           " bytes");
+  }
 
   switch (decoder.state()) {
     case MessageDecoder::State::INITIAL:
@@ -402,7 +414,7 @@ static Result<std::unique_ptr<Message>> ReadMessageInternal(
         // requested field ranges into it.
         ARROW_ASSIGN_OR_RAISE(
             body, AllocateBuffer(decoder.next_required_size(), default_memory_pool()));
-        RETURN_NOT_OK(ReadFieldsSubset(offset, metadata_length, file, fields_loader,
+        RETURN_NOT_OK(ReadFieldsSubset(offset, file, fields_loader,
                                        SliceBuffer(metadata, 0, metadata_length),
                                        decoder.next_required_size(), body));
       } else if (body_length.has_value()) {
@@ -469,9 +481,14 @@ Future<std::shared_ptr<Message>> ReadMessageAsync(int64_t offset, int32_t metada
       ->ReadAsync(context, offset, metadata_length + body_length,
                   /*allow_short_read=*/false)
       .Then([=](std::shared_ptr<Buffer> metadata) -> Result<std::shared_ptr<Message>> {
+        // TODO reconcile + factor out with blocking ReadMessage logic
         DCHECK_EQ(metadata->size(), metadata_length + body_length);
         ARROW_RETURN_NOT_OK(
             state->decoder->Consume(SliceBuffer(metadata, 0, metadata_length)));
+        if (state->decoder->buffered_size() > 0) {
+          return Status::Invalid("Message metadata too long by ",
+                                 state->decoder->buffered_size(), " bytes");
+        }
         switch (state->decoder->state()) {
           case MessageDecoder::State::INITIAL:
             return std::move(state->result);
@@ -484,12 +501,15 @@ Future<std::shared_ptr<Message>> ReadMessageAsync(int64_t offset, int32_t metada
                                    " invalid. File offset: ", offset,
                                    ", metadata length: ", metadata_length);
           case MessageDecoder::State::BODY: {
-            auto body = SliceBuffer(metadata, metadata_length, body_length);
-            if (body->size() < state->decoder->next_required_size()) {
-              return Status::IOError("Expected to be able to read ",
-                                     state->decoder->next_required_size(),
-                                     " bytes for message body, got ", body->size());
+            if (body_length != state->decoder->next_required_size()) {
+              // The streaming decoder got out of sync with the actual advertised
+              // metadata and body size, which signals an invalid IPC file.
+              return Status::IOError(
+                  "Invalid IPC file: advertised body size is ", body_length,
+                  ", but message decoder expects to read ",
+                  state->decoder->next_required_size(), " bytes instead");
             }
+            auto body = SliceBuffer(metadata, metadata_length, body_length);
             RETURN_NOT_OK(state->decoder->Consume(body));
             return std::move(state->result);
           }
@@ -755,6 +775,8 @@ class MessageDecoder::MessageDecoderImpl {
 
   int64_t next_required_size() const { return next_required_size_ - buffered_size_; }
 
+  int64_t buffered_size() const { return buffered_size_; }
+
   MessageDecoder::State state() const { return state_; }
 
  private:
@@ -1020,6 +1042,8 @@ Status MessageDecoder::Consume(std::shared_ptr<Buffer> buffer) {
 
 int64_t MessageDecoder::next_required_size() const { return impl_->next_required_size(); }
 
+int64_t MessageDecoder::buffered_size() const { return impl_->buffered_size(); }
+
 MessageDecoder::State MessageDecoder::state() const { return impl_->state(); }
 
 // ----------------------------------------------------------------------

cpp/src/arrow/ipc/message.h

@@ -419,6 +419,11 @@ class ARROW_EXPORT MessageDecoder {
   /// \return the current state
   State state() const;
 
+  /// \brief Return the number of bytes buffered in the decoder.
+  ///
+  /// This method is mainly useful for testing and debugging.
+  int64_t buffered_size() const;
+
  private:
   class MessageDecoderImpl;
   std::unique_ptr<MessageDecoderImpl> impl_;
@@ -453,9 +458,12 @@ using FieldsLoaderFunction = std::function<Status(const void*, io::RandomAccessF
 ///
 /// Read a length-prefixed message flatbuffer starting at the indicated file
 /// offset. If the message has a body with non-zero length, it will also be
-/// read
+/// read.
 ///
-/// The metadata_length includes at least the length prefix and the flatbuffer
+/// The metadata_length includes the IPC encapsulation prefix and the
+/// Flatbuffers-serialized message.
+///
+/// This function should only be used when a RecordBatch message is expected.
 ///
 /// \param[in] offset the position in the file where the message starts. The
 /// first 4 bytes after the offset are the message length
@@ -474,7 +482,8 @@ Result<std::unique_ptr<Message>> ReadMessage(
 /// Read a length-prefixed message flatbuffer starting at the indicated file
 /// offset.
 ///
-/// The metadata_length includes at least the length prefix and the flatbuffer
+/// The metadata_length includes the IPC encapsulation prefix and the
+/// Flatbuffers-serialized message.
 ///
 /// \param[in] offset the position in the file where the message starts. The
 /// first 4 bytes after the offset are the message length
@@ -578,6 +587,7 @@ Status DecodeMessage(MessageDecoder* decoder, io::InputStream* stream);
 /// \param[out] message_length the total size of the payload written including
 /// padding
 /// \return Status
+ARROW_EXPORT
 Status WriteMessage(const Buffer& message, const IpcWriteOptions& options,
                     io::OutputStream* file, int32_t* message_length);