fix: harden remaining review findings

src/shared/protocol.c

@@ -13,6 +13,7 @@
 
 #define RECEIVE_TIMEOUT_SEC 60 /* 60 second per-message timeout */
 #define SEND_TIMEOUT_SEC 60
+#define MAX_CONNECTION_MEMORY (256ULL * 1024 * 1024) /* bounded cumulative receive budget */
 
 static __thread int io_read_fd = -1;
 static __thread int io_write_fd = -1;
@@ -24,12 +25,15 @@ static struct timespec bw_last_refill = {0, 0};
 static mtx_t bw_mutex;
 static once_flag bw_mutex_once = ONCE_FLAG_INIT;
 
+static __thread unsigned long long total_allocated_bytes;
+
 void io_set_fds(int read_fd, int write_fd) {
   io_read_fd = read_fd;
   io_write_fd = write_fd;
   /* A descriptor switch starts a new transport; never reuse a TLS object
      belonging to a previous connection or test pipe. */
   io_ssl = NULL;
+  total_allocated_bytes = 0;
 }
 
 static void bw_mutex_init(void) {
@@ -243,7 +247,8 @@ char* receive_str(int file_descriptor) {
   size_t size;
   if (!receive_n_data(file_descriptor, &size, sizeof(size_t)))
     return NULL;
-  if (size > MAX_STRING_SIZE || size > SIZE_MAX - 1) {
+  if (size > MAX_STRING_SIZE || size > SIZE_MAX - 1 ||
+      size + 1 > MAX_CONNECTION_MEMORY - total_allocated_bytes) {
     log_message(LOG_LEVEL_ERROR, "String size %zu exceeds maximum %llu", size,
                 (unsigned long long)MAX_STRING_SIZE);
     return NULL;
@@ -261,6 +266,7 @@ char* receive_str(int file_descriptor) {
     return NULL;
   }
   data[size] = '\0';
+  total_allocated_bytes += size + 1;
   log_message(LOG_LEVEL_DEBUG, "Received String: %s", data);
   return data;
 }
@@ -285,6 +291,10 @@ Data* receive_data(int file_descriptor) {
     return NULL;
   }
   size_t allocation_size = size == 0 ? 1 : (size_t)size;
+  if (allocation_size > MAX_CONNECTION_MEMORY - total_allocated_bytes) {
+    log_message(LOG_LEVEL_ERROR, "Per-connection memory limit exceeded");
+    return NULL;
+  }
   void* data = malloc(allocation_size);
   if (data == NULL)
     return NULL;
@@ -292,10 +302,12 @@ Data* receive_data(int file_descriptor) {
     free(data);
     return NULL;
   }
+  total_allocated_bytes += allocation_size;
   log_message(LOG_LEVEL_DEBUG, "Received %lld data", size);
   Data* result = data_create(data, (size_t)size);
   if (!result) {
     free(data);
+    total_allocated_bytes -= allocation_size;
   }
   return result;
 }