feat: add rsync-compatible max-alloc limit

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>
@@ -139,7 +140,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

@@ -33,6 +33,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->use_delta = false;
   config->delta_block_size = DELTA_BLOCK_SIZE_DEFAULT;
@@ -134,7 +135,8 @@ 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->max_delete >= 0;
+         config->delta_max_file_size <= DELTA_MAX_FILE_SIZE && config->max_delete >= 0 &&
+         config->max_alloc > 0;
 }
 
 Config* config_create(void) {
@@ -218,11 +220,11 @@ void config_delete(Config* config) {
  * helper call order in config_send and config_receive unchanged when adding
  * fields. */
 static bool send_core_fields(int fd, const Config* c) {
-  return send_str(fd, c->version) && 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) && send_int(fd, c->use_metadata) &&
-         send_int(fd, c->compression_level) &&
+  return send_str(fd, c->version) && send_n_data(fd, &c->max_alloc, sizeof(c->max_alloc)) &&
+         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) &&
+         send_int(fd, c->use_metadata) && send_int(fd, c->compression_level) &&
          send_n_data(fd, &c->chunk_size, sizeof(c->chunk_size)) && send_int(fd, c->use_sendfile);
 }
 
@@ -258,6 +260,9 @@ static bool send_resume_options(int fd, const Config* c) {
 
 static bool receive_core_fields(int fd, Config* c) {
   int value;
+  if (!receive_n_data(fd, &c->max_alloc, sizeof(c->max_alloc)) || c->max_alloc == 0)
+    return false;
+  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)
@@ -334,6 +339,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

@@ -33,6 +33,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 use_delta;
   uint32_t delta_block_size;

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;
 
@@ -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

@@ -357,7 +357,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

@@ -55,7 +55,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;
@@ -122,7 +122,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/protocol.c

@@ -19,7 +19,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;
@@ -45,6 +46,7 @@ void io_set_fds(int read_fd, int write_fd) {
   legacy_io_session.write_fd = write_fd;
   legacy_io_session.ssl = NULL;
   legacy_io_session.total_allocated_bytes = 0;
+  legacy_io_session.max_alloc = DEFAULT_MAX_ALLOC;
   protocol_session_set_bwlimit(&legacy_io_session, global_bwlimit());
 }
 
@@ -54,9 +56,33 @@ 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;
   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(size_t size) {
+  const ProtocolSession* session = bound_session ? bound_session : &legacy_io_session;
+  return (unsigned long long)size <= session->max_alloc;
+}
+
+void* protocol_alloc(size_t size) {
+  if (!allocation_allowed(size))
+    return NULL;
+  return malloc(size);
+}
+
+void* protocol_realloc(void* ptr, size_t size) {
+  if (!allocation_allowed(size))
+    return NULL;
+  return realloc(ptr, size);
+}
+
 void protocol_session_bind(ProtocolSession* session) {
   bound_session = session;
 }
@@ -154,6 +180,7 @@ static ProtocolSession* legacy_session(int read_fd, int 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;
+    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());
@@ -342,7 +369,7 @@ char* protocol_receive_str(ProtocolSession* session) {
                 (unsigned long long)MAX_STRING_SIZE);
     return NULL;
   }
-  char* data = (char*)malloc(size + 1);
+  char* data = (char*)protocol_alloc(size + 1);
   if (data == NULL)
     return NULL;
   if (!protocol_receive_n_data(session, data, size)) {
@@ -385,6 +412,8 @@ 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) {
     log_message(LOG_LEVEL_ERROR, "Per-connection memory limit exceeded (%llu + %llu > %llu)",
@@ -392,7 +421,7 @@ Data* protocol_receive_data_limited(ProtocolSession* session, unsigned long long
                 (unsigned long long)MAX_CONNECTION_MEMORY);
     return NULL;
   }
-  void* data = malloc(allocation_size);
+  void* data = protocol_alloc(allocation_size);
   if (data == NULL)
     return NULL;
   if (!protocol_receive_n_data(session, data, (size_t)size)) {

src/shared/protocol.h

@@ -18,6 +18,7 @@
 #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)
 
 typedef struct ssl_st SSL;
 
@@ -36,6 +37,7 @@ typedef struct ProtocolSession {
   long long bw_last_refill_sec;
   long bw_last_refill_nsec;
   unsigned long long total_allocated_bytes;
+  unsigned long long max_alloc;
 } ProtocolSession;
 
 typedef int Status;
@@ -65,6 +67,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);
 bool protocol_send_n_data(ProtocolSession* session, const void* data, size_t data_size);
 bool protocol_receive_n_data(ProtocolSession* session, void* data, size_t data_size);
 bool protocol_send_str(ProtocolSession* session, const char* data);