feat: add rsync-compatible max-alloc limit

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)) {