diff --git a/src/serialization.cpp b/src/serialization.cpp index 36ddb467c..310604f54 100644 --- a/src/serialization.cpp +++ b/src/serialization.cpp @@ -99,7 +99,7 @@ void compressZlib(const std::string &data, std::ostream &os, int level) compressZlib((u8*)data.c_str(), data.size(), os, level); } -void decompressZlib(std::istream &is, std::ostream &os) +void decompressZlib(std::istream &is, std::ostream &os, size_t limit) { z_stream z; const s32 bufsize = 16384; @@ -108,6 +108,7 @@ void decompressZlib(std::istream &is, std::ostream &os) int status = 0; int ret; int bytes_read = 0; + int bytes_written = 0; int input_buffer_len = 0; z.zalloc = Z_NULL; @@ -124,8 +125,20 @@ void decompressZlib(std::istream &is, std::ostream &os) for(;;) { + int output_size = bufsize; z.next_out = (Bytef*)output_buffer; - z.avail_out = bufsize; + z.avail_out = output_size; + + if (limit) { + int limit_remaining = limit - bytes_written; + if (limit_remaining <= 0) { + // we're aborting ahead of time - throw an error? + break; + } + if (limit_remaining < output_size) { + z.avail_out = output_size = limit_remaining; + } + } if(z.avail_in == 0) { @@ -153,10 +166,11 @@ void decompressZlib(std::istream &is, std::ostream &os) zerr(status); throw SerializationError("decompressZlib: inflate failed"); } - int count = bufsize - z.avail_out; + int count = output_size - z.avail_out; //dstream<<"count="< &data, std::ostream &os, u8 version); diff --git a/src/unittest/test_compression.cpp b/src/unittest/test_compression.cpp index 7d0378131..dfcadd4b2 100644 --- a/src/unittest/test_compression.cpp +++ b/src/unittest/test_compression.cpp @@ -37,6 +37,8 @@ public: void testRLECompression(); void testZlibCompression(); void testZlibLargeData(); + void testZlibLimit(); + void _testZlibLimit(u32 size, u32 limit); }; static TestCompression g_test_instance; @@ -46,6 +48,7 @@ void TestCompression::runTests(IGameDef *gamedef) TEST(testRLECompression); TEST(testZlibCompression); TEST(testZlibLargeData); + TEST(testZlibLimit); } //////////////////////////////////////////////////////////////////////////////// @@ -170,3 +173,63 @@ void TestCompression::testZlibLargeData() i, str_decompressed[i], i, data_in[i]); } } + +void TestCompression::testZlibLimit() +{ + // edge cases + _testZlibLimit(1024, 1023); + _testZlibLimit(1024, 1024); + _testZlibLimit(1024, 1025); + + // test around buffer borders + u32 bufsize = 16384; // as in implementation + for (int s = -1; s <= 1; s++) + { + for (int l = -1; l <= 1; l++) + { + _testZlibLimit(bufsize + s, bufsize + l); + } + } + // span multiple buffers + _testZlibLimit(35000, 22000); + _testZlibLimit(22000, 35000); +} + +void TestCompression::_testZlibLimit(u32 size, u32 limit) +{ + infostream << "Test: Testing zlib wrappers with a decompression " + "memory limit of " << limit << std::endl; + + infostream << "Test: Input size of compressZlib for limit is " + << size << std::endl; + + // how much data we expect to get + u32 expected = size < limit ? size : limit; + + // create recognizable data + std::string data_in; + data_in.resize(size); + for (u32 i = 0; i < size; i++) + data_in[i] = (u8)(i % 256); + + std::ostringstream os_compressed(std::ios::binary); + compressZlib(data_in, os_compressed); + infostream << "Test: Output size of compressZlib for limit is " + << os_compressed.str().size()<