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);