From 6daf425aa99ee1e78028c4d4dae45a690c0c2385 Mon Sep 17 00:00:00 2001 From: Chinmay Shrivastava <67478015+JonSnow1807@users.noreply.github.com> Date: Sat, 7 Nov 2026 00:51:30 -0600 Subject: [PATCH] GH-55428: [C++][S3] Preserve retryable errors when ranged-read buffers overflow --- cpp/src/arrow/filesystem/s3_stream_internal.h | 169 ++++++++++++++++++ cpp/src/arrow/filesystem/s3fs.cc | 12 +- cpp/src/arrow/filesystem/s3fs_test.cc | 137 ++++++++++++++ python/pyarrow/tests/test_fs.py | 133 ++++++++++++++ 4 files changed, 449 insertions(+), 2 deletions(-) create mode 100644 cpp/src/arrow/filesystem/s3_stream_internal.h diff --git a/cpp/src/arrow/filesystem/s3_stream_internal.h b/cpp/src/arrow/filesystem/s3_stream_internal.h new file mode 100644 index 000000000000..2ddc5b5c3b0d --- /dev/null +++ b/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 +#include +#include +#include +#include +#include + +#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(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(count); + while (remaining != 0) { + if (pptr() == epptr() && !spill_) { + Spill(); + } + const auto n = std::min(remaining, static_cast(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(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(Written())) + : pos_type(off_type(-1)); + } + if (which != std::ios_base::in) { + return pos_type(off_type(-1)); + } + const auto size = static_cast(Written()); + const auto base = dir == std::ios_base::beg ? off_type(0) + : dir == std::ios_base::cur + ? static_cast(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(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(pptr() - Storage()); } + + void Advance(size_t n) { + while (n != 0) { + const auto step = static_cast( + std::min(n, static_cast(std::numeric_limits::max()))); + pbump(step); + n -= step; + } + } + + void Spill() { + overflowed_ = true; + const auto kept = std::min(Written(), kSpillSize); + const auto read_pos = std::min(static_cast(gptr() - eback()), kept); + auto spill = std::make_unique(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 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 diff --git a/cpp/src/arrow/filesystem/s3fs.cc b/cpp/src/arrow/filesystem/s3fs.cc index 76b1ce616143..1d5ca581f6ac 100644 --- a/cpp/src/arrow/filesystem/s3fs.cc +++ b/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("", data, nbytes); }; + return [=]() { + return Aws::New("", data, + static_cast(nbytes)); + }; } Result GetObjectRange(Aws::S3::S3Client* client, @@ -1341,7 +1345,11 @@ Result 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(result.GetBody()).CheckForOverflow()); + return std::move(result); } template diff --git a/cpp/src/arrow/filesystem/s3fs_test.cc b/cpp/src/arrow/filesystem/s3fs_test.cc index 06d25287a4f3..eee4b20be585 100644 --- a/cpp/src/arrow/filesystem/s3fs_test.cc +++ b/cpp/src/arrow/filesystem/s3fs_test.cc @@ -16,7 +16,9 @@ // under the License. #include +#include #include +#include #include #include #include @@ -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(stream), + std::istreambuf_iterator()); +} + +class DirectWriteRangeBuffer : public internal::GetObjectRangeStreamBuf { + public: + using internal::GetObjectRangeStreamBuf::GetObjectRangeStreamBuf; + + void WriteDirectly(const char* data, size_t size) { + ASSERT_LE(size, static_cast(epptr() - pptr())); + std::memcpy(pptr(), data, size); + pbump(static_cast(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 = "SlowDownretry"; + 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(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(input), + std::istreambuf_iterator())); +} + +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(minio_env_https); diff --git a/python/pyarrow/tests/test_fs.py b/python/pyarrow/tests/test_fs.py index 5bf1950c0654..a5e88640fe5b 100644 --- a/python/pyarrow/tests/test_fs.py +++ b/python/pyarrow/tests/test_fs.py @@ -16,9 +16,11 @@ # under the License. from datetime import datetime, timezone, timedelta +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import gzip import os import pathlib +import threading from urllib.request import urlopen import subprocess import sys @@ -1234,6 +1236,137 @@ def test_gcs_options(pickle_with_and_without_subtree_filesystem): GcsFileSystem(credential_token_expiration=datetime.now()) +@pytest.fixture +def s3_range_server(monkeypatch): + from pyarrow.fs import AwsStandardS3RetryStrategy, S3FileSystem + + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + servers = [] + + def start(size, responses): + ranges = [] + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_HEAD(self): + self.send_response(200) + self.send_header("Content-Length", str(size)) + self.end_headers() + + def do_GET(self): + ranges.append(self.headers.get("Range")) + response_index = min(len(ranges) - 1, len(responses) - 1) + status, body, chunked = responses[response_index] + self.send_response(status) + if status in (200, 206): + self.send_header("Content-Type", "application/octet-stream") + if status == 206: + self.send_header("Content-Range", + f"bytes 0-{size - 1}/{size}") + else: + self.send_header("Content-Type", "application/xml") + if chunked: + self.send_header("Transfer-Encoding", "chunked") + else: + self.send_header("Content-Length", str(len(body))) + self.end_headers() + if chunked: + self.wfile.write(f"{len(body):X}\r\n".encode("ascii")) + self.wfile.write(body + b"\r\n0\r\n\r\n") + else: + self.wfile.write(body) + + def log_message(self, *_args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + servers.append((server, thread)) + endpoint = f"127.0.0.1:{server.server_port}" + fs = S3FileSystem(access_key="test", secret_key="test", + region="us-west-2", scheme="http", + endpoint_override=endpoint, + retry_strategy=AwsStandardS3RetryStrategy( + max_attempts=3), connect_timeout=2) + return fs, ranges + + yield start + + for server, thread in servers: + server.shutdown() + server.server_close() + thread.join() + + +_S3_SLOW_DOWN_XML = ( + b"SlowDown" + b"Please reduce your request rate." + b"rid" +) + + +@pytest.mark.s3 +@pytest.mark.parametrize("status, size", [(429, 121), (503, 121), (429, 8192)]) +def test_s3_ranged_read_retries_large_error_body(s3_range_server, status, size): + xml = _S3_SLOW_DOWN_XML + assert 121 < len(xml) < 8192 + fs, ranges = s3_range_server(size, [(status, xml, False)]) + with pytest.raises(OSError) as excinfo: + with fs.open_input_file("repro-bucket/tiny.bin") as source: + source.read() + + assert "SLOW_DOWN" in str(excinfo.value) + assert "Failed to flush response stream" not in str(excinfo.value) + assert ranges == [f"bytes=0-{size - 1}"] * 3 + + +@pytest.mark.s3 +@pytest.mark.parametrize("status", [429, 503]) +def test_s3_ranged_read_recovers_after_slowdown(s3_range_server, status): + data = bytes(range(121)) + responses = [(status, _S3_SLOW_DOWN_XML, False)] * 2 + responses.append((206, data, False)) + fs, ranges = s3_range_server(len(data), responses) + + with fs.open_input_file("repro-bucket/tiny.bin") as source: + assert source.read() == data + + assert ranges == [f"bytes=0-{len(data) - 1}"] * 3 + + +@pytest.mark.s3 +@pytest.mark.parametrize( + "status, chunked", + [(200, False), (206, False), (200, True), (206, True)]) +def test_s3_rejects_oversized_range(s3_range_server, status, chunked): + data = b"a" * 17 + fs, ranges = s3_range_server(16, [(status, data, chunked)]) + + with pytest.raises( + OSError, match="GetObject response exceeded the requested range"): + with fs.open_input_file("repro-bucket/tiny.bin") as source: + source.read() + + assert ranges == ["bytes=0-15"] + + +@pytest.mark.s3 +@pytest.mark.parametrize("status", [429, 503]) +def test_s3_retries_error_beyond_spill(s3_range_server, status): + xml = (b"SlowDown" + b"x" * (64 * 1024) + + b"") + fs, ranges = s3_range_server(121, [(status, xml, False)]) + + with pytest.raises(OSError) as excinfo: + with fs.open_input_file("repro-bucket/tiny.bin") as source: + source.read() + + assert "Failed to flush response stream" not in str(excinfo.value) + assert ranges == ["bytes=0-120"] * 3 + + @pytest.mark.s3 def test_s3_options(pickle_with_and_without_subtree_filesystem): pickle_module = pickle_with_and_without_subtree_filesystem