GH-55428: [C++][S3] Preserve retryable errors when ranged-read buffers overflow

cpp/src/arrow/filesystem/s3_stream_internal.h

@@ -0,0 +1,169 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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 <algorithm>
+#include <cstring>
+#include <iostream>
+#include <limits>
+#include <memory>
+#include <streambuf>
+
+#include "arrow/status.h"
+
+namespace arrow {
+namespace fs {
+namespace internal {
+
+// The SDK writes error bodies through the GetObject response stream too. Keep the
+// caller's buffer as the put area for normal reads; capture a bounded prefix if
+// an error body exceeds the requested range.
+class GetObjectRangeStreamBuf : public std::streambuf {
+ public:
+  static constexpr size_t kSpillSize = 64 * 1024;
+
+  GetObjectRangeStreamBuf(void* data, size_t size) : data_(static_cast<char*>(data)) {
+    setp(data_, data_ + size);
+    setg(data_, data_, data_);
+  }
+
+  bool overflowed() const { return overflowed_; }
+
+ protected:
+  std::streamsize xsputn(const char* source, std::streamsize count) override {
+    if (count <= 0) {
+      return 0;
+    }
+    auto remaining = static_cast<size_t>(count);
+    while (remaining != 0) {
+      if (pptr() == epptr() && !spill_) {
+        Spill();
+      }
+      const auto n = std::min(remaining, static_cast<size_t>(epptr() - pptr()));
+      if (n == 0) {
+        break;
+      }
+      std::memcpy(pptr(), source, n);
+      Advance(n);
+      source += n;
+      remaining -= n;
+    }
+    // A short write makes the SDK lose the HTTP error and its retry decision.
+    return count;
+  }
+
+  int_type overflow(int_type ch) override {
+    if (traits_type::eq_int_type(ch, traits_type::eof())) {
+      return traits_type::not_eof(ch);
+    }
+    const char c = traits_type::to_char_type(ch);
+    xsputn(&c, 1);
+    return ch;
+  }
+
+  int_type underflow() override {
+    if (gptr() == egptr()) {
+      const auto pos = static_cast<size_t>(gptr() - eback());
+      setg(Storage(), Storage() + pos, Storage() + Written());
+    }
+    return gptr() == egptr() ? traits_type::eof() : traits_type::to_int_type(*gptr());
+  }
+
+  pos_type seekoff(off_type offset, std::ios_base::seekdir dir,
+                   std::ios_base::openmode which) override {
+    if (which == std::ios_base::out) {
+      return dir == std::ios_base::cur && offset == 0
+                 ? pos_type(static_cast<off_type>(Written()))
+                 : pos_type(off_type(-1));
+    }
+    if (which != std::ios_base::in) {
+      return pos_type(off_type(-1));
+    }
+    const auto size = static_cast<off_type>(Written());
+    const auto base = dir == std::ios_base::beg ? off_type(0)
+                      : dir == std::ios_base::cur
+                          ? static_cast<off_type>(gptr() - eback())
+                          : size;
+    if (dir != std::ios_base::beg && dir != std::ios_base::cur &&
+        dir != std::ios_base::end) {
+      return pos_type(off_type(-1));
+    }
+    if (offset < -base || offset > size - base) {
+      return pos_type(off_type(-1));
+    }
+    const auto pos = base + offset;
+    if (pos < 0 || pos > size) {
+      return pos_type(off_type(-1));
+    }
+    setg(Storage(), Storage() + pos, Storage() + size);
+    return pos_type(pos);
+  }
+
+  pos_type seekpos(pos_type pos, std::ios_base::openmode which) override {
+    return seekoff(static_cast<off_type>(pos), std::ios_base::beg, which);
+  }
+
+  int sync() override { return 0; }
+
+ private:
+  char* Storage() const { return spill_ ? spill_.get() : data_; }
+  size_t Written() const { return static_cast<size_t>(pptr() - Storage()); }
+
+  void Advance(size_t n) {
+    while (n != 0) {
+      const auto step = static_cast<int>(
+          std::min(n, static_cast<size_t>(std::numeric_limits<int>::max())));
+      pbump(step);
+      n -= step;
+    }
+  }
+
+  void Spill() {
+    overflowed_ = true;
+    const auto kept = std::min(Written(), kSpillSize);
+    const auto read_pos = std::min(static_cast<size_t>(gptr() - eback()), kept);
+    auto spill = std::make_unique<char[]>(kSpillSize);
+    if (kept != 0) {
+      std::memcpy(spill.get(), data_, kept);
+    }
+    spill_ = std::move(spill);
+    setp(spill_.get(), spill_.get() + kSpillSize);
+    Advance(kept);
+    setg(spill_.get(), spill_.get() + read_pos, spill_.get() + kept);
+  }
+
+  char* data_;
+  std::unique_ptr<char[]> spill_;
+  bool overflowed_ = false;
+};
+
+class GetObjectRangeStream : private GetObjectRangeStreamBuf, public std::iostream {
+ public:
+  GetObjectRangeStream(void* data, size_t size)
+      : GetObjectRangeStreamBuf(data, size), std::iostream(this) {}
+
+  Status CheckForOverflow() const {
+    if (overflowed()) {
+      return Status::IOError("GetObject response exceeded the requested range");
+    }
+    return Status::OK();
+  }
+};
+
+}  // namespace internal
+}  // namespace fs
+}  // namespace arrow

cpp/src/arrow/filesystem/s3fs.cc

@@ -103,6 +103,7 @@
 #include "arrow/filesystem/filesystem.h"
 #include "arrow/filesystem/path_util.h"
 #include "arrow/filesystem/s3_internal.h"
+#include "arrow/filesystem/s3_stream_internal.h"
 #include "arrow/filesystem/util_internal.h"
 #include "arrow/io/interfaces.h"
 #include "arrow/io/memory.h"
@@ -1327,7 +1328,10 @@ class StringViewStream : Aws::Utils::Stream::PreallocatedStreamBuf, public std::
 // See https://github.com/aws/aws-sdk-cpp/issues/64 for an alternative but
 // functionally similar recipe.
 Aws::IOStreamFactory AwsWriteableStreamFactory(void* data, int64_t nbytes) {
-  return [=]() { return Aws::New<StringViewStream>("", data, nbytes); };
+  return [=]() {
+    return Aws::New<internal::GetObjectRangeStream>("", data,
+                                                    static_cast<size_t>(nbytes));
+  };
 }
 
 Result<S3Model::GetObjectResult> GetObjectRange(Aws::S3::S3Client* client,
@@ -1341,7 +1345,11 @@ Result<S3Model::GetObjectResult> GetObjectRange(Aws::S3::S3Client* client,
   RETURN_NOT_OK(SetSSECustomerKey(&req, sse_customer_key));
   req.SetRange(ToAwsString(FormatRange(start, length)));
   req.SetResponseStreamFactory(AwsWriteableStreamFactory(out, length));
-  return OutcomeToResult("GetObject", client->GetObject(req));
+  ARROW_ASSIGN_OR_RAISE(auto result,
+                        OutcomeToResult("GetObject", client->GetObject(req)));
+  RETURN_NOT_OK(
+      static_cast<internal::GetObjectRangeStream&>(result.GetBody()).CheckForOverflow());
+  return std::move(result);
 }
 
 template <typename ObjectResult>

cpp/src/arrow/filesystem/s3fs_test.cc

@@ -16,7 +16,9 @@
 // under the License.
 
 #include <algorithm>
+#include <cstring>
 #include <exception>
+#include <iterator>
 #include <memory>
 #include <sstream>
 #include <string>
@@ -53,6 +55,7 @@
 
 #include "arrow/filesystem/filesystem.h"
 #include "arrow/filesystem/s3_internal.h"
+#include "arrow/filesystem/s3_stream_internal.h"
 #include "arrow/filesystem/s3_test_util.h"
 #include "arrow/filesystem/s3fs.h"
 #include "arrow/filesystem/test_util.h"
@@ -104,6 +107,140 @@ ::testing::Environment* minio_env =
 ::testing::Environment* minio_env_https =
     ::testing::AddGlobalTestEnvironment(new MinioTestEnvironment(/*enable_tls=*/true));
 
+namespace {
+
+std::string ReadRangeResponse(std::iostream& stream) {
+  return std::string(std::istreambuf_iterator<char>(stream),
+                     std::istreambuf_iterator<char>());
+}
+
+class DirectWriteRangeBuffer : public internal::GetObjectRangeStreamBuf {
+ public:
+  using internal::GetObjectRangeStreamBuf::GetObjectRangeStreamBuf;
+
+  void WriteDirectly(const char* data, size_t size) {
+    ASSERT_LE(size, static_cast<size_t>(epptr() - pptr()));
+    std::memcpy(pptr(), data, size);
+    pbump(static_cast<int>(size));
+  }
+};
+
+TEST(GetObjectRangeStreamTest, SuccessfulReadsUseCallerBuffer) {
+  std::string caller(9, '!');
+  internal::GetObjectRangeStream stream(caller.data(), 8);
+  stream.write("abcdefgh", 8);
+  stream.flush();
+
+  ASSERT_TRUE(stream.good());
+  ASSERT_OK(stream.CheckForOverflow());
+  EXPECT_EQ("abcdefgh!", caller);
+  EXPECT_EQ(std::streampos(8), stream.tellp());
+  EXPECT_EQ("abcdefgh", ReadRangeResponse(stream));
+
+  std::string short_caller(8, '!');
+  internal::GetObjectRangeStream short_stream(short_caller.data(), short_caller.size());
+  short_stream.write("abc", 3);
+  EXPECT_EQ("abc", ReadRangeResponse(short_stream));
+}
+
+TEST(GetObjectRangeStreamTest, ErrorBodyOverflowsInBulkAndSplitWrites) {
+  const std::string xml = "<Error><Code>SlowDown</Code><Message>retry</Message></Error>";
+  for (const bool split : {false, true}) {
+    std::string caller(12, '!');
+    internal::GetObjectRangeStream stream(caller.data(), caller.size());
+    if (split) {
+      for (size_t pos = 0; pos < xml.size(); pos += 7) {
+        stream.write(xml.data() + pos, std::min(size_t{7}, xml.size() - pos));
+      }
+    } else {
+      stream.write(xml.data(), xml.size());
+    }
+    stream.flush();
+
+    ASSERT_TRUE(stream.good());
+    EXPECT_EQ(std::streampos(xml.size()), stream.tellp());
+    EXPECT_TRUE(stream.CheckForOverflow().IsIOError());
+    EXPECT_EQ(xml.substr(0, caller.size()), caller);
+    EXPECT_EQ(xml, ReadRangeResponse(stream));
+    stream.clear();
+    stream.seekg(0);
+    EXPECT_EQ(xml, ReadRangeResponse(stream));
+  }
+}
+
+TEST(GetObjectRangeStreamTest, OversizedSuccessfulResponseIsRejected) {
+  std::string caller(4, '!');
+  internal::GetObjectRangeStream stream(caller.data(), caller.size());
+  stream.write("abcdef", 6);
+
+  ASSERT_TRUE(stream.good());
+  EXPECT_EQ("abcd", caller);
+  EXPECT_TRUE(stream.CheckForOverflow().IsIOError());
+}
+
+TEST(GetObjectRangeStreamTest, CharacterWritesAndDirectPutArea) {
+  const std::string body = "abcdefghijklmnop";
+  std::string caller(4, '!');
+  internal::GetObjectRangeStream stream(caller.data(), caller.size());
+  for (char c : body) {
+    stream.put(c);
+  }
+  ASSERT_TRUE(stream.good());
+  EXPECT_EQ(body, ReadRangeResponse(stream));
+
+  std::string direct_caller(4, '!');
+  DirectWriteRangeBuffer buffer(direct_caller.data(), direct_caller.size());
+  buffer.WriteDirectly(body.data(), direct_caller.size());
+  EXPECT_EQ(std::streampos(direct_caller.size()),
+            buffer.pubseekoff(0, std::ios_base::cur, std::ios_base::out));
+  EXPECT_EQ(static_cast<std::streamsize>(body.size() - direct_caller.size()),
+            buffer.sputn(body.data() + direct_caller.size(),
+                         body.size() - direct_caller.size()));
+  std::istream input(&buffer);
+  EXPECT_TRUE(buffer.overflowed());
+  EXPECT_EQ(body, std::string(std::istreambuf_iterator<char>(input),
+                              std::istreambuf_iterator<char>()));
+}
+
+TEST(GetObjectRangeStreamTest, CapsAnOversizedErrorBody) {
+  const auto cap = internal::GetObjectRangeStreamBuf::kSpillSize;
+  for (const size_t body_size : {cap, cap + 1}) {
+    const std::string body(body_size, 'x');
+    std::string caller(8, '!');
+    internal::GetObjectRangeStream stream(caller.data(), caller.size());
+    stream.write(body.data(), body.size());
+    stream.flush();
+
+    ASSERT_TRUE(stream.good());
+    EXPECT_TRUE(stream.CheckForOverflow().IsIOError());
+    EXPECT_EQ(std::streampos(cap), stream.tellp());
+    EXPECT_EQ(body.substr(0, cap), ReadRangeResponse(stream));
+    stream.clear();
+    stream.seekg(0);
+    EXPECT_EQ(body.substr(0, cap), ReadRangeResponse(stream));
+    stream.clear();
+    stream.seekg(cap + 1);
+    EXPECT_TRUE(stream.fail());
+  }
+}
+
+TEST(GetObjectRangeStreamTest, CapsPrefixWhenCallerBufferExceedsCap) {
+  const auto cap = internal::GetObjectRangeStreamBuf::kSpillSize;
+  std::string caller(cap + 8, '!');
+  internal::GetObjectRangeStream stream(caller.data(), caller.size());
+  const std::string prefix(caller.size(), 'a');
+  stream.write(prefix.data(), prefix.size());
+  ASSERT_OK(stream.CheckForOverflow());
+  stream.write("suffix", 6);
+
+  ASSERT_TRUE(stream.good());
+  EXPECT_TRUE(stream.CheckForOverflow().IsIOError());
+  EXPECT_EQ(std::streampos(cap), stream.tellp());
+  EXPECT_EQ(std::string(cap, 'a'), ReadRangeResponse(stream));
+}
+
+}  // namespace
+
 MinioTestEnvironment* GetMinioEnv(bool enable_tls) {
   if (enable_tls) {
     return ::arrow::internal::checked_cast<MinioTestEnvironment*>(minio_env_https);