diff --git a/cpp/BUILD b/cpp/BUILD index a5946ec..2b697ad 100644 --- a/cpp/BUILD +++ b/cpp/BUILD @@ -296,7 +296,15 @@ cc_test( "@abseil-cpp//absl/strings", "@googletest//:gtest_main", "@riegeli//riegeli/base:initializer", + "@riegeli//riegeli/bytes:chain_writer", "@riegeli//riegeli/bytes:string_reader", "@riegeli//riegeli/bytes:string_writer", + "@riegeli//riegeli/chunk_encoding:chunk", + "@riegeli//riegeli/chunk_encoding:chunk_decoder", + "@riegeli//riegeli/chunk_encoding:compressor_options", + "@riegeli//riegeli/chunk_encoding:constants", + "@riegeli//riegeli/chunk_encoding:simple_encoder", + "@riegeli//riegeli/records:chunk_reader", + "@riegeli//riegeli/records:chunk_writer", ], ) diff --git a/cpp/array_record_reader.cc b/cpp/array_record_reader.cc index dcd1afe..484385f 100644 --- a/cpp/array_record_reader.cc +++ b/cpp/array_record_reader.cc @@ -409,6 +409,10 @@ void ArrayRecordReaderBase::Initialize() { } if (!state_->chunk_offsets->empty()) { + if (state_->record_group_size <= 0) { + Fail(InvalidArgumentError("record_group_size cannot be zero")); + return; + } // Finds minimal chunk_group_size that is larger equals to the readahead // buffer. A chunk_group corresponds to a PRead call. Smaller // chunk_group_size is better for random access, the converse is better @@ -624,6 +628,9 @@ absl::Status ArrayRecordReaderBase::ParallelReadRecordsWithIndices( state_->num_records); } uint64_t chunk_idx = record_idx / state_->record_group_size; + if (chunk_idx >= per_chunk_indices.size()) { + return OutOfRangeError("chunk_idx out of bounds"); + } uint64_t local_idx = record_idx - chunk_idx * state_->record_group_size; per_chunk_indices[chunk_idx].emplace_back(local_idx, indices_idx); } diff --git a/cpp/array_record_reader_test.cc b/cpp/array_record_reader_test.cc index f3efad5..55a9a56 100644 --- a/cpp/array_record_reader_test.cc +++ b/cpp/array_record_reader_test.cc @@ -35,8 +35,16 @@ limitations under the License. #include "cpp/test_utils.h" #include "cpp/thread_pool.h" #include "riegeli/base/maker.h" +#include "riegeli/bytes/chain_writer.h" #include "riegeli/bytes/string_reader.h" #include "riegeli/bytes/string_writer.h" +#include "riegeli/chunk_encoding/chunk.h" +#include "riegeli/chunk_encoding/chunk_decoder.h" +#include "riegeli/chunk_encoding/compressor_options.h" +#include "riegeli/chunk_encoding/constants.h" +#include "riegeli/chunk_encoding/simple_encoder.h" +#include "riegeli/records/chunk_reader.h" +#include "riegeli/records/chunk_writer.h" constexpr uint32_t kDatasetSize = 3210; @@ -313,5 +321,134 @@ TEST(ArrayRecordReaderOptionTest, ParserTest) { } } +std::string CorruptFooterGroupSize(absl::string_view valid_encoded, + int new_group_size) { + std::string corrupted_encoded; + riegeli::StringReader<> string_reader(valid_encoded); + riegeli::StringWriter<> string_writer(&corrupted_encoded); + riegeli::DefaultChunkReader<> chunk_reader(&string_reader); + riegeli::DefaultChunkWriter<> chunk_writer(&string_writer); + + bool past_footer = false; + riegeli::Chunk chunk; + while (chunk_reader.ReadChunk(chunk)) { + if (past_footer) { + chunk_writer.WriteChunk(chunk); + continue; + } + + riegeli::ChunkDecoder decoder; + if (!decoder.Decode(chunk)) { + chunk_writer.WriteChunk(chunk); + continue; + } + + decoder.SetIndex(0); + absl::string_view first_record; + if (!decoder.ReadRecord(first_record)) { + chunk_writer.WriteChunk(chunk); + continue; + } + + RiegeliFooterMetadata metadata; + if (metadata.ParsePartialFromString(first_record) && + metadata.has_array_record_metadata()) { + riegeli::SimpleEncoder footer_encoder( + riegeli::CompressorOptions().set_uncompressed()); + footer_encoder.AddRecord(first_record); + + absl::string_view footer_record; + while (decoder.ReadRecord(footer_record)) { + ArrayRecordFooter footer; + if (footer.ParsePartialFromString(footer_record) && + footer.has_chunk_offset()) { + footer.set_num_records(new_group_size); + footer_encoder.AddRecord(footer.SerializeAsString()); + } else { + footer_encoder.AddRecord(footer_record); + } + } + riegeli::Chunk corrupted_chunk; + riegeli::ChunkType chunk_type; + uint64_t num_records; + uint64_t decoded_data_size; + riegeli::ChainWriter<> chain_writer(&corrupted_chunk.data); + footer_encoder.EncodeAndClose(chain_writer, chunk_type, num_records, + decoded_data_size); + chain_writer.Close(); + corrupted_chunk.header = riegeli::ChunkHeader( + corrupted_chunk.data, chunk_type, num_records, decoded_data_size); + chunk_writer.WriteChunk(corrupted_chunk); + past_footer = true; + } else { + chunk_writer.WriteChunk(chunk); + } + } + chunk_writer.Close(); + string_writer.Close(); + + if (corrupted_encoded.size() < 64 * 1024) { + corrupted_encoded.resize(64 * 1024, '\0'); + } + return corrupted_encoded; +} + +TEST(ArrayRecordReaderTest, GroupSizeZero) { + std::string encoded; + auto writer_options = + ArrayRecordWriterBase::Options().set_group_size(10).set_uncompressed(); + auto writer = ArrayRecordWriter( + riegeli::Maker(&encoded), writer_options, nullptr); + for (int i = 0; i < 5; ++i) { + EXPECT_TRUE(writer.WriteRecord("test")); + } + ASSERT_TRUE(writer.Close()); + + std::string corrupted = CorruptFooterGroupSize(encoded, 0); + + auto reader_opt = ArrayRecordReaderBase::Options(); + auto reader = ArrayRecordReader( + riegeli::Maker(corrupted), reader_opt, nullptr); + EXPECT_FALSE(reader.status().ok()) << reader.status().message(); + EXPECT_EQ(reader.status().code(), absl::StatusCode::kInvalidArgument) + << reader.status().message(); +} + +TEST(ArrayRecordReaderTest, OutOfBoundsChunkIdx) { + std::string encoded; + auto writer_options = + ArrayRecordWriterBase::Options().set_group_size(10).set_uncompressed(); + auto writer = ArrayRecordWriter( + riegeli::Maker(&encoded), writer_options, nullptr); + for (int i = 0; i < 5; ++i) { + EXPECT_TRUE(writer.WriteRecord("test")); + } + ASSERT_TRUE(writer.Close()); + + // Valid file has 1 chunk. group_size in footer is 5. + // We corrupt it to group_size = 1. + // The metadata says num_records = 5, num_chunks = 1. + // Reader will think record_group_size = 1. + // So chunk_idx = record_idx / 1. + // If we read index 4: chunk_idx = 4. + // But per_chunk_indices size = num_chunks = 1. + // 4 >= 1, triggering OutOfRangeError! + std::string corrupted = CorruptFooterGroupSize(encoded, 1); + + auto reader_opt = ArrayRecordReaderBase::Options(); + auto reader = ArrayRecordReader( + riegeli::Maker(corrupted), reader_opt, nullptr); + ASSERT_TRUE(reader.status().ok()) << reader.status().message(); + + std::vector indices = {4}; + auto status = reader.ParallelReadRecordsWithIndices( + indices, [&](uint64_t, absl::string_view) -> absl::Status { + return absl::OkStatus(); + }); + + EXPECT_FALSE(status.ok()); + EXPECT_EQ(status.code(), absl::StatusCode::kOutOfRange); +} + } // namespace } // namespace array_record