diff options
| -rw-r--r-- | source/ByteBuffer.cpp | 130 | ||||
| -rw-r--r-- | source/ByteBuffer.h | 22 |
2 files changed, 147 insertions, 5 deletions
diff --git a/source/ByteBuffer.cpp b/source/ByteBuffer.cpp index c82951271..78cc1385e 100644 --- a/source/ByteBuffer.cpp +++ b/source/ByteBuffer.cpp @@ -13,8 +13,54 @@ -#define NEEDBYTES(Num) if (!CanReadBytes(Num)) return false; -#define PUTBYTES(Num) if (!CanWriteBytes(Num)) return false; +// If a string sent over the protocol is larger than this, a warning is emitted to the console +#define MAX_STRING_SIZE (512 KiB) + +#define NEEDBYTES(Num) if (!CanReadBytes(Num)) return false; // Check if at least Num bytes can be read from the buffer, return false if not +#define PUTBYTES(Num) if (!CanWriteBytes(Num)) return false; // Check if at least Num bytes can be written to the buffer, return false if not + + + + + +#if 0 + +/// Self-test of the VarInt-reading and writing code +class cByteBufferSelfTest +{ +public: + cByteBufferSelfTest(void) + { + TestRead(); + TestWrite(); + } + + void TestRead(void) + { + cByteBuffer buf(50); + buf.Write("\x05\xac\x02\x00", 4); + UInt64 v1; + ASSERT(buf.ReadVarInt(v1) && (v1 == 5)); + UInt64 v2; + ASSERT(buf.ReadVarInt(v2) && (v2 == 300)); + UInt64 v3; + ASSERT(buf.ReadVarInt(v3) && (v3 == 0)); + } + + void TestWrite(void) + { + cByteBuffer buf(50); + buf.WriteVarInt(5); + buf.WriteVarInt(300); + buf.WriteVarInt(0); + AString All; + buf.ReadAll(All); + ASSERT(All.size() == 4); + ASSERT(memcmp(All.data(), "\x05\xac\x02\x00", All.size()) == 0); + } +} g_ByteBufferTest; + +#endif @@ -328,6 +374,48 @@ bool cByteBuffer::ReadBEUTF16String16(AString & a_Value) +bool cByteBuffer::ReadVarInt(UInt64 & a_Value) +{ + CHECK_THREAD; + CheckValid(); + UInt64 Value = 0; + int Shift = 0; + unsigned char b = 0; + do + { + NEEDBYTES(1); + ReadBuf(&b, 1); + Value = Value | (((Int64)(b & 0x7f)) << Shift); + Shift += 7; + } while ((b & 0x80) != 0); + a_Value = Value; + return true; +} + + + + + +bool cByteBuffer::ReadVarUTF8String(AString & a_Value) +{ + CHECK_THREAD; + CheckValid(); + UInt64 Size = 0; + if (!ReadVarInt(Size)) + { + return false; + } + if (Size > MAX_STRING_SIZE) + { + LOGWARNING("%s: String too large: %llu (%llu KiB)", __FUNCTION__, Size, Size / 1024); + } + return ReadString(a_Value, (int)Size); +} + + + + + bool cByteBuffer::WriteChar(char a_Value) { CHECK_THREAD; @@ -446,6 +534,44 @@ bool cByteBuffer::WriteBEUTF16String16(const AString & a_Value) +bool cByteBuffer::WriteVarInt(UInt64 a_Value) +{ + CHECK_THREAD; + CheckValid(); + + // A 64-bit integer can be encoded by at most 10 bytes: + unsigned char b[10]; + int idx = 0; + do + { + b[idx] = (a_Value & 0x7f) | ((a_Value > 0x7f) ? 0x80 : 0x00); + a_Value = a_Value >> 7; + idx++; + } while (a_Value > 0); + + return WriteBuf(b, idx); +} + + + + +bool cByteBuffer::WriteVarUTF8String(AString & a_Value) +{ + CHECK_THREAD; + CheckValid(); + PUTBYTES(a_Value.size() + 1); // This is a lower-bound on the bytes that will be actually written. Fail early. + bool res = WriteVarInt(a_Value.size()); + if (!res) + { + return false; + } + return WriteBuf(a_Value.data(), a_Value.size()); +} + + + + + bool cByteBuffer::ReadBuf(void * a_Buffer, int a_Count) { CHECK_THREAD; diff --git a/source/ByteBuffer.h b/source/ByteBuffer.h index 650eda5b0..eb5ce5910 100644 --- a/source/ByteBuffer.h +++ b/source/ByteBuffer.h @@ -57,8 +57,22 @@ public: bool ReadBEFloat (float & a_Value); bool ReadBEDouble (double & a_Value); bool ReadBool (bool & a_Value); - bool ReadBEUTF16String16(AString & a_Value); - + bool ReadBEUTF16String16(AString & a_Value); // string length as BE short, then string as UTF-16BE + bool ReadVarInt (UInt64 & a_Value); + bool ReadVarUTF8String (AString & a_Value); // string length as VarInt, then string as UTF-8 + + /// Reads VarInt, assigns it to anything that can be assigned from an UInt64 (unsigned int, short, char, Byte, double, ...) + template <typename T> bool ReadVarUInt(T & a_Value) + { + UInt64 v; + bool res = ReadVarInt(v); + if (res) + { + a_Value = v; + } + return res; + } + // Write the specified datatype; return true if successfully written bool WriteChar (char a_Value); bool WriteByte (unsigned char a_Value); @@ -68,7 +82,9 @@ public: bool WriteBEFloat (float a_Value); bool WriteBEDouble (double a_Value); bool WriteBool (bool a_Value); - bool WriteBEUTF16String16(const AString & a_Value); + bool WriteBEUTF16String16(const AString & a_Value); // string length as BE short, then string as UTF-16BE + bool WriteVarInt (UInt64 a_Value); + bool WriteVarUTF8String (AString & a_Value); // string length as VarInt, then string as UTF-8 /// Reads a_Count bytes into a_Buffer; returns true if successful bool ReadBuf(void * a_Buffer, int a_Count); |