diff --git a/dCommon/GeneralUtils.cpp b/dCommon/GeneralUtils.cpp index a11d991a9..7004b2fdc 100644 --- a/dCommon/GeneralUtils.cpp +++ b/dCommon/GeneralUtils.cpp @@ -292,10 +292,12 @@ std::vector GeneralUtils::SplitString(const std::string_view str, c } std::u16string GeneralUtils::ReadWString(RakNet::BitStream& inStream) { - uint32_t length; + uint32_t length = 0; inStream.Read(length); std::u16string string; + if (!StreamHolds(inStream, length, sizeof(uint16_t))) return string; + for (uint32_t i = 0; i < length; ++i) { uint16_t c; inStream.Read(c); diff --git a/dCommon/GeneralUtils.h b/dCommon/GeneralUtils.h index 3c345241e..150c6165d 100644 --- a/dCommon/GeneralUtils.h +++ b/dCommon/GeneralUtils.h @@ -143,6 +143,18 @@ namespace GeneralUtils { bool ReplaceInString(std::string& str, const std::string_view from, const std::string_view to); + /** + * Whether a stream still holds a number of values of a size. A length read from a stream is checked with this + * before that many values are read, as the stream can not hold more values than it has bytes left. + * @param inStream The stream the values are read from + * @param count The number of values + * @param valueSize The size of one value in bytes + */ + inline bool StreamHolds(const RakNet::BitStream& inStream, const uint64_t count, const uint64_t valueSize = 1) { + // Neither side can wrap: the count is at most 32 bits wide wherever it is read from a stream + return count <= UINT32_MAX && count * valueSize * 8 <= static_cast(inStream.GetNumberOfUnreadBits()); + } + std::u16string ReadWString(RakNet::BitStream& inStream); std::vector SplitString(const std::u16string_view str, const char16_t delimiter); diff --git a/tests/dCommonTests/TestGeneralUtils.cpp b/tests/dCommonTests/TestGeneralUtils.cpp index 682c77fa8..6af7f95fe 100644 --- a/tests/dCommonTests/TestGeneralUtils.cpp +++ b/tests/dCommonTests/TestGeneralUtils.cpp @@ -33,3 +33,36 @@ TEST(GeneralUtilsTest, UnpredictableNumberFillsEveryByte) { EXPECT_EQ(seenBits, 0xFFFFFFFFFFFFFFFFULL); } + +TEST(GeneralUtilsTest, StreamHoldsCountsTheBytesLeft) { + RakNet::BitStream stream; + stream.Write(4); + for (uint16_t i = 0; i < 4; i++) stream.Write(u'a' + i); + + uint32_t length = 0; + stream.Read(length); + EXPECT_TRUE(GeneralUtils::StreamHolds(stream, 0)); + EXPECT_TRUE(GeneralUtils::StreamHolds(stream, 8)); + EXPECT_FALSE(GeneralUtils::StreamHolds(stream, 9)); + EXPECT_TRUE(GeneralUtils::StreamHolds(stream, 4, sizeof(uint16_t))); + EXPECT_FALSE(GeneralUtils::StreamHolds(stream, 5, sizeof(uint16_t))); + EXPECT_FALSE(GeneralUtils::StreamHolds(stream, UINT32_MAX, sizeof(uint16_t))); + EXPECT_FALSE(GeneralUtils::StreamHolds(stream, UINT64_MAX, sizeof(uint16_t))); + // What a negative 32 bit length becomes + EXPECT_FALSE(GeneralUtils::StreamHolds(stream, static_cast(-1), sizeof(uint16_t))); +} + +TEST(GeneralUtilsTest, ReadWStringReadsWhatTheStreamHolds) { + RakNet::BitStream stream; + stream.Write(4); + for (uint16_t i = 0; i < 4; i++) stream.Write(u'a' + i); + EXPECT_EQ(GeneralUtils::ReadWString(stream), u"abcd"); + + RakNet::BitStream tooLong; + tooLong.Write(0xFFFFFFFF); + tooLong.Write(u'a'); + EXPECT_EQ(GeneralUtils::ReadWString(tooLong), u""); + + RakNet::BitStream empty; + EXPECT_EQ(GeneralUtils::ReadWString(empty), u""); +}