diff --git a/components/nif/nifstream.cpp b/components/nif/nifstream.cpp index cd362892a6..51ba2c9bac 100644 --- a/components/nif/nifstream.cpp +++ b/components/nif/nifstream.cpp @@ -57,6 +57,7 @@ namespace Nif std::string NIFStream::getSizedString(size_t length) { + checkStreamSize(length); std::string str(length, '\0'); mStream->read(str.data(), length); if (mStream->fail()) @@ -90,6 +91,7 @@ namespace Nif std::string NIFStream::getStringPalette() { size_t size = get(); + checkStreamSize(size); std::string str(size, '\0'); mStream->read(str.data(), size); if (mStream->fail()) @@ -235,6 +237,7 @@ namespace Nif { if (getVersion() < generateVersion(4, 1, 0, 0)) { + checkStreamSize(size * sizeof(int32_t)); std::vector buf(size); read(buf.data(), size); for (size_t i = 0; i < size; ++i) @@ -242,6 +245,7 @@ namespace Nif } else { + checkStreamSize(size * sizeof(int8_t)); std::vector buf(size); read(buf.data(), size); for (size_t i = 0; i < size; ++i) @@ -264,4 +268,9 @@ namespace Nif } } + void NIFStream::checkStreamSize(std::size_t size) + { + if (size > mStreamSize) + throw std::runtime_error(std::format("Trying to read more than stream size: {} max={}", size, mStreamSize)); + } } diff --git a/components/nif/nifstream.hpp b/components/nif/nifstream.hpp index 6862bf882a..e5c2ce41d6 100644 --- a/components/nif/nifstream.hpp +++ b/components/nif/nifstream.hpp @@ -16,6 +16,7 @@ #include #include +#include #include #include @@ -78,6 +79,7 @@ namespace Nif Files::IStreamPtr mStream; const ToUTF8::StatelessUtf8Encoder* mEncoder; std::string mBuffer; + std::size_t mStreamSize; public: explicit NIFStream( @@ -85,6 +87,7 @@ namespace Nif : mReader(reader) , mStream(std::move(stream)) , mEncoder(encoder) + , mStreamSize(static_cast(Files::getStreamSizeLeft(*mStream))) { } @@ -129,6 +132,9 @@ namespace Nif { if (size == 0) return; + + checkStreamSize(size * sizeof(T)); + vec.resize(size); read(vec.data(), size); } @@ -159,6 +165,9 @@ namespace Nif /// Read a sequence of null-terminated strings std::string getStringPalette(); + + private: + void checkStreamSize(std::size_t size); }; template <>