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

@@ -330,6 +330,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 +388,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:
@@ -469,9 +477,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 +497,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 +771,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 +1038,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

@@ -386,6 +386,9 @@ class ARROW_EXPORT MessageDecoder {
   /// decoder
   int64_t next_required_size() const;
 
+  // TODO document this
+  int64_t buffered_size() const;
+
   /// \brief Return the current state of the decoder.
   ///
   /// This method is provided for users who want to optimize performance.
@@ -578,6 +581,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);