fix: honor explicit protocol allocation sessions

src/shared/protocol.c

@@ -40,8 +40,7 @@ static bool protocol_reserve_memory(ProtocolSession* session, size_t charge) {
   }
 }
 
-void protocol_release_memory(size_t charge) {
-  ProtocolSession* session = bound_session ? bound_session : &legacy_io_session;
+static void protocol_release_memory_for_session(ProtocolSession* session, size_t charge) {
   unsigned long long allocated = atomic_load(&session->total_allocated_bytes);
   while (true) {
     unsigned long long remaining = (unsigned long long)charge >= allocated ? 0 : allocated - charge;
@@ -49,6 +48,11 @@ void protocol_release_memory(size_t charge) {
       break;
   }
 }
+
+void protocol_release_memory(size_t charge) {
+  ProtocolSession* session = bound_session ? bound_session : &legacy_io_session;
+  protocol_release_memory_for_session(session, charge);
+}
 void io_set_fds(int read_fd, int write_fd) {
   bound_session = NULL;
   io_read_fd = read_fd;
@@ -81,23 +85,32 @@ void protocol_session_set_max_alloc(ProtocolSession* session, unsigned long long
   session->max_alloc = max_alloc;
 }
 
-static bool allocation_allowed(size_t size) {
-  const ProtocolSession* session = bound_session ? bound_session : &legacy_io_session;
+static bool allocation_allowed(const ProtocolSession* session, size_t size) {
   return (unsigned long long)size <= session->max_alloc;
 }
 
-void* protocol_alloc(size_t size) {
-  if (!allocation_allowed(size))
+static void* protocol_alloc_for_session(const ProtocolSession* session, size_t size) {
+  if (!allocation_allowed(session, size))
     return NULL;
   return malloc(size);
 }
 
-void* protocol_realloc(void* ptr, size_t size) {
-  if (!allocation_allowed(size))
+static void* protocol_realloc_for_session(const ProtocolSession* session, void* ptr, size_t size) {
+  if (!allocation_allowed(session, size))
     return NULL;
   return realloc(ptr, size);
 }
 
+void* protocol_alloc(size_t size) {
+  const ProtocolSession* session = bound_session ? bound_session : &legacy_io_session;
+  return protocol_alloc_for_session(session, size);
+}
+
+void* protocol_realloc(void* ptr, size_t size) {
+  const ProtocolSession* session = bound_session ? bound_session : &legacy_io_session;
+  return protocol_realloc_for_session(session, ptr, size);
+}
+
 void protocol_session_bind(ProtocolSession* session) {
   bound_session = session;
 }
@@ -383,7 +396,7 @@ char* protocol_receive_str(ProtocolSession* session) {
                 (unsigned long long)MAX_STRING_SIZE);
     return NULL;
   }
-  char* data = (char*)protocol_alloc(size + 1);
+  char* data = (char*)protocol_alloc_for_session(session, size + 1);
   if (data == NULL)
     return NULL;
   if (!protocol_receive_n_data(session, data, size)) {
@@ -434,20 +447,20 @@ Data* protocol_receive_data_limited(ProtocolSession* session, unsigned long long
                 (unsigned long long)MAX_CONNECTION_MEMORY);
     return NULL;
   }
-  void* data = protocol_alloc(allocation_size);
+  void* data = protocol_alloc_for_session(session, allocation_size);
   if (data == NULL) {
-    protocol_release_memory(allocation_size);
+    protocol_release_memory_for_session(session, allocation_size);
     return NULL;
   }
   if (!protocol_receive_n_data(session, data, (size_t)size)) {
     free(data);
-    protocol_release_memory(allocation_size);
+    protocol_release_memory_for_session(session, allocation_size);
     return NULL;
   }
   log_message(LOG_LEVEL_DEBUG, "Received %lld data", size);
   Data* result = data_create(data, (size_t)size);
   if (!result) {
-    protocol_release_memory(allocation_size);
+    protocol_release_memory_for_session(session, allocation_size);
     return NULL;
   }
   result->protocol_charge = allocation_size;