fix: address PR 212 review issues

src/shared/chunk.c

@@ -1,4 +1,5 @@
 #include <stddef.h>
+#include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -14,6 +15,7 @@
 
 /* Maximum individual file data size within a chunk (64 MB) */
 #define MAX_FILE_DATA_SIZE (64ULL * 1024 * 1024)
+#define MAX_CHUNK_FILES (1024 * 1024)
 
 Chunk* chunk_create(File** items, int element_count) {
   Chunk* chunk = (Chunk*)malloc(sizeof(Chunk));
@@ -92,6 +94,11 @@ Chunk* chunk_deserialize(Data* data, bool use_metadata) {
   size_t remaining_size = data->size;
 
   while (remaining_size > 0) {
+    if (files->size >= MAX_CHUNK_FILES) {
+      log_message(LOG_LEVEL_ERROR, "Chunk contains too many files");
+      array_list_delete(files);
+      return NULL;
+    }
     if (remaining_size < sizeof(size_t)) {
       log_message(LOG_LEVEL_ERROR, "Invalid chunk format: not enough data for path length");
       array_list_delete(files);
@@ -103,7 +110,7 @@ Chunk* chunk_deserialize(Data* data, bool use_metadata) {
     data_pointer += sizeof(size_t);
     remaining_size -= sizeof(size_t);
 
-    if (remaining_size < path_len) {
+    if (path_len > SIZE_MAX - 1 || remaining_size < path_len) {
       log_message(LOG_LEVEL_ERROR, "Invalid chunk format: not enough data for path");
       array_list_delete(files);
       return NULL;
@@ -122,10 +129,15 @@ Chunk* chunk_deserialize(Data* data, bool use_metadata) {
 
     File* file = file_create(path);
     free(path);
+    if (file == NULL) {
+      array_list_delete(files);
+      return NULL;
+    }
 
     if (use_metadata) {
       if (remaining_size < sizeof(int)) {
         log_message(LOG_LEVEL_ERROR, "Invalid chunk format: not enough data for metadata");
+        file_destroy(file);
         array_list_delete(files);
         return NULL;
       }
@@ -134,6 +146,7 @@ Chunk* chunk_deserialize(Data* data, bool use_metadata) {
       memcpy(&present_flag, data_pointer, sizeof(int));
       if (present_flag && remaining_size < sizeof(int) + FILE_METADATA_WIRE_SIZE) {
         log_message(LOG_LEVEL_ERROR, "Invalid chunk format: not enough data for metadata body");
+        file_destroy(file);
         array_list_delete(files);
         return NULL;
       }
@@ -141,10 +154,16 @@ Chunk* chunk_deserialize(Data* data, bool use_metadata) {
       remaining_size -= sizeof(int);
       if (file->metadata)
         remaining_size -= FILE_METADATA_WIRE_SIZE;
+      else if (present_flag) {
+        file_destroy(file);
+        array_list_delete(files);
+        return NULL;
+      }
     }
 
     if (remaining_size < sizeof(size_t)) {
       log_message(LOG_LEVEL_ERROR, "Invalid chunk format: not enough data for data size");
+      file_destroy(file);
       array_list_delete(files);
       return NULL;
     }
@@ -156,6 +175,7 @@ Chunk* chunk_deserialize(Data* data, bool use_metadata) {
 
     if (remaining_size < file_data_size) {
       log_message(LOG_LEVEL_ERROR, "Invalid chunk format: not enough data for file content");
+      file_destroy(file);
       array_list_delete(files);
       return NULL;
     }
@@ -164,29 +184,48 @@ Chunk* chunk_deserialize(Data* data, bool use_metadata) {
     if (file_data_size > MAX_FILE_DATA_SIZE) {
       log_message(LOG_LEVEL_ERROR, "File data size %zu exceeds maximum %llu", file_data_size,
                   (unsigned long long)MAX_FILE_DATA_SIZE);
+      file_destroy(file);
       array_list_delete(files);
       return NULL;
     }
 
-    void* file_data = malloc(file_data_size);
+    void* file_data = malloc(file_data_size > 0 ? file_data_size : 1);
     if (file_data == NULL) {
       perror("Could not allocate memory for file data");
+      file_destroy(file);
       array_list_delete(files);
       return NULL;
     }
     memcpy(file_data, data_pointer, file_data_size);
     data_destroy(file->data);
     file->data = data_create(file_data, file_data_size);
+    if (file->data == NULL) {
+      file_destroy(file);
+      array_list_delete(files);
+      return NULL;
+    }
     data_pointer += file_data_size;
     remaining_size -= file_data_size;
 
-    array_list_add(files, file);
+    if (!array_list_add(files, file)) {
+      file_destroy(file);
+      array_list_delete(files);
+      return NULL;
+    }
   }
 
   File** file_array = (File**)array_list_to_array(files);
+  if (file_array == NULL) {
+    array_list_delete(files);
+    return NULL;
+  }
   Chunk* chunk = chunk_create(file_array, files->size);
 
   free(file_array);
+  if (chunk == NULL) {
+    array_list_delete(files);
+    return NULL;
+  }
   files->item_destroyer = NULL;
   array_list_delete(files);
 

src/shared/file_store.c

@@ -130,9 +130,20 @@ bool file_store_write_secure(const char* path, const void* data, unsigned long l
         ok = file_restore_metadata_fd(fd, metadata);
     }
   } else {
-    char tmp[NAME_MAX];
+    int tmp_size = snprintf(NULL, 0, ".%s.tmp.%ld.%u", leaf, (long)getpid(), 99U);
+    if (tmp_size < 0) {
+      close(dirfd);
+      free(leaf);
+      return false;
+    }
+    char* tmp = malloc((size_t)tmp_size + 1);
+    if (!tmp) {
+      close(dirfd);
+      free(leaf);
+      return false;
+    }
     for (unsigned int i = 0; i < 100 && !ok; ++i) {
-      snprintf(tmp, sizeof(tmp), ".%s.tmp.%ld.%u", leaf, (long)getpid(), i);
+      snprintf(tmp, (size_t)tmp_size + 1, ".%s.tmp.%ld.%u", leaf, (long)getpid(), i);
       fd = openat(dirfd, tmp, O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC | O_NOFOLLOW, 0600);
       if (fd < 0)
         continue;
@@ -150,6 +161,7 @@ bool file_store_write_secure(const char* path, const void* data, unsigned long l
       if (!ok)
         unlinkat(dirfd, tmp, 0);
     }
+    free(tmp);
   }
   if (fd >= 0)
     close(fd);

src/shared/protocol.c

@@ -19,6 +19,7 @@ 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 unsigned long long io_bwlimit = 0;
 static long long bw_tokens = 0;
@@ -26,8 +27,6 @@ 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 = 0;
-
 void io_set_fds(int read_fd, int write_fd) {
   bound_session = NULL;
   io_read_fd = read_fd;
@@ -35,7 +34,11 @@ void io_set_fds(int read_fd, int 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;
+  legacy_io_session.read_fd = read_fd;
+  legacy_io_session.write_fd = write_fd;
+  legacy_io_session.ssl = NULL;
+  legacy_io_session.total_allocated_bytes = 0;
+  protocol_session_set_bwlimit(&legacy_io_session, io_bwlimit);
 }
 
 void protocol_session_init(ProtocolSession* session, int read_fd, int write_fd) {
@@ -126,32 +129,30 @@ SSL* io_get_ssl(void) {
   return io_ssl;
 }
 
-static ProtocolSession* legacy_session(void) {
-  static __thread ProtocolSession session;
+static ProtocolSession* legacy_session(int read_fd, int write_fd) {
   if (bound_session)
     return bound_session;
-  session.read_fd = io_read_fd;
-  session.write_fd = io_write_fd;
-  session.ssl = io_ssl;
-  session.bwlimit = io_bwlimit;
-  session.bw_tokens = (unsigned long long)(bw_tokens < 0 ? 0 : bw_tokens);
-  session.bw_last_refill_sec = bw_last_refill.tv_sec;
-  session.bw_last_refill_nsec = bw_last_refill.tv_nsec;
-  return &session;
+  int target_read_fd = io_read_fd != -1 ? io_read_fd : read_fd;
+  int target_write_fd = io_write_fd != -1 ? io_write_fd : write_fd;
+  if (legacy_io_session.read_fd != target_read_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;
+    protocol_session_set_bwlimit(&legacy_io_session, io_bwlimit);
+  } else if (legacy_io_session.bwlimit != io_bwlimit) {
+    protocol_session_set_bwlimit(&legacy_io_session, io_bwlimit);
+  }
+  legacy_io_session.ssl = io_ssl;
+  return &legacy_io_session;
 }
 
 bool send_n_data(int file_descriptor, const void* data, size_t data_size) {
-  ProtocolSession* session = legacy_session();
-  if (session->write_fd == -1)
-    session->write_fd = file_descriptor;
-  return protocol_send_n_data(session, data, data_size);
+  return protocol_send_n_data(legacy_session(-1, file_descriptor), data, data_size);
 }
 
 bool receive_n_data(int file_descriptor, void* data, size_t data_size) {
-  ProtocolSession* session = legacy_session();
-  if (session->read_fd == -1)
-    session->read_fd = file_descriptor;
-  return protocol_receive_n_data(session, data, data_size);
+  return protocol_receive_n_data(legacy_session(file_descriptor, -1), data, data_size);
 }
 
 static int deadline_remaining_ms(const struct timespec* deadline) {
@@ -403,34 +404,26 @@ bool protocol_receive_status(ProtocolSession* session, Status* status) {
 }
 
 bool send_str(int fd, const char* data) {
-  (void)fd;
-  return protocol_send_str(legacy_session(), data);
+  return protocol_send_str(legacy_session(-1, fd), data);
 }
 char* receive_str(int fd) {
-  (void)fd;
-  return protocol_receive_str(legacy_session());
+  return protocol_receive_str(legacy_session(fd, -1));
 }
 bool send_data(int fd, const Data* data) {
-  (void)fd;
-  return protocol_send_data(legacy_session(), data);
+  return protocol_send_data(legacy_session(-1, fd), data);
 }
 Data* receive_data(int fd) {
-  (void)fd;
-  return protocol_receive_data(legacy_session());
+  return protocol_receive_data(legacy_session(fd, -1));
 }
 bool send_int(int fd, int data) {
-  (void)fd;
-  return protocol_send_int(legacy_session(), data);
+  return protocol_send_int(legacy_session(-1, fd), data);
 }
 bool receive_int(int fd, int* data) {
-  (void)fd;
-  return protocol_receive_int(legacy_session(), data);
+  return protocol_receive_int(legacy_session(fd, -1), data);
 }
 bool send_status(int fd, Status status) {
-  (void)fd;
-  return protocol_send_status(legacy_session(), status);
+  return protocol_send_status(legacy_session(-1, fd), status);
 }
 bool receive_status(int fd, Status* status) {
-  (void)fd;
-  return protocol_receive_status(legacy_session(), status);
+  return protocol_receive_status(legacy_session(fd, -1), status);
 }