end the zlib stream when readDataBlock and sync throw

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<typename EndFunction>
+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<Codec codec>
 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<InputStream> 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<InputStream> 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<InputStream> in = memoryInputStream(

lang/c++/test/DataFileTests.cc

@@ -24,6 +24,7 @@
 
 #include <chrono>
 #include <filesystem>
+#include <fstream>
 #include <thread>
 
 #include <sstream>
@@ -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<ComplexInteger> 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<char>(in), std::istreambuf_iterator<char>());
+    }
+    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<unsigned char>(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<std::streamsize>(bytes.size()));
+    }
+
+    const auto readCorrupted = [&] {
+        avro::DataFileReader<ComplexInteger> 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));