diff --git a/src/network/NetworkMessage.h b/src/network/NetworkMessage.h index 42a924b..0138102 100644 --- a/src/network/NetworkMessage.h +++ b/src/network/NetworkMessage.h @@ -3,6 +3,9 @@ #include #include #include +#include + +#define VARINT_MAX_SIZE 5 namespace Feather::Network { @@ -45,7 +48,9 @@ namespace Feather::Network NetworkMessageCounter packetSizeCounter; packetSizeCounter.WriteVarInt(size); - m_data.reserve(size_t(size) + size_t(packetSizeCounter.GetSize())); + size_t initialSize = size_t(size) + size_t(packetSizeCounter.GetSize()); + + m_data.reserve(initialSize); WriteVarInt(size); } @@ -72,10 +77,13 @@ namespace Feather::Network std::memcpy(&m_data[oldSize], string, size_t(stringSize)); } - const uint8_t* GetData() const { return m_data.data(); } - size_t GetDataSize() const { return m_data.size(); } + inline virtual const uint8_t* GetData() const { return m_data.data(); } + inline virtual size_t GetDataSize() const { return m_data.size(); } + + protected: + // Used by subclasses with their own constructors + NetworkMessage() {} - private: std::vector m_data; }; @@ -89,4 +97,78 @@ namespace Feather::Network x(message, args...); \ return message; \ } + + // This class provides an alternative to NetworkMessageCounter, instead performing + // one count and prepending the size marker to the start of the buffer. + // + // To do this, you must call Finalize() before GetDataSize() or GetData() + class DynamicNetworkMessage : public NetworkMessage + { + public: + DynamicNetworkMessage(int32_t expectedSize) : NetworkMessage() + { + size_t initialSize = VARINT_MAX_SIZE + expectedSize; + + m_data.reserve(initialSize); + m_data.resize(VARINT_MAX_SIZE); + } + + void Finalize() + { + PrependVarIntSize(); + m_finalized = true; + } + + // Inlining and 'final' should ideally avert vtable lookups + // TOOD: Make these separate functions if vtable is not averted + + inline virtual const uint8_t* GetData() const final + { + // You must call Finalize() first + assert(m_finalized); + return m_data.data() + m_startOffset; + } + + inline virtual size_t GetDataSize() const final + { + // You must call Finalize() first + assert(m_finalized); + return m_data.size() - m_startOffset; + } + private: + int m_startOffset = 0; + bool m_finalized = false; + + void PrependVarIntSize() + { + uint8_t sizeBytes[VARINT_MAX_SIZE]; + + // compute the size of the packet excluding the size marker + int value = m_data.size() - VARINT_MAX_SIZE; + int index = 0; + + // perform WriteVarInt into the static buffer sizeBytes + do { + uint8_t temp = uint8_t(value & 0b01111111); + value >>= 7; + if (value != 0) + temp |= 0b10000000; + + sizeBytes[index++] = temp; + } while (value != 0 && index < VARINT_MAX_SIZE); + + // where to start writing the size marker + // if we used all 5 bytes, this is zero + // if we used 3 bytes, this would be 2 + int startIndex = VARINT_MAX_SIZE - index; + + // Copy sizeBytes into our actual buffer, but align it to the right + for (int i = 0; i < index; i++) + { + m_data[startIndex + i] = sizeBytes[i]; + } + + m_startOffset = startIndex; + } + }; } \ No newline at end of file