1 files changed, 20 insertions, 10 deletions

diff --git a/src/packet/buffer.c b/src/packet/buffer.c
index 28dd73b..b609476 100644
--- a/src/packet/buffer.c
+++ b/src/packet/buffer.c
@@ -14,6 +14,7 @@ struct PacketBuffer {
 PacketBuffer* buffer_create(int capacity) {
     PacketBuffer* buffer = malloc(sizeof(PacketBuffer));
     buffer->arr = malloc(capacity);
+    memset(buffer->arr, 0, capacity);
     buffer->capacity = capacity;
     buffer->size = 0;
     buffer->index = 0;
@@ -30,7 +31,7 @@ void buffer_seek(PacketBuffer* buffer, int index) {
 }
 
 uint8_t buffer_read(PacketBuffer* buffer) {
-    if (buffer->index > buffer->size) {
+    if (buffer->index >= buffer->size) {
         return 0;
     }
     uint8_t data = buffer->arr[buffer->index];
@@ -53,7 +54,7 @@ uint32_t buffer_read_int(PacketBuffer* buffer) {
 }
 
 uint8_t buffer_get(PacketBuffer* buffer, int index) {
-    if (index > buffer->size) {
+    if (index >= buffer->size) {
         return 0;
     }
     uint8_t data = buffer->arr[index];
@@ -73,8 +74,12 @@ uint16_t buffer_get_size(PacketBuffer* buffer) {
 }
 
 static void write(uint8_t** buffer, uint8_t* size, uint8_t* capacity, uint8_t data) {
-    if (*size == *capacity) {
-        *capacity *= 2;
+    if (*size >= *capacity) {
+        if (*capacity >= 128) {
+            *capacity = 255;
+        } else {
+            *capacity *= 2;
+        }
         *buffer = realloc(*buffer, *capacity); 
     }
     (*buffer)[*size] = data;
@@ -144,7 +149,16 @@ void buffer_read_string(PacketBuffer* buffer, uint8_t** out) {
     buffer_read_n(buffer, out, len);
 }
 
+static void buffer_expand(PacketBuffer* buffer, int capacity) {
+    if (buffer->capacity >= capacity) return;
+
+    buffer->arr = realloc(buffer->arr, capacity);
+    memset(buffer->arr + buffer->capacity, 0, capacity - buffer->capacity);
+    buffer->capacity = capacity;
+}
+
 void buffer_read_n(PacketBuffer* buffer, uint8_t** out, uint8_t len) {
+    buffer_expand(buffer, buffer->index + len + 1);
     *out = malloc(len + 1);
     *out[0] = len;
     memcpy(*out + 1, buffer->arr + buffer->index, len);
@@ -152,7 +166,7 @@ void buffer_read_n(PacketBuffer* buffer, uint8_t** out, uint8_t len) {
 }
 
 void buffer_write(PacketBuffer* buffer, uint8_t data) {
-    if(buffer->index == buffer->capacity) {
+    if(buffer->index >= buffer->capacity) {
         buffer->capacity *= 2;
         buffer->arr = realloc(buffer->arr, buffer->capacity);
     }
@@ -205,11 +219,7 @@ void buffer_write_string(PacketBuffer* buffer, uint8_t* in) {
 }
 
 void buffer_write_n(PacketBuffer* buffer, uint8_t* in, int len) {
-    if (buffer->size + len >= buffer->capacity) {
-        buffer->capacity *= 2;
-        buffer->capacity += len;
-        buffer->arr = realloc(buffer->arr, buffer->capacity);
-    }
+    buffer_expand(buffer, buffer->index + len + 1);
     memcpy(buffer->arr + buffer->index, in, len);
     buffer->size += len;
     buffer->index += len;