From 3397c93362a99fd2f296bb3d3ebf6c418eb27f52 Mon Sep 17 00:00:00 2001
From: RD42 <42702181+dashr9230@users.noreply.github.com>
Date: Mon, 12 Feb 2024 20:33:38 +0800
Subject: [PATCH] [raknet] Implement `BitStream::SetNumberOfBitsAllocated()`

---
 raknet/BitStream.cpp | 9 +++++++++
 raknet/BitStream.h   | 5 +++++
 2 files changed, 14 insertions(+)

diff --git a/raknet/BitStream.cpp b/raknet/BitStream.cpp
index 2eade3a..4e1c8ac 100644
--- a/raknet/BitStream.cpp
+++ b/raknet/BitStream.cpp
@@ -110,6 +110,15 @@ BitStream::BitStream( char* _dataC, unsigned int lengthInBytes, bool _copyData )
 		data = ( unsigned char* ) _data;
 }
 
+// Use this if you pass a pointer copy to the constructor (_copyData==false) and want to overallocate to prevent reallocation
+void BitStream::SetNumberOfBitsAllocated( const unsigned int lengthInBits )
+{
+#ifdef _DEBUG
+	assert( lengthInBits >= ( unsigned int ) numberOfBitsAllocated );
+#endif
+	numberOfBitsAllocated = lengthInBits;
+}
+
 BitStream::~BitStream()
 {
 	if ( copyData && numberOfBitsAllocated > BITSTREAM_STACK_ALLOCATION_SIZE << 3)
diff --git a/raknet/BitStream.h b/raknet/BitStream.h
index cf9bbe2..b5e3390 100644
--- a/raknet/BitStream.h
+++ b/raknet/BitStream.h
@@ -5,6 +5,7 @@
 
 #include "Export.h"
 #include "NetworkTypes.h"
+#include <assert.h>
 
 /// Arbitrary size, just picking something likely to be larger than most packets
 #define BITSTREAM_STACK_ALLOCATION_SIZE 256
@@ -46,6 +47,10 @@ namespace RakNet
 		/// Destructor 
 		~BitStream();
 
+		/// Use this if you pass a pointer copy to the constructor
+		/// *(_copyData==false) and want to overallocate to prevent
+		/// *reallocation
+		void SetNumberOfBitsAllocated( const unsigned int lengthInBits );
 
 
 	private: