Merge remote-tracking branch 'origin/feat/max-alloc' into dev

src/shared/array_list.c

@@ -1,17 +1,18 @@
 #include "log.h"
 #include "array_list.h"
+#include "protocol.h"
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
 
 ArrayList* array_list_create(void (*item_destroyer)(void* item)) {
-  ArrayList* list = (ArrayList*)malloc(sizeof(ArrayList));
+  ArrayList* list = (ArrayList*)protocol_alloc(sizeof(ArrayList));
   if (list == NULL) {
     log_perror("ERROR: Could not allocate memory for array list struct");
     return NULL;
   }
 
-  list->items = malloc(INITIAL_ARRAY_SIZE * sizeof(void*));
+  list->items = protocol_alloc(INITIAL_ARRAY_SIZE * sizeof(void*));
   if (list->items == NULL) {
     free(list);
     return NULL;
@@ -41,7 +42,7 @@ static bool array_list_extend(ArrayList* array_list) {
   int new_capacity = array_list->capacity * 2;
   if (new_capacity == 0)
     new_capacity = INITIAL_ARRAY_SIZE;
-  void* new_items = realloc(array_list->items, new_capacity * sizeof(void*));
+  void* new_items = protocol_realloc(array_list->items, new_capacity * sizeof(void*));
   if (new_items == NULL) {
     log_perror("ERROR: Could not reallocate memory for array list items");
     return false;
@@ -67,7 +68,7 @@ void** array_list_to_array(const ArrayList* array_list) {
   if (array_list == NULL) {
     return NULL;
   }
-  void** array = malloc(array_list->size * sizeof(void*));
+  void** array = protocol_alloc(array_list->size * sizeof(void*));
   if (array == NULL) {
     log_perror("Could not malloc space for array from array list!");
     return NULL;

src/shared/chunk.c

@@ -22,7 +22,7 @@
 Chunk* chunk_create(File** items, int element_count) {
   if (element_count < 0 || (element_count > 0 && items == NULL))
     return NULL;
-  Chunk* chunk = (Chunk*)malloc(sizeof(Chunk));
+  Chunk* chunk = (Chunk*)protocol_alloc(sizeof(Chunk));
   if (chunk == NULL) {
     log_perror("ERROR: Could not allocate memory for chunk structure");
     return NULL;
@@ -35,7 +35,7 @@ Chunk* chunk_create(File** items, int element_count) {
       free(chunk);
       return NULL;
     }
-    chunk->items = (File**)malloc((size_t)element_count * sizeof(File*));
+    chunk->items = (File**)protocol_alloc((size_t)element_count * sizeof(File*));
     if (chunk->items == NULL) {
       free(chunk);
       return NULL;
@@ -158,7 +158,7 @@ Chunk* chunk_deserialize(Data* data, bool use_metadata) {
       array_list_delete(files);
       return NULL;
     }
-    char* path = malloc(path_len + 1);
+    char* path = protocol_alloc(path_len + 1);
     if (path == NULL) {
       log_perror("Could not allocate memory for file path");
       array_list_delete(files);
@@ -245,7 +245,7 @@ Chunk* chunk_deserialize(Data* data, bool use_metadata) {
     }
 
     size_t allocation_size = file_data_size > 0 ? file_data_size : 1;
-    void* file_data = malloc(allocation_size);
+    void* file_data = protocol_alloc(allocation_size);
     if (file_data == NULL) {
       log_perror("Could not allocate memory for file data");
       file_destroy(file);

src/shared/compression.c

@@ -1,6 +1,7 @@
 #include "compression.h"
 #include "data.h"
 #include "log.h"
+#include "protocol.h"
 #include <stdlib.h>
 #include <limits.h>
 #include <stdint.h>
@@ -182,7 +183,7 @@ Data* data_decompress_limited(Data* compressed_data, size_t maximum_size) {
       buf_size *= 2;
       if (buf_size > hard_limit)
         buf_size = (size_t)hard_limit;
-      void* new_data = realloc(uncompressed_data->data, buf_size);
+      void* new_data = protocol_realloc(uncompressed_data->data, buf_size);
       if (!new_data) {
         log_message(LOG_LEVEL_ERROR, "Failed to grow decompression buffer");
         ZSTD_freeDCtx(dctx);

src/shared/config.c

@@ -37,6 +37,7 @@ static void config_set_defaults(Config* config) {
   config->include_count = 0;
   config->max_size = 0;
   config->min_size = 0;
+  config->max_alloc = DEFAULT_MAX_ALLOC;
   config->use_incremental = false;
   config->ignore_times = false;
   config->size_only = false;
@@ -155,9 +156,9 @@ static bool validate_received_config(const Config* config) {
          config->chunk_size > 0 && config->chunk_size <= MAX_CHUNK_SIZE &&
          config->delta_block_size >= DELTA_BLOCK_SIZE_MIN &&
          config->delta_block_size <= DELTA_BLOCK_SIZE_MAX &&
-         config->delta_max_file_size <= DELTA_MAX_FILE_SIZE && config->modify_window >= 0 &&
+config->delta_max_file_size <= DELTA_MAX_FILE_SIZE && config->modify_window >= 0 &&
          config->max_delete >= 0 && config->skip_compress_count >= 0 &&
-         config->skip_compress_count <= 10000 &&
+         config->skip_compress_count <= 10000 && config->max_alloc > 0 &&
          (!config->chmod_spec || !*config->chmod_spec ||
           chmod_apply(0, config->chmod_spec, &(mode_t){0}));
 }
@@ -252,6 +253,8 @@ static bool send_core_fields(int fd, const Config* c) {
   if (!send_str(fd, c->version) || !send_int(fd, c->eight_bit_output))
     return false;
   protocol_set_8_bit_output(c->eight_bit_output);
+  if (!send_n_data(fd, &c->max_alloc, sizeof(c->max_alloc)))
+    return false;
   return send_str(fd, c->send_directory) && send_str(fd, c->receive_root_directory) &&
          send_int(fd, c->save_to_disk) && send_int(fd, c->use_multithreading) &&
          send_int(fd, c->use_chunk_serialization) && send_int(fd, c->use_compression) &&
@@ -311,6 +314,11 @@ static bool receive_core_fields(int fd, Config* c) {
   if (!receive_wire_bool(fd, &c->eight_bit_output))
     return false;
   protocol_set_8_bit_output(c->eight_bit_output);
+  if (!receive_n_data(fd, &c->max_alloc, sizeof(c->max_alloc)) || c->max_alloc == 0)
+    return false;
+  if (c->max_alloc > MAX_SERVER_ALLOC)
+    c->max_alloc = MAX_SERVER_ALLOC;
+  protocol_session_set_max_alloc(NULL, c->max_alloc);
   c->send_directory = receive_str(fd);
   c->receive_root_directory = receive_str(fd);
   if (!c->send_directory || !c->receive_root_directory)
@@ -413,6 +421,7 @@ static bool receive_resume_options(int fd, Config* c) {
 }
 
 bool config_send(int file_descriptor, const Config* config) {
+  protocol_session_set_max_alloc(NULL, config->max_alloc);
   if (!send_core_fields(file_descriptor, config) || !send_delta_fields(file_descriptor, config) ||
       !send_file_options(file_descriptor, config) ||
       !send_selection_options(file_descriptor, config) ||

src/shared/config.h

@@ -36,6 +36,7 @@ typedef struct Config {
   int include_count;
   unsigned long long max_size;
   unsigned long long min_size;
+  unsigned long long max_alloc;
   bool use_incremental;
   bool ignore_times;
   bool size_only;
@@ -144,7 +145,7 @@ typedef struct Config {
   bool skip_compress_set;
 } Config;
 
-#define PROTOCOL_VERSION "2.3.0"
+#define PROTOCOL_VERSION "2.4.0"
 #define DEFAULT_CHUNK_SIZE (10 * 1024 * 1024)
 
 Config* config_create(void);

src/shared/data.c

@@ -1,11 +1,12 @@
 #include "data.h"
 #include "log.h"
+#include "protocol.h"
 #include <stdlib.h>
 
 Data* data_create_empty(size_t data_size) {
   /* malloc(0) is UB; allocate at least 1 byte but preserve requested size */
   size_t alloc_size = data_size > 0 ? data_size : 1;
-  void* data = malloc(alloc_size);
+  void* data = protocol_alloc(alloc_size);
   if (data == NULL) {
     log_message(LOG_LEVEL_ERROR, "Could not allocate memory for empty data");
     return NULL;
@@ -14,7 +15,7 @@ Data* data_create_empty(size_t data_size) {
 }
 
 Data* data_create_reserve(size_t size) {
-  Data* d = malloc(sizeof(Data));
+  Data* d = protocol_alloc(sizeof(Data));
   if (d == NULL) {
     log_message(LOG_LEVEL_ERROR, "Could not allocate memory for data");
     return NULL;
@@ -26,7 +27,7 @@ Data* data_create_reserve(size_t size) {
 }
 
 Data* data_create(void* data, size_t data_size) {
-  Data* new_data = malloc(sizeof(Data));
+  Data* new_data = protocol_alloc(sizeof(Data));
   if (new_data == NULL) {
     log_message(LOG_LEVEL_ERROR, "Could not allocate memory for data");
     free(data);

src/shared/delta.c

@@ -1,5 +1,6 @@
 #include "delta.h"
 #include "log.h"
+#include "protocol.h"
 #include <stdint.h>
 #include <limits.h>
 #include <stdlib.h>
@@ -43,7 +44,7 @@ DeltaSignature* delta_signature_create(const void* old_file_data, uint64_t old_f
 
   uint32_t block_count = (uint32_t)((old_file_size + block_size - 1) / block_size);
 
-  DeltaSignature* sig = malloc(sizeof(DeltaSignature));
+  DeltaSignature* sig = protocol_alloc(sizeof(DeltaSignature));
   if (!sig)
     return NULL;
 
@@ -54,7 +55,7 @@ DeltaSignature* delta_signature_create(const void* old_file_data, uint64_t old_f
     free(sig);
     return NULL;
   }
-  sig->blocks = malloc((size_t)block_count * sizeof(DeltaBlockSig));
+  sig->blocks = protocol_alloc((size_t)block_count * sizeof(DeltaBlockSig));
   if (!sig->blocks) {
     free(sig);
     return NULL;
@@ -82,7 +83,7 @@ Data* delta_signature_serialize(const DeltaSignature* sig) {
       total > SIZE_MAX)
     return NULL;
 
-  uint8_t* buf = malloc((size_t)total);
+  uint8_t* buf = protocol_alloc((size_t)total);
   if (!buf)
     return NULL;
 
@@ -111,7 +112,7 @@ DeltaSignature* delta_signature_deserialize(const Data* data) {
   const uint8_t* buf = (const uint8_t*)data->data;
   size_t pos = 0;
 
-  DeltaSignature* sig = malloc(sizeof(DeltaSignature));
+  DeltaSignature* sig = protocol_alloc(sizeof(DeltaSignature));
   if (!sig)
     return NULL;
 
@@ -149,7 +150,7 @@ DeltaSignature* delta_signature_deserialize(const Data* data) {
     free(sig);
     return NULL;
   }
-  sig->blocks = malloc((size_t)blocks_size);
+  sig->blocks = protocol_alloc((size_t)blocks_size);
   if (!sig->blocks) {
     free(sig);
     return NULL;
@@ -178,7 +179,7 @@ static bool ensure_capacity(DeltaInstruction** instrs, uint32_t* capacity, uint3
   if (*capacity > MAX_DELTA_INSTRUCTIONS / 2)
     return false;
   uint32_t new_cap = *capacity * 2;
-  DeltaInstruction* tmp = realloc(*instrs, (size_t)new_cap * sizeof(DeltaInstruction));
+  DeltaInstruction* tmp = protocol_realloc(*instrs, (size_t)new_cap * sizeof(DeltaInstruction));
   if (!tmp)
     return false;
   *instrs = tmp;
@@ -195,7 +196,7 @@ static bool flush_literal(DeltaInstruction** instrs, uint32_t* capacity, uint32_
   uint32_t lit_len = (uint32_t)(end - start);
   if (!ensure_capacity(instrs, capacity, *count))
     return false;
-  uint8_t* lit_data = malloc(lit_len);
+  uint8_t* lit_data = protocol_alloc(lit_len);
   if (!lit_data)
     return false;
   memcpy(lit_data, data + start, lit_len);
@@ -225,7 +226,7 @@ Delta* delta_compute(const void* new_file_data, uint64_t new_file_size, const De
 
   uint32_t capacity = 64;
   uint32_t count = 0;
-  DeltaInstruction* instrs = malloc((size_t)capacity * sizeof(DeltaInstruction));
+  DeltaInstruction* instrs = protocol_alloc((size_t)capacity * sizeof(DeltaInstruction));
   if (!instrs)
     return NULL;
 
@@ -309,7 +310,7 @@ Delta* delta_compute(const void* new_file_data, uint64_t new_file_size, const De
     }
   }
 
-  Delta* delta = malloc(sizeof(Delta));
+  Delta* delta = protocol_alloc(sizeof(Delta));
   if (!delta) {
     free_instructions(instrs, count);
     return NULL;
@@ -355,7 +356,7 @@ Data* delta_serialize(const Delta* delta) {
   if (delta->delta_size > UINT64_MAX - header_size || header_size + delta->delta_size > SIZE_MAX)
     return NULL;
   uint64_t total = header_size + delta->delta_size;
-  uint8_t* buf = malloc((size_t)total);
+  uint8_t* buf = protocol_alloc((size_t)total);
   if (!buf)
     return NULL;
 
@@ -395,7 +396,7 @@ Delta* delta_deserialize(const Data* data) {
   const uint8_t* buf = (const uint8_t*)data->data;
   size_t pos = 0;
 
-  Delta* delta = malloc(sizeof(Delta));
+  Delta* delta = protocol_alloc(sizeof(Delta));
   if (!delta)
     return NULL;
 
@@ -412,9 +413,10 @@ Delta* delta_deserialize(const Data* data) {
     return NULL;
   }
 
-  delta->instructions = delta->instruction_count == 0
-                            ? NULL
-                            : malloc((size_t)delta->instruction_count * sizeof(DeltaInstruction));
+  delta->instructions =
+      delta->instruction_count == 0
+          ? NULL
+          : protocol_alloc((size_t)delta->instruction_count * sizeof(DeltaInstruction));
   if (delta->instruction_count > 0 && !delta->instructions) {
     free(delta);
     return NULL;
@@ -465,7 +467,7 @@ Delta* delta_deserialize(const Data* data) {
         free(delta);
         return NULL;
       }
-      delta->instructions[i].literal.data = malloc(lit_len ? lit_len : 1);
+      delta->instructions[i].literal.data = protocol_alloc(lit_len ? lit_len : 1);
       if (!delta->instructions[i].literal.data) {
         log_message(LOG_LEVEL_ERROR, "Failed to allocate %u bytes for literal data", lit_len);
         free_instructions(delta->instructions, i);
@@ -492,7 +494,7 @@ void* delta_apply(const void* old_data, uint64_t old_size, const Delta* delta,
       delta->new_file_size > DELTA_MAX_FILE_SIZE || delta->new_file_size > SIZE_MAX)
     return NULL;
 
-  void* output = malloc(delta->new_file_size ? (size_t)delta->new_file_size : 1);
+  void* output = protocol_alloc(delta->new_file_size ? (size_t)delta->new_file_size : 1);
   if (!output)
     return NULL;
 

src/shared/file.c

@@ -14,6 +14,7 @@
 #include "log.h"
 #include "metadata.h"
 #include "utils.h"
+#include "protocol.h"
 
 static bool write_all(int fd, const void* data, unsigned long long size) {
   const unsigned char* p = data;
@@ -45,14 +46,14 @@ bool file_checksum(File* file, uint64_t* checksum) {
 File* file_create(const char* path) {
   if (!path)
     return NULL;
-  File* file = (File*)malloc(sizeof(File));
+  File* file = (File*)protocol_alloc(sizeof(File));
   if (file == NULL) {
     log_perror("ERROR: Could not allocate memory for file struct");
     return NULL;
   }
 
   size_t path_len = strlen(path);
-  file->path = (char*)malloc(path_len + 1);
+  file->path = (char*)protocol_alloc(path_len + 1);
   if (file->path == NULL) {
     free(file);
     return NULL;
@@ -85,7 +86,7 @@ void file_destroy(void* item) {
 }
 
 FileMetadata* file_metadata_create(const struct stat* stats) {
-  FileMetadata* m = malloc(sizeof(FileMetadata));
+  FileMetadata* m = protocol_alloc(sizeof(FileMetadata));
   if (m == NULL) {
     log_perror("ERROR: Could not allocate memory for file metadata");
     return NULL;
@@ -112,7 +113,7 @@ bool file_load_data(File* file) {
   if (file->data->data == NULL) {
     if (file->data->size == 0)
       return true;
-    file->data->data = malloc(file->data->size);
+    file->data->data = protocol_alloc(file->data->size);
     if (file->data->data == NULL) {
       log_perror("Could not allocate memory for file data");
       return false;

src/shared/file_receive.c

@@ -421,7 +421,7 @@ File* receive_incremental_check(int fd, const Config* config, bool* skipped) {
   unsigned long long old_size = has_old_file ? (unsigned long long)st.st_size : 0;
   void* old_data = NULL;
   if (has_old_file && old_size > 0 && old_size <= MAX_RECEIVE_FILE_SIZE && old_size <= SIZE_MAX) {
-    old_data = malloc((size_t)old_size);
+    old_data = protocol_alloc((size_t)old_size);
     if (old_data) {
       size_t got = 0;
       while (got < (size_t)old_size) {

src/shared/metadata.c

@@ -77,7 +77,7 @@ FileMetadata* metadata_from_buf(char** buf) {
     return NULL;
   if (!present)
     return NULL;
-  FileMetadata* m = malloc(sizeof(FileMetadata));
+  FileMetadata* m = protocol_alloc(sizeof(FileMetadata));
   if (m == NULL)
     return NULL;
   int32_t mode;
@@ -144,7 +144,7 @@ FileMetadata* metadata_receive(int file_descriptor, int* ok) {
       *ok = 0;
     return NULL;
   }
-  FileMetadata* m = malloc(sizeof(FileMetadata));
+  FileMetadata* m = protocol_alloc(sizeof(FileMetadata));
   if (m == NULL) {
     if (ok)
       *ok = 0;

src/shared/multiprocessing.c

@@ -30,6 +30,8 @@ PipelineContextSender* pipeline_context_sender_create(Config* config, Queue* que
   context->progress_bytes = 0;
   context->sender_done = false;
   atomic_init(&context->cancelled, false);
+  protocol_session_init(&context->allocation_session, -1, -1);
+  protocol_session_set_max_alloc(&context->allocation_session, config->max_alloc);
   int init = 0;
   if (mtx_init(&context->mutex_scanner, mtx_plain) != thrd_success)
     goto fail;
@@ -182,6 +184,7 @@ int receive_thread(void* pipeline_context) {
 
 int write_thread(void* pipeline_context) {
   PipelineContextReceiver* context = (PipelineContextReceiver*)pipeline_context;
+  protocol_session_bind(&context->session);
   mtx_lock(&context->mutex);
   bool save_to_disk = context->config->save_to_disk;
   char* root_directory = str_dup(context->config->receive_root_directory);
@@ -193,6 +196,7 @@ int write_thread(void* pipeline_context) {
     cnd_broadcast(&context->condition_not_full);
     cnd_broadcast(&context->condition_not_empty);
     mtx_unlock(&context->mutex);
+    protocol_session_unbind();
     return thrd_error;
   }
 
@@ -202,6 +206,7 @@ int write_thread(void* pipeline_context) {
                                     &context->condition_not_full, &context->receiver_done);
     if (file == NULL) {
       free(root_directory);
+      protocol_session_unbind();
       return thrd_success;
     }
     if (save_to_disk && !file_save_to_disk(root_directory, file, context->config)) {
@@ -213,6 +218,7 @@ int write_thread(void* pipeline_context) {
       cnd_broadcast(&context->condition_not_empty);
       mtx_unlock(&context->mutex);
       free(root_directory);
+      protocol_session_unbind();
       return thrd_error;
     }
     file_destroy(file);

src/shared/multiprocessing.h

@@ -29,6 +29,7 @@ typedef struct {
   unsigned long long progress_bytes;
   bool sender_done;
   atomic_bool cancelled;
+  ProtocolSession allocation_session;
 } PipelineContextSender;
 
 typedef struct PipelineContextReceiver {

src/shared/protocol.c

@@ -20,7 +20,8 @@ static __thread int io_read_fd = -1;
 static __thread int io_write_fd = -1;
 static __thread SSL* io_ssl;
 static __thread ProtocolSession* bound_session;
-static __thread ProtocolSession legacy_io_session = {.read_fd = -1, .write_fd = -1};
+static __thread ProtocolSession legacy_io_session = {
+    .read_fd = -1, .write_fd = -1, .max_alloc = DEFAULT_MAX_ALLOC};
 
 static unsigned long long io_bwlimit = 0;
 static mtx_t bw_mutex;
@@ -28,12 +29,30 @@ static once_flag bw_mutex_once = ONCE_FLAG_INIT;
 
 static unsigned long long global_bwlimit(void);
 
+static bool protocol_reserve_memory(ProtocolSession* session, size_t charge) {
+  unsigned long long allocated = atomic_load(&session->total_allocated_bytes);
+  while (true) {
+    if (allocated > MAX_CONNECTION_MEMORY ||
+        (unsigned long long)charge > MAX_CONNECTION_MEMORY - allocated)
+      return false;
+    if (atomic_compare_exchange_weak(&session->total_allocated_bytes, &allocated,
+                                     allocated + (unsigned long long)charge))
+      return true;
+  }
+}
+
+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;
+    if (atomic_compare_exchange_weak(&session->total_allocated_bytes, &allocated, remaining))
+      break;
+  }
+}
+
 void protocol_release_memory(size_t charge) {
   ProtocolSession* session = bound_session ? bound_session : &legacy_io_session;
-  if ((unsigned long long)charge >= session->total_allocated_bytes)
-    session->total_allocated_bytes = 0;
-  else
-    session->total_allocated_bytes -= charge;
+  protocol_release_memory_for_session(session, charge);
 }
 void io_set_fds(int read_fd, int write_fd) {
   bound_session = NULL;
@@ -45,8 +64,9 @@ void io_set_fds(int read_fd, int write_fd) {
   legacy_io_session.read_fd = read_fd;
   legacy_io_session.write_fd = write_fd;
   legacy_io_session.ssl = NULL;
-  legacy_io_session.eight_bit_output = false;
-  legacy_io_session.total_allocated_bytes = 0;
+legacy_io_session.eight_bit_output = false;
+  atomic_store(&legacy_io_session.total_allocated_bytes, 0);
+  legacy_io_session.max_alloc = DEFAULT_MAX_ALLOC;
   protocol_session_set_bwlimit(&legacy_io_session, global_bwlimit());
 }
 
@@ -56,9 +76,43 @@ void protocol_session_init(ProtocolSession* session, int read_fd, int write_fd)
   memset(session, 0, sizeof(*session));
   session->read_fd = read_fd;
   session->write_fd = write_fd;
+  session->max_alloc = DEFAULT_MAX_ALLOC;
+  atomic_init(&session->total_allocated_bytes, 0);
   protocol_session_set_bwlimit(session, global_bwlimit());
 }
 
+void protocol_session_set_max_alloc(ProtocolSession* session, unsigned long long max_alloc) {
+  if (!session)
+    session = bound_session ? bound_session : &legacy_io_session;
+  session->max_alloc = max_alloc;
+}
+
+static bool allocation_allowed(const ProtocolSession* session, size_t size) {
+  return (unsigned long long)size <= session->max_alloc;
+}
+
+static void* protocol_alloc_for_session(const ProtocolSession* session, size_t size) {
+  if (!allocation_allowed(session, size))
+    return NULL;
+  return malloc(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;
   log_set_8_bit_output(session && session->eight_bit_output);
@@ -169,7 +223,8 @@ static ProtocolSession* legacy_session(int read_fd, int write_fd) {
       legacy_io_session.write_fd != target_write_fd) {
     legacy_io_session.read_fd = target_read_fd;
     legacy_io_session.write_fd = target_write_fd;
-    legacy_io_session.total_allocated_bytes = 0;
+    atomic_store(&legacy_io_session.total_allocated_bytes, 0);
+    legacy_io_session.max_alloc = DEFAULT_MAX_ALLOC;
     protocol_session_set_bwlimit(&legacy_io_session, global_bwlimit());
   } else if (legacy_io_session.bwlimit != global_bwlimit()) {
     protocol_session_set_bwlimit(&legacy_io_session, global_bwlimit());
@@ -352,13 +407,12 @@ char* protocol_receive_str(ProtocolSession* session) {
   size_t size;
   if (!protocol_receive_n_data(session, &size, sizeof(size_t)))
     return NULL;
-  if (size > MAX_STRING_SIZE || size > SIZE_MAX - 1 ||
-      size + 1 > MAX_CONNECTION_MEMORY - session->total_allocated_bytes) {
+  if (size > MAX_STRING_SIZE || size > SIZE_MAX - 1) {
     log_message(LOG_LEVEL_ERROR, "String size %zu exceeds maximum %llu", size,
                 (unsigned long long)MAX_STRING_SIZE);
     return NULL;
   }
-  char* data = (char*)malloc(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)) {
@@ -371,7 +425,6 @@ char* protocol_receive_str(ProtocolSession* session) {
     return NULL;
   }
   data[size] = '\0';
-  session->total_allocated_bytes += size + 1;
   log_debug_message(LOG_DEBUG_PROTO, "Received String: %s", data);
   return data;
 }
@@ -401,25 +454,29 @@ Data* protocol_receive_data_limited(ProtocolSession* session, unsigned long long
                 (unsigned long long)MAX_DATA_PAYLOAD_SIZE);
     return NULL;
   }
+  if (size > SIZE_MAX)
+    return NULL;
   size_t allocation_size = size == 0 ? 1 : (size_t)size;
-  if (allocation_size > MAX_CONNECTION_MEMORY - session->total_allocated_bytes) {
+  if (!protocol_reserve_memory(session, allocation_size)) {
     log_message(LOG_LEVEL_ERROR, "Per-connection memory limit exceeded (%llu + %llu > %llu)",
-                (unsigned long long)session->total_allocated_bytes, size,
+                (unsigned long long)atomic_load(&session->total_allocated_bytes), size,
                 (unsigned long long)MAX_CONNECTION_MEMORY);
     return NULL;
   }
-  void* data = malloc(allocation_size);
-  if (data == NULL)
-    return NULL;
-  if (!protocol_receive_n_data(session, data, (size_t)size)) {
-    free(data);
+  void* data = protocol_alloc_for_session(session, allocation_size);
+  if (data == NULL) {
+    protocol_release_memory_for_session(session, allocation_size);
     return NULL;
   }
-  session->total_allocated_bytes += allocation_size;
-  log_debug_message(LOG_DEBUG_PROTO, "Received %lld data", size);
+  if (!protocol_receive_n_data(session, data, (size_t)size)) {
+    free(data);
+    protocol_release_memory_for_session(session, allocation_size);
+    return NULL;
+  }
+log_debug_message(LOG_DEBUG_PROTO, "Received %lld data", size);
   Data* result = data_create(data, (size_t)size);
   if (!result) {
-    session->total_allocated_bytes -= allocation_size;
+    protocol_release_memory_for_session(session, allocation_size);
     return NULL;
   }
   result->protocol_charge = allocation_size;

src/shared/protocol.h

@@ -4,6 +4,7 @@
 #include "data.h"
 #include <stdbool.h>
 #include <stddef.h>
+#include <stdatomic.h>
 
 /* Maximum allowed string size for receive_str (64 KB) */
 #define MAX_STRING_SIZE (64 * 1024)
@@ -18,6 +19,9 @@
 #define MAX_MANIFEST_ENTRIES (1024 * 1024)
 /* Aggregate bytes retained by one received deletion manifest. */
 #define MAX_MANIFEST_BYTES (16ULL * 1024 * 1024)
+#define DEFAULT_MAX_ALLOC (1ULL * 1024 * 1024 * 1024)
+/* Server policy ceiling for a client-provided allocation limit. */
+#define MAX_SERVER_ALLOC (256ULL * 1024 * 1024)
 
 typedef struct ssl_st SSL;
 
@@ -35,8 +39,9 @@ typedef struct ProtocolSession {
   long long bw_tokens;
   long long bw_last_refill_sec;
   long bw_last_refill_nsec;
-  unsigned long long total_allocated_bytes;
+  atomic_ullong total_allocated_bytes;
   bool eight_bit_output;
+  unsigned long long max_alloc;
 } ProtocolSession;
 
 typedef int Status;
@@ -66,6 +71,9 @@ void protocol_session_bind(ProtocolSession* session);
 void protocol_session_unbind(void);
 void protocol_session_set_ssl(ProtocolSession* session, SSL* ssl);
 void protocol_session_set_bwlimit(ProtocolSession* session, unsigned long long bytes_per_sec);
+void protocol_session_set_max_alloc(ProtocolSession* session, unsigned long long max_alloc);
+void* protocol_alloc(size_t size);
+void* protocol_realloc(void* ptr, size_t size);
 void protocol_session_set_8_bit_output(ProtocolSession* session, bool enabled);
 void protocol_set_8_bit_output(bool enabled);
 bool protocol_send_n_data(ProtocolSession* session, const void* data, size_t data_size);