diff --git a/lang/c++/impl/DataFile.cc b/lang/c++/impl/DataFile.cc index 661ee583d08..b654dfd1d44 100644 --- a/lang/c++/impl/DataFile.cc +++ b/lang/c++/impl/DataFile.cc @@ -55,6 +55,24 @@ const size_t maxSyncInterval = 1u << 30; // Recommended by https://www.zlib.net/zlib_how.html const size_t zlibBufGrowSize = 128 * 1024; +// Releases a zlib stream when leaving the enclosing scope, so the state +// allocated by deflateInit2()/inflateInit2() is freed on the paths that throw +// as well as on the one that runs to completion. +template +class ZlibStreamGuard { +public: + ZlibStreamGuard(z_stream &zs, EndFunction end) : zs_(zs), end_(end) {} + ~ZlibStreamGuard() { + (void) end_(&zs_); + } + ZlibStreamGuard(const ZlibStreamGuard &) = delete; + ZlibStreamGuard &operator=(const ZlibStreamGuard &) = delete; + +private: + z_stream &zs_; + EndFunction end_; +}; + template struct codec_trait { static std::string name() { @@ -284,6 +302,7 @@ void DataFileWriterBase::sync() { if (ret != Z_OK) { throw Exception("Failed to initialize deflate, error: {}", ret); } + ZlibStreamGuard guard(zs, deflateEnd); std::unique_ptr input = memoryInputStream(*buffer_); const uint8_t *data; @@ -309,7 +328,6 @@ void DataFileWriterBase::sync() { } buf.resize(zs.total_out); - (void) deflateEnd(&zs); } // make sure all is flushed std::unique_ptr in = memoryInputStream(buf.data(), buf.size()); int64_t byteCount = buf.size(); @@ -624,6 +642,7 @@ void DataFileReaderBase::readDataBlock() { if (ret != Z_OK) { throw Exception("Failed to initialize inflate, error: {}", ret); } + ZlibStreamGuard guard(zs, inflateEnd); const uint8_t *data; size_t len; @@ -647,7 +666,6 @@ void DataFileReaderBase::readDataBlock() { } uncompressed.resize(zs.total_out); - (void) inflateEnd(&zs); } std::unique_ptr in = memoryInputStream( diff --git a/lang/c++/test/DataFileTests.cc b/lang/c++/test/DataFileTests.cc index ed4e93c0663..39355dba055 100644 --- a/lang/c++/test/DataFileTests.cc +++ b/lang/c++/test/DataFileTests.cc @@ -24,6 +24,7 @@ #include #include +#include #include #include @@ -1469,6 +1470,58 @@ void testZstdCompressionRoundTrip() { } #endif +void testCorruptDeflateBlock() { + BOOST_TEST_CHECKPOINT(__func__); + + avro::ValidSchema schema = avro::compileJsonSchemaFromString(sch); + const char *filename = "test_corrupt_deflate.df"; + + { + avro::DataFileWriter writer( + filename, schema, 16 * 1024, avro::DEFLATE_CODEC); + for (int i = 0; i < 100; ++i) { + writer.write(ComplexInteger(i, 2 * i)); + } + writer.close(); + } + + string bytes; + { + std::ifstream in(filename, std::ios::binary); + bytes.assign(std::istreambuf_iterator(in), std::istreambuf_iterator()); + } + BOOST_REQUIRE_GT(bytes.size(), 16u); + + // The 16-byte sync marker ends both the header and every block, so its + // first occurrence is followed by the first block: the object count and + // the block size, both varints, and then the deflate payload. + const size_t header = bytes.find(bytes.substr(bytes.size() - 16)); + BOOST_REQUIRE(header != string::npos); + size_t payload = header + 16; + for (int varint = 0; varint < 2; ++varint) { + while (payload < bytes.size() && (static_cast(bytes[payload]) & 0x80) != 0) { + ++payload; + } + ++payload; + } + BOOST_REQUIRE_LT(payload, bytes.size()); + // BTYPE 3 is reserved in RFC 1951, so inflate() rejects the block outright. + bytes[payload] = '\xff'; + { + std::ofstream out(filename, std::ios::binary | std::ios::trunc); + out.write(bytes.data(), static_cast(bytes.size())); + } + + const auto readCorrupted = [&] { + avro::DataFileReader reader(filename, schema); + ComplexInteger record; + reader.read(record); + }; + BOOST_CHECK_THROW(readCorrupted(), avro::Exception); + + std::filesystem::remove(filename); +} + void testCodecEnumValues() { BOOST_TEST_CHECKPOINT(__func__); @@ -1700,6 +1753,8 @@ init_unit_test_suite(int, char *[]) { boost::unit_test::framework::master_test_suite().add(BOOST_TEST_CASE(&testMetadataWithZstdCodec)); #endif + boost::unit_test::framework::master_test_suite().add(BOOST_TEST_CASE(&testCorruptDeflateBlock)); + // Codec enum and isCodecAvailable tests boost::unit_test::framework::master_test_suite().add(BOOST_TEST_CASE(&testCodecEnumValues)); boost::unit_test::framework::master_test_suite().add(BOOST_TEST_CASE(&testIsCodecAvailable));