diff --git a/dCommon/Sd0.cpp b/dCommon/Sd0.cpp index 8fbc4232f..c5b388d42 100644 --- a/dCommon/Sd0.cpp +++ b/dCommon/Sd0.cpp @@ -66,6 +66,11 @@ Sd0::Sd0(std::istream& buffer) { LOG("Failed to read chunk size from stream %lld %zu", buffer.tellg(), m_Chunks.size()); break; } + if (m_Chunks.size() >= MAX_CHUNKS) { + LOG("More than %u chunks in an sd0 buffer, rejecting.", MAX_CHUNKS); + m_Chunks.clear(); + break; + } auto& chunk = m_Chunks.emplace_back(); bool firstBuffer = m_Chunks.size() == 1; auto dataOffset = GetDataOffset(firstBuffer); @@ -155,7 +160,8 @@ std::string Sd0::GetAsStringUncompressed() const { break; } - auto oldSize = toReturn.size(); + // Each chunk is inflated right after the one before it + const auto oldSize = toReturn.size(); toReturn.resize(oldSize + MAX_UNCOMPRESSED_CHUNK_SIZE); int32_t error{}; const auto uncompressedSize = ZCompression::Decompress( @@ -171,6 +177,14 @@ std::string Sd0::GetAsStringUncompressed() const { } totalSize += uncompressedSize; + toReturn.resize(totalSize); + + if (totalSize > MAX_UNCOMPRESSED_SIZE) { + LOG("An sd0 buffer inflates to more than %u bytes, aborting", MAX_UNCOMPRESSED_SIZE); + toReturn = ""; + totalSize = 0; + break; + } } toReturn.resize(totalSize); diff --git a/dCommon/Sd0.h b/dCommon/Sd0.h index b4dae02ff..807839ed2 100644 --- a/dCommon/Sd0.h +++ b/dCommon/Sd0.h @@ -4,7 +4,9 @@ #ifndef SD0_H #define SD0_H +#include #include +#include #include #include #include @@ -23,6 +25,16 @@ public: */ static constexpr inline size_t MAX_UNCOMPRESSED_CHUNK_SIZE = 1024 * 256; + /** + * @brief Max number of zlib chunks in one sd0 buffer + */ + static constexpr inline uint32_t MAX_CHUNKS = 4096; + + /** + * @brief Max size of an sd0 buffer once all of its chunks are inflated, far above the size of a model + */ + static constexpr inline uint32_t MAX_UNCOMPRESSED_SIZE = 64 * 1024 * 1024; + // Empty, for FromData Sd0() = default; diff --git a/tests/dCommonTests/Sd0Tests.cpp b/tests/dCommonTests/Sd0Tests.cpp index 1e0b53e55..810f146dd 100644 --- a/tests/dCommonTests/Sd0Tests.cpp +++ b/tests/dCommonTests/Sd0Tests.cpp @@ -4,7 +4,10 @@ #include #include #include +#include +#include "Game.h" +#include "Logger.h" #include "Sd0.h" #include "ZCompression.h" @@ -65,3 +68,76 @@ TEST(Sd0, EmptyDataIsTheHeaderOnly) { Sd0 empty; EXPECT_TRUE(empty.GetAsVector().empty()); } + +class Sd0Test : public ::testing::Test { +protected: + void SetUp() override { + Game::logger = new Logger("./testing.log", true, true); + } + + void TearDown() override { + delete Game::logger; + Game::logger = nullptr; + } + + // One zlib chunk as it sits in an sd0 buffer: its size and then its data + static std::string MakeChunk(const std::string& data) { + std::vector compressed(ZCompression::GetMaxCompressedLength(data.size())); + const auto compressedSize = ZCompression::Compress( + reinterpret_cast(data.data()), data.size(), + compressed.data(), compressed.size()); + EXPECT_NE(compressedSize, -1); + + std::string chunk; + const uint32_t size = compressedSize; + chunk.append(reinterpret_cast(&size), sizeof(size)); + chunk.append(reinterpret_cast(compressed.data()), size); + return chunk; + } + + static std::string Inflate(const std::string& sd0Data) { + std::istringstream stream(sd0Data); + const Sd0 sd0(stream); + return sd0.GetAsStringUncompressed(); + } + + const std::string m_Header = std::string(Sd0::SD0_HEADER, 5); +}; + +TEST_F(Sd0Test, DataOfSeveralChunksSurvivesARoundTrip) { + std::string data; + for (uint32_t i = 0; data.size() < Sd0::MAX_UNCOMPRESSED_CHUNK_SIZE * 2 + 1000; i++) data += std::to_string(i) + ","; + + std::istringstream empty; + Sd0 sd0(empty); + sd0.FromData(reinterpret_cast(data.data()), data.size()); + EXPECT_EQ(sd0.GetAsVector().size(), 3); + EXPECT_EQ(sd0.GetAsStringUncompressed(), data); + + EXPECT_EQ(Inflate(sd0.GetAsStream().str()), data); +} + +TEST_F(Sd0Test, ChunksThatAreNotFullFollowOneAnother) { + EXPECT_EQ(Inflate(m_Header + MakeChunk("abc") + MakeChunk("def") + MakeChunk("ghi")), "abcdefghi"); +} + +TEST_F(Sd0Test, TooManyChunksAreRefused) { + const auto chunk = MakeChunk("a"); + + std::string allowed = m_Header; + for (uint32_t i = 0; i < Sd0::MAX_CHUNKS; i++) allowed += chunk; + EXPECT_EQ(Inflate(allowed), std::string(Sd0::MAX_CHUNKS, 'a')); + + EXPECT_EQ(Inflate(allowed + chunk), ""); +} + +TEST_F(Sd0Test, TooMuchInflatedDataIsRefused) { + const auto chunk = MakeChunk(std::string(Sd0::MAX_UNCOMPRESSED_CHUNK_SIZE, 'a')); + const uint32_t chunksToFill = Sd0::MAX_UNCOMPRESSED_SIZE / Sd0::MAX_UNCOMPRESSED_CHUNK_SIZE; + + std::string allowed = m_Header; + for (uint32_t i = 0; i < chunksToFill; i++) allowed += chunk; + EXPECT_EQ(Inflate(allowed).size(), Sd0::MAX_UNCOMPRESSED_SIZE); + + EXPECT_EQ(Inflate(allowed + MakeChunk("a")), ""); +}