diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index 2e5c67e07b6e..ec8b6c1b32f2 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -511,6 +511,7 @@ set(ARROW_UTIL_SRCS util/float16.cc util/formatting.cc util/future.cc + util/fuzz_internal.cc util/hashing.cc util/int_util.cc util/io_util.cc diff --git a/cpp/src/arrow/ipc/file_fuzz.cc b/cpp/src/arrow/ipc/file_fuzz.cc index 840d19a4efae..c3ce9c2dea6f 100644 --- a/cpp/src/arrow/ipc/file_fuzz.cc +++ b/cpp/src/arrow/ipc/file_fuzz.cc @@ -19,10 +19,11 @@ #include "arrow/ipc/reader.h" #include "arrow/status.h" +#include "arrow/util/fuzz_internal.h" #include "arrow/util/macros.h" extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) { auto status = arrow::ipc::internal::FuzzIpcFile(data, static_cast(size)); - ARROW_UNUSED(status); + arrow::internal::LogFuzzStatus(status, data, static_cast(size)); return 0; } diff --git a/cpp/src/arrow/ipc/reader.cc b/cpp/src/arrow/ipc/reader.cc index 1ec283662670..4910b1596c30 100644 --- a/cpp/src/arrow/ipc/reader.cc +++ b/cpp/src/arrow/ipc/reader.cc @@ -53,6 +53,7 @@ #include "arrow/util/checked_cast.h" #include "arrow/util/compression.h" #include "arrow/util/endian.h" +#include "arrow/util/fuzz_internal.h" #include "arrow/util/key_value_metadata.h" #include "arrow/util/logging_internal.h" #include "arrow/util/parallel.h" @@ -2618,6 +2619,12 @@ Status ValidateFuzzBatch(const RecordBatch& batch) { return st; } +IpcReadOptions FuzzingOptions() { + IpcReadOptions options; + options.memory_pool = ::arrow::internal::fuzzing_memory_pool(); + return options; +} + } // namespace Status FuzzIpcStream(const uint8_t* data, int64_t size) { @@ -2625,7 +2632,8 @@ Status FuzzIpcStream(const uint8_t* data, int64_t size) { io::BufferReader buffer_reader(buffer); std::shared_ptr batch_reader; - ARROW_ASSIGN_OR_RAISE(batch_reader, RecordBatchStreamReader::Open(&buffer_reader)); + ARROW_ASSIGN_OR_RAISE(batch_reader, + RecordBatchStreamReader::Open(&buffer_reader, FuzzingOptions())); Status st; while (true) { @@ -2645,7 +2653,8 @@ Status FuzzIpcFile(const uint8_t* data, int64_t size) { io::BufferReader buffer_reader(buffer); std::shared_ptr batch_reader; - ARROW_ASSIGN_OR_RAISE(batch_reader, RecordBatchFileReader::Open(&buffer_reader)); + ARROW_ASSIGN_OR_RAISE(batch_reader, + RecordBatchFileReader::Open(&buffer_reader, FuzzingOptions())); Status st; const int n_batches = batch_reader->num_record_batches(); diff --git a/cpp/src/arrow/ipc/stream_fuzz.cc b/cpp/src/arrow/ipc/stream_fuzz.cc index e26f3d1f4e10..04b12863beda 100644 --- a/cpp/src/arrow/ipc/stream_fuzz.cc +++ b/cpp/src/arrow/ipc/stream_fuzz.cc @@ -19,10 +19,11 @@ #include "arrow/ipc/reader.h" #include "arrow/status.h" +#include "arrow/util/fuzz_internal.h" #include "arrow/util/macros.h" extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) { auto status = arrow::ipc::internal::FuzzIpcStream(data, static_cast(size)); - ARROW_UNUSED(status); + arrow::internal::LogFuzzStatus(status, data, static_cast(size)); return 0; } diff --git a/cpp/src/arrow/memory_pool.h b/cpp/src/arrow/memory_pool.h index 19a938c33601..a8f964200ed4 100644 --- a/cpp/src/arrow/memory_pool.h +++ b/cpp/src/arrow/memory_pool.h @@ -26,6 +26,7 @@ #include "arrow/result.h" #include "arrow/status.h" #include "arrow/type_fwd.h" +#include "arrow/util/macros.h" #include "arrow/util/visibility.h" namespace arrow { @@ -245,6 +246,73 @@ class ARROW_EXPORT ProxyMemoryPool : public MemoryPool { std::unique_ptr impl_; }; +/// EXPERIMENTAL MemoryPool wrapper with an upper limit +/// +/// Checking for limits is not done in a fully thread-safe way, therefore +/// multi-threaded allocations might be able to go successfully above the +/// configured limit. +class ARROW_EXPORT CappedMemoryPool : public MemoryPool { + public: + CappedMemoryPool(MemoryPool* wrapped_pool, int64_t bytes_allocated_limit) + : wrapped_(wrapped_pool), bytes_allocated_limit_(bytes_allocated_limit) {} + + using MemoryPool::Allocate; + using MemoryPool::Reallocate; + + Status Allocate(int64_t size, int64_t alignment, uint8_t** out) override { + // XXX Another thread may allocate memory between the limit check and + // the `Allocate` call. It is possible for the two allocations to be successful + // while going above the limit. + // Solving this issue would require refactoring the `MemoryPool` implementation + // to delegate the limit check to `MemoryPoolStats`. + const auto attempted = size + wrapped_->bytes_allocated(); + if (ARROW_PREDICT_FALSE(attempted > bytes_allocated_limit_)) { + return OutOfMemory(attempted); + } + return wrapped_->Allocate(size, alignment, out); + } + + Status Reallocate(int64_t old_size, int64_t new_size, int64_t alignment, + uint8_t** ptr) override { + const auto attempted = new_size - old_size + wrapped_->bytes_allocated(); + if (ARROW_PREDICT_FALSE(attempted > bytes_allocated_limit_)) { + return OutOfMemory(attempted); + } + return wrapped_->Reallocate(old_size, new_size, alignment, ptr); + } + + void Free(uint8_t* buffer, int64_t size, int64_t alignment) override { + return wrapped_->Free(buffer, size, alignment); + } + + void ReleaseUnused() override { wrapped_->ReleaseUnused(); } + + void PrintStats() override { wrapped_->PrintStats(); } + + int64_t bytes_allocated() const override { return wrapped_->bytes_allocated(); } + + int64_t max_memory() const override { return wrapped_->max_memory(); } + + int64_t total_bytes_allocated() const override { + return wrapped_->total_bytes_allocated(); + } + + int64_t num_allocations() const override { return wrapped_->num_allocations(); } + + std::string backend_name() const override { return wrapped_->backend_name(); } + + private: + Status OutOfMemory(int64_t value) { + return Status::OutOfMemory( + "MemoryPool bytes_allocated cap exceeded: " + "limit=", + bytes_allocated_limit_, ", attempted=", value); + } + + MemoryPool* wrapped_; + const int64_t bytes_allocated_limit_; +}; + /// \brief Return a process-wide memory pool based on the system allocator. ARROW_EXPORT MemoryPool* system_memory_pool(); diff --git a/cpp/src/arrow/memory_pool_test.cc b/cpp/src/arrow/memory_pool_test.cc index 3f0a85287671..20006ebeb49a 100644 --- a/cpp/src/arrow/memory_pool_test.cc +++ b/cpp/src/arrow/memory_pool_test.cc @@ -17,6 +17,7 @@ #include #include +#include #include @@ -290,4 +291,93 @@ TEST(Jemalloc, GetAllocationStats) { #endif } +class TestCappedMemoryPool : public ::arrow::TestMemoryPoolBase { + public: + MemoryPool* memory_pool() override { return InitPool(/*limit=*/1'000'000'000LL); } + + MemoryPool* InitPool(int64_t limit) { + proxy_memory_pool_ = std::make_shared(default_memory_pool()); + capped_memory_pool_ = + std::make_shared(proxy_memory_pool_.get(), limit); + return capped_memory_pool_.get(); + } + + protected: + std::shared_ptr proxy_memory_pool_; + std::shared_ptr capped_memory_pool_; +}; + +TEST_F(TestCappedMemoryPool, MemoryTracking) { this->TestMemoryTracking(); } + +TEST_F(TestCappedMemoryPool, OOM) { + // CappedMemoryPool rejects the huge allocation without hitting the underlying + // allocator, so this should work even under Address Sanitizer. + this->TestOOM(); +} + +TEST_F(TestCappedMemoryPool, Reallocate) { this->TestReallocate(); } + +TEST_F(TestCappedMemoryPool, Alignment) { this->TestAlignment(); } + +TEST_F(TestCappedMemoryPool, AllocateLimit) { + auto pool = InitPool(/*limit=*/1000); + + uint8_t* data1; + uint8_t* data2; + ASSERT_OK(pool->Allocate(600, &data1)); + ASSERT_EQ(600, pool->bytes_allocated()); + ASSERT_EQ(600, pool->total_bytes_allocated()); + ASSERT_EQ(600, pool->max_memory()); + + ASSERT_OK(pool->Allocate(400, &data2)); + ASSERT_EQ(1000, pool->bytes_allocated()); + ASSERT_EQ(1000, pool->total_bytes_allocated()); + ASSERT_EQ(1000, pool->max_memory()); + pool->Free(data2, 400); + ASSERT_EQ(600, pool->bytes_allocated()); + ASSERT_EQ(1000, pool->total_bytes_allocated()); + ASSERT_EQ(1000, pool->max_memory()); + + ASSERT_OK(pool->Allocate(300, &data2)); + ASSERT_EQ(900, pool->bytes_allocated()); + ASSERT_EQ(1300, pool->total_bytes_allocated()); + ASSERT_EQ(1000, pool->max_memory()); + pool->Free(data2, 300); + ASSERT_EQ(600, pool->bytes_allocated()); + ASSERT_EQ(1300, pool->total_bytes_allocated()); + ASSERT_EQ(1000, pool->max_memory()); + + ASSERT_RAISES(OutOfMemory, pool->Allocate(401, &data2)); + ASSERT_EQ(600, pool->bytes_allocated()); + ASSERT_EQ(1300, pool->total_bytes_allocated()); + ASSERT_EQ(1000, pool->max_memory()); + + pool->Free(data1, 600); +} + +TEST_F(TestCappedMemoryPool, ReallocateLimit) { + auto pool = InitPool(/*limit=*/1000); + + uint8_t* data1; + uint8_t* data2; + ASSERT_OK(pool->Allocate(600, &data1)); + ASSERT_OK(pool->Allocate(400, &data2)); + ASSERT_EQ(1000, pool->bytes_allocated()); + ASSERT_EQ(1000, pool->total_bytes_allocated()); + ASSERT_EQ(1000, pool->max_memory()); + + ASSERT_OK(pool->Reallocate(400, 300, &data2)); + ASSERT_EQ(900, pool->bytes_allocated()); + ASSERT_EQ(1000, pool->total_bytes_allocated()); + ASSERT_EQ(1000, pool->max_memory()); + + ASSERT_RAISES(OutOfMemory, pool->Reallocate(300, 401, &data2)); + ASSERT_EQ(900, pool->bytes_allocated()); + ASSERT_EQ(1000, pool->total_bytes_allocated()); + ASSERT_EQ(1000, pool->max_memory()); + + pool->Free(data1, 600); + pool->Free(data2, 300); +} + } // namespace arrow diff --git a/cpp/src/arrow/util/fuzz_internal.cc b/cpp/src/arrow/util/fuzz_internal.cc new file mode 100644 index 000000000000..c4eddc51265d --- /dev/null +++ b/cpp/src/arrow/util/fuzz_internal.cc @@ -0,0 +1,40 @@ +// 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. + +#include "arrow/util/fuzz_internal.h" + +#include "arrow/memory_pool.h" +#include "arrow/status.h" +#include "arrow/util/logging_internal.h" + +namespace arrow::internal { + +MemoryPool* fuzzing_memory_pool() { + static auto pool = std::make_shared<::arrow::CappedMemoryPool>( + ::arrow::default_memory_pool(), /*bytes_allocated_limit=*/kFuzzingMemoryLimit); + return pool.get(); +} + +void LogFuzzStatus(const Status& st, const uint8_t* data, int64_t size) { + // Most fuzz inputs will be invalid and generate errors, only log potential OOMs + if (st.IsOutOfMemory()) { + ARROW_LOG(WARNING) << "Fuzzing input with size=" << size + << " hit allocation failure: " << st.ToString(); + } +} + +} // namespace arrow::internal diff --git a/cpp/src/arrow/util/fuzz_internal.h b/cpp/src/arrow/util/fuzz_internal.h new file mode 100644 index 000000000000..f3f104eec602 --- /dev/null +++ b/cpp/src/arrow/util/fuzz_internal.h @@ -0,0 +1,37 @@ +// 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 "arrow/type_fwd.h" +#include "arrow/util/macros.h" + +namespace arrow::internal { + +// The default rss_limit_mb on OSS-Fuzz is 2560 MB and we want to fail allocations +// before that limit is reached, otherwise the fuzz target gets killed (GH-48105). +constexpr int64_t kFuzzingMemoryLimit = 2200LL * 1000 * 1000; + +/// Return a memory pool that will not allocate more than kFuzzingMemoryLimit bytes. +ARROW_EXPORT MemoryPool* fuzzing_memory_pool(); + +// Optionally log the outcome of fuzzing an input +ARROW_EXPORT void LogFuzzStatus(const Status&, const uint8_t* data, int64_t size); + +} // namespace arrow::internal diff --git a/cpp/src/parquet/arrow/fuzz.cc b/cpp/src/parquet/arrow/fuzz.cc index f1c724508c43..8a4cfbceffb4 100644 --- a/cpp/src/parquet/arrow/fuzz.cc +++ b/cpp/src/parquet/arrow/fuzz.cc @@ -16,10 +16,11 @@ // under the License. #include "arrow/status.h" +#include "arrow/util/fuzz_internal.h" #include "parquet/arrow/reader.h" extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) { auto status = parquet::arrow::internal::FuzzReader(data, static_cast(size)); - ARROW_UNUSED(status); + arrow::internal::LogFuzzStatus(status, data, static_cast(size)); return 0; } diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index 2933664b27ea..c32e58868802 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -29,6 +29,7 @@ #include "arrow/buffer.h" #include "arrow/extension_type.h" #include "arrow/io/memory.h" +#include "arrow/memory_pool.h" #include "arrow/record_batch.h" #include "arrow/table.h" #include "arrow/type.h" @@ -36,6 +37,7 @@ #include "arrow/util/async_generator.h" #include "arrow/util/bit_util.h" #include "arrow/util/future.h" +#include "arrow/util/fuzz_internal.h" #include "arrow/util/iterator.h" #include "arrow/util/logging_internal.h" #include "arrow/util/parallel.h" @@ -1403,7 +1405,7 @@ namespace internal { namespace { -Status FuzzReader(std::unique_ptr reader) { +Status FuzzReadData(std::unique_ptr reader) { auto st = Status::OK(); for (int i = 0; i < reader->num_row_groups(); ++i) { std::shared_ptr table; @@ -1490,7 +1492,7 @@ Status FuzzReader(const uint8_t* data, int64_t size) { auto buffer = std::make_shared<::arrow::Buffer>(data, size); auto file = std::make_shared<::arrow::io::BufferReader>(buffer); - auto pool = ::arrow::default_memory_pool(); + auto pool = ::arrow::internal::fuzzing_memory_pool(); auto reader_properties = default_reader_properties(); std::default_random_engine rng(/*seed*/ 42); @@ -1562,7 +1564,7 @@ Status FuzzReader(const uint8_t* data, int64_t size) { std::unique_ptr reader; RETURN_NOT_OK(FileReader::Make(pool, std::move(pq_file_reader), properties, &reader)); - st &= FuzzReader(std::move(reader)); + st &= FuzzReadData(std::move(reader)); } return st; }