diff --git a/include/Engine/Network/Package.h b/include/Engine/Network/Package.h index 9b45101d..2d9f766f 100644 --- a/include/Engine/Network/Package.h +++ b/include/Engine/Network/Package.h @@ -20,6 +20,10 @@ public: template void AddPrimitive(T val) { + // Check if we are trying to add more than the package can fit. + if (m_MaxPacketSize < m_Offset + sizeof(T)) { + LOG_WARNING("Packet AddPrimitive(): You are trying to add more than we have allocated for!"); + } memcpy(m_Data + m_Offset, &val, sizeof(T)); m_Offset += sizeof(T); } @@ -28,7 +32,7 @@ public: T PopFrontPrimitive() { if (m_Offset < m_ReturnDataOffset + sizeof(T)) { - LOG_WARNING("Package PopFrontPrimitive(): You are trying to remove more than what exists in this package!"); + LOG_WARNING("Packet PopFrontPrimitive(): You are trying to remove more than what exists in this package!"); return -1; } T returnValue; @@ -51,6 +55,7 @@ private: char* m_Data; unsigned int m_ReturnDataOffset = 0; int m_Offset = 0; + unsigned int m_MaxPacketSize = 128; }; #endif \ No newline at end of file diff --git a/src/Engine/Network/Package.cpp b/src/Engine/Network/Package.cpp index 43c28837..be99e4a3 100644 --- a/src/Engine/Network/Package.cpp +++ b/src/Engine/Network/Package.cpp @@ -16,7 +16,11 @@ Package::Package(MessageType type, unsigned int& packageID) Package::Package(char* data, const int sizeOfPackage) { // Create message + + // Resize message + m_MaxPacketSize = sizeOfPackage; m_Data = new char[sizeOfPackage]; + // Copy data newly allocated memory memcpy(m_Data, data, sizeOfPackage); m_Offset = sizeOfPackage; } @@ -29,14 +33,18 @@ Package::~Package() void Package::AddString(std::string str) { // Message, add one extra byte for null terminator - memcpy(m_Data + m_Offset, str.data(), (str.size() + 1) * sizeof(char)); - m_Offset += (str.size() + 1) * sizeof(char); + int sizeOfString = str.size() + 1; + if (m_Offset + sizeOfString > m_MaxPacketSize) { + LOG_WARNING("Package::AddString(): Data size in packet exceeded maximum package size.\n"); + } + memcpy(m_Data + m_Offset, str.data(), sizeOfString * sizeof(char)); + m_Offset += sizeOfString * sizeof(char); } void Package::AddData(char * data, int sizeOfData) { - if (m_Offset + sizeOfData > 128) { - LOG_WARNING("Package::AddData(): Data size in package exceeded maximum package size.\n"); + if (m_Offset + sizeOfData > m_MaxPacketSize) { + LOG_WARNING("Package::AddData(): Data size in packet exceeded maximum package size.\n"); } memcpy(m_Data + m_Offset, data, sizeOfData); m_Offset += sizeOfData; @@ -46,7 +54,7 @@ std::string Package::PopFrontString() { std::string returnValue(m_Data + m_ReturnDataOffset); if (m_Offset < m_ReturnDataOffset + returnValue.size()){ - LOG_WARNING("Package PopFrontString(): Oh no! You are trying to remove things outside my memory kingdom"); + LOG_WARNING("packet PopFrontString(): Oh no! You are trying to remove things outside my memory kingdom"); return "PopFrontString Failed"; } // +1 for null terminator. @@ -56,6 +64,10 @@ std::string Package::PopFrontString() char * Package::PopData(int SizeOfData) { + if (m_Offset < m_ReturnDataOffset + SizeOfData) { + LOG_WARNING("packet PopData(): Oh no! You are trying to remove things outside my memory kingdom"); + return nullptr; + } unsigned int oldReturnDataOffset = m_ReturnDataOffset; m_ReturnDataOffset += SizeOfData; return (m_Data + oldReturnDataOffset);