3 files changed, 125 insertions, 25 deletions

diff --git a/analyzer/utils/memory.c b/analyzer/utils/memory.c
new file mode 100644
index 0000000..0a43bd3
--- /dev/null
+++ b/analyzer/utils/memory.c
@@ -0,0 +1,62 @@
+#include "memory.h"
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <unistd.h>
+
+void *safe_malloc(size_t size) {
+    for (size_t i = 0; i < MAX_ALLOC_ATTEMPTS; i++) {
+        void *p = malloc(size);
+        if (p)
+            return p;
+        if (i < MAX_ALLOC_ATTEMPTS - 1) {
+            printf("Failed malloc. Retrying in 1s...\n");
+            sleep(1);
+        }
+    }
+
+    printf("ERROR: Failed malloc %d times. Exiting...", MAX_ALLOC_ATTEMPTS);
+    exit(EXIT_FAILURE);
+}
+
+void *safe_realloc(void *ptr, size_t size) {
+    if (!ptr)
+        return safe_malloc(size);
+
+    for (size_t i = 0; i < MAX_ALLOC_ATTEMPTS; i++) {
+        void *p = realloc(ptr, size);
+        if (p)
+            return p;
+        if (i < MAX_ALLOC_ATTEMPTS - 1) {
+            printf("Failed realloc. Retrying in 1s...\n");
+            sleep(1);
+        }
+    }
+
+    printf("ERROR: Failed realloc %d times. Exiting...", MAX_ALLOC_ATTEMPTS);
+    exit(EXIT_FAILURE);
+}
+
+// expand an array
+// returns pointer to the expanded array
+void *array_expand(void *arr, size_t el_size, size_t *cap, size_t to_fit) {
+    while (*cap < to_fit)
+        *cap *= 2;
+    return safe_realloc(arr, *cap * el_size);
+}
+
+void *array_append(void *arr, size_t el_size, size_t *len, size_t *cap, void *val) {
+    if (*len >= *cap)
+        arr = array_expand(arr, el_size, cap, *len + 1);
+    memcpy((char *)arr + (*len)++ * el_size, val, el_size);
+    return arr;
+}
+
+void *array_concat(void *arr1, size_t el_size, size_t *len1, size_t *cap, void *arr2, size_t len2) {
+    if (*len1 + len2 > *cap)
+        arr1 = array_expand(arr1, el_size, cap, *len1 + len2);
+    memcpy((char *)arr1 + *len1 * el_size, arr2, len2 * el_size);
+    *len1 += len2;
+    return arr1;
+}
diff --git a/analyzer/utils/memory.h b/analyzer/utils/memory.h
index 1de58b9..0973623 100644
--- a/analyzer/utils/memory.h
+++ b/analyzer/utils/memory.h
@@ -1,4 +1,46 @@
 // utils/memory.h
 // helper functions for dealing with memory management
 
+#include <stddef.h>
+
 #define ARR_LEN(arr) (sizeof(arr) / sizeof(arr[0]))
+
+#define MAX_ALLOC_ATTEMPTS 3
+
+// tries malloc MAX_ALLOC_ATTEMPTS times, if all fail - exits the program
+// params:
+//      - size - count of bytes to allocate
+// returns:
+//      pointer to the allocated memory
+extern void *safe_malloc(size_t size);
+
+// tries realloc MAX_ALLOC_ATTEMPTS times, if all fail - exits the program
+// params:
+//      - ptr - pointer to the memory to reallocate
+//      - size - new size of the memory
+// returns:
+//      pointer to the reallocated memory
+extern void *safe_realloc(void *ptr, size_t size);
+
+// append to array, expand if necessary
+// params:
+//      - arr - array to append to
+//      - el_size - size of an array element
+//      - len - pointer to the length of the array
+//      - cap - pointer to the capacity of the array
+//      - val - pointer to the value to append to the array
+// returns:
+//      new pointer to the array (different if it was realloced)
+extern void *array_append(void *arr, size_t el_size, size_t *len, size_t *cap, void *val);
+
+// concatenate 2 arrays, expand if necessary
+// params:
+//      - arr1 - first array (destination)
+//      - el_size - size of an array element
+//      - len1 - pointer to the length of arr1
+//      - cap - pointer to the capacity of arr1
+//      - arr2 - second array
+//      - len2 - length of arr2
+// returns:
+//      new pointer to the array (different if it was realloced)
+extern void *array_concat(void *arr1, size_t el_size, size_t *len1, size_t *cap, void *arr2, size_t len2);
diff --git a/analyzer/utils/tcp.c b/analyzer/utils/tcp.c
index b925ecf..6309e64 100644
--- a/analyzer/utils/tcp.c
+++ b/analyzer/utils/tcp.c
@@ -6,15 +6,19 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "utils/memory.h"
 #include "utils/tcp_seq_math.h"
 
 // initalize default TCP segment
 TCPSegment new_segment() {
-    TCPSegment segment;
-    segment.payload = NULL;
-    segment.len = 0;
-    segment.seq = 0;
-    segment.port = 0;
+    // clang-format off
+    TCPSegment segment = {
+        .payload = NULL,
+        .len = 0,
+        .seq = 0,
+        .port = 0
+    };
+    // clang-format on
     return segment;
 }
 
@@ -61,19 +65,22 @@ TCPSegment get_segment_from_packet(const byte_t *packet, const struct pcap_pkthd
     // get payload
     int total_header_len = eth_header_len + ip_header_len + tcp_header_len;
     segment.len = packet_data->caplen - total_header_len;
-    segment.payload = malloc(segment.len);
+    segment.payload = safe_malloc(segment.len);
     memcpy(segment.payload, packet + total_header_len, segment.len);
     return segment;
 }
 
 TCPStream new_stream() {
-    TCPStream stream;
-    stream.data = malloc(INIT_STREAM_CAPACITY);
-    stream.len = 0;
-    stream.capacity = INIT_STREAM_CAPACITY;
-    stream.seq = 0;
-    stream.port = 0;
-    stream.pending = NULL;
+    // clang-format off
+    TCPStream stream = {
+        .data = safe_malloc(INIT_STREAM_CAPACITY),
+        .len = 0,
+        .capacity = INIT_STREAM_CAPACITY,
+        .seq = 0,
+        .port = 0,
+        .pending = NULL
+    };
+    // clang-format on
     return stream;
 }
 
@@ -107,20 +114,9 @@ int in_acceptable_range(TCPStream stream, TCPSegment segment) {
     return seq_in_range_ei(segment.seq, stream_expected, max_acceptable_seq);
 }
 
-// expand stream capacity exponentially to fit new_len
-void expand_stream_capacity(TCPStream *stream, size_t new_len) {
-    while (stream->capacity < new_len)
-        stream->capacity <<= 1;
-    stream->data = realloc(stream->data, stream->capacity);
-}
-
 // add segment directly to stream
 void add_segment(TCPStream *stream, TCPSegment segment, size_t data_start) {
-    size_t new_stream_len = stream->len + segment.len - data_start;
-    if (new_stream_len > stream->capacity)
-        expand_stream_capacity(stream, new_stream_len);
-    memcpy(stream->data + stream->len, segment.payload + data_start, segment.len - data_start);
-    stream->len = new_stream_len;
+    stream->data = array_concat(stream->data, 1, &stream->len, &stream->capacity, segment.payload + data_start, segment.len - data_start);
 }
 
 // add pending segments to the stream