diff --git a/src/libs/ice/ICEByteBuffer.cpp b/src/libs/ice/ICEByteBuffer.cpp index 85b67d2c..679a54ba 100644 --- a/src/libs/ice/ICEByteBuffer.cpp +++ b/src/libs/ice/ICEByteBuffer.cpp @@ -497,12 +497,12 @@ size_t BufferReader::count() const } // -------------- BufferWriter ---------------------- -BufferWriter::BufferWriter(ByteBuffer &buffer) - :mData(buffer.mutableData()), mIndex(0) +BufferWriter::BufferWriter(ByteBuffer &buffer, size_t startOffset) + :mBuffer(&buffer), mData(buffer.mutableData() + startOffset), mBase(startOffset), mIndex(0), mCapacity(0) {} -BufferWriter::BufferWriter(void *output) - :mData(reinterpret_cast(output)), mIndex(0) +BufferWriter::BufferWriter(void *output, size_t capacity) + :mBuffer(nullptr), mData(reinterpret_cast(output)), mBase(0), mIndex(0), mCapacity(capacity) {} @@ -545,6 +545,30 @@ void BufferWriter::writeIp(const NetworkAddress& ip) void BufferWriter::writeBuffer(const void* dataPtr, size_t dataSize) { + if (dataSize == 0) + return; + + if (mBuffer) + { + // Growable mode: ensure the backing buffer is large enough, then refresh + // the cached pointer in case resize() reallocated the storage. + size_t needed = mBase + mIndex + dataSize; + if (needed > mBuffer->size()) + mBuffer->resize(needed); + mData = mBuffer->mutableData() + mBase; + } + else + { + // Bounded mode: never write past the caller-supplied capacity. + if (mIndex + dataSize > mCapacity) + { + assert(0 && "BufferWriter capacity exceeded"); + dataSize = mIndex < mCapacity ? mCapacity - mIndex : 0; + if (dataSize == 0) + return; + } + } + memmove(mData + mIndex, dataPtr, dataSize); mIndex += dataSize; } diff --git a/src/libs/ice/ICEByteBuffer.h b/src/libs/ice/ICEByteBuffer.h index 4e767d53..90801cad 100644 --- a/src/libs/ice/ICEByteBuffer.h +++ b/src/libs/ice/ICEByteBuffer.h @@ -172,12 +172,17 @@ namespace ice class BufferWriter { protected: - uint8_t* mData; - size_t mIndex; + ByteBuffer* mBuffer; // Non-null => growable: writes resize this buffer to fit + uint8_t* mData; // Cached write pointer (refreshed on growth in growable mode) + size_t mBase; // Start offset within mBuffer (growable mode) + size_t mIndex; // Bytes written since mBase + size_t mCapacity; // Usable capacity from mData (raw/bounded mode) public: - BufferWriter(void* output); - BufferWriter(ByteBuffer& buffer); + // Bounded writer over caller-owned memory. Writes past 'capacity' are refused. + BufferWriter(void* output, size_t capacity); + // Growable writer: the backing buffer is resized as needed so writes never overflow. + BufferWriter(ByteBuffer& buffer, size_t startOffset = 0); void writeUInt(uint32_t value); void writeUShort(uint16_t value); diff --git a/src/libs/ice/ICEStunMessage.cpp b/src/libs/ice/ICEStunMessage.cpp index 76d0d729..3d627455 100644 --- a/src/libs/ice/ICEStunMessage.cpp +++ b/src/libs/ice/ICEStunMessage.cpp @@ -216,8 +216,10 @@ void StunMessage::buildPacket(ByteBuffer& buffer, const std::string& password) for (size_t i=0; i<4; i++) bitstream.writeBit(bit(msgtype, 3-i)); - // Enqueue 2 bytes of length - now it is zero - BufferWriter stream(buffer.mutableData() + bitstream.count() / 8); + // Enqueue 2 bytes of length - now it is zero. + // Growable writer starting right after the bits already emitted; it resizes + // 'buffer' as attributes are appended so an oversized message cannot overflow. + BufferWriter stream(buffer, bitstream.count() / 8); stream.writeUShort(0); // Enqueue magic cookie value