fix: memory leak, unitialized memory access, FILE pointer leak bugs

This commit fixes numerous general code bugs, improving reliability and consistency of ReZygisk.

loader/src/common/daemon.c

@@ -108,43 +108,32 @@ void rezygiskd_get_info(struct rezygisk_info *info) {
 
   read_uint32_t(fd, (uint32_t *)&info->pid);
 
-  read_size_t(fd, &info->modules->modules_count);
+  read_size_t(fd, &info->modules.modules_count);
-  if (info->modules->modules_count == 0) {
+  if (info->modules.modules_count == 0) {
-    info->modules->modules = NULL;
+    info->modules.modules = NULL;
 
     close(fd);
 
     return;
   }
 
-  info->modules->modules = (char **)malloc(sizeof(char *) * info->modules->modules_count);
+  info->modules.modules = (char **)malloc(sizeof(char *) * info->modules.modules_count);
-  if (info->modules->modules == NULL) {
+  if (!info->modules.modules) {
     PLOGE("allocating modules name memory");
 
-    free(info->modules);
+    info->modules.modules_count = 0;
-    info->modules = NULL;
-    info->modules->modules_count = 0;
 
     close(fd);
 
     return;
   }
 
-  for (size_t i = 0; i < info->modules->modules_count; i++) {
+  for (size_t i = 0; i < info->modules.modules_count; i++) {
     char *module_name = read_string(fd);
     if (module_name == NULL) {
       PLOGE("reading module name");
 
-      info->modules->modules_count = i;
+      goto info_cleanup;
-
-      free_rezygisk_info(info);
-
-      info->modules = NULL;
-      info->modules->modules_count = 0;
-
-      close(fd);
-
-      return;
     }
 
     char module_path[PATH_MAX];
@@ -156,43 +145,49 @@ void rezygiskd_get_info(struct rezygisk_info *info) {
     if (!module_prop) {
       PLOGE("failed to open module prop file %s", module_path);
 
-      info->modules->modules_count = i;
+      goto info_cleanup;
-
-      free_rezygisk_info(info);
-
-      info->modules = NULL;
-      info->modules->modules_count = 0;
-
-      close(fd);
-
-      return;
     }
 
+    info->modules.modules[i] = NULL;
+
     char line[1024];
     while (fgets(line, sizeof(line), module_prop) != NULL) {
       if (strncmp(line, "name=", strlen("name=")) != 0) continue;
 
-      info->modules->modules[i] = strndup(line + 5, strlen(line) - 6);
+      info->modules.modules[i] = strndup(line + 5, strlen(line) - 6);
 
       break;
     }
 
+    if (info->modules.modules[i] == NULL) {
+      PLOGE("failed to read module name from %s", module_path);
+
       fclose(module_prop);
+
+      goto info_cleanup;
+    }
+
+    fclose(module_prop);
+
+    continue;
+
+    info_cleanup:
+      info->modules.modules_count = i;
+      free_rezygisk_info(info);
+
+      break;
   }
 
   close(fd);
 }
 
 void free_rezygisk_info(struct rezygisk_info *info) {
-  if (info->modules->modules) {
+  for (size_t i = 0; i < info->modules.modules_count; i++) {
-    for (size_t i = 0; i < info->modules->modules_count; i++) {
+    free(info->modules.modules[i]);
-      free(info->modules->modules[i]);
   }
 
-    free(info->modules->modules);
+  free(info->modules.modules);
-  }
+  info->modules.modules = NULL;
-
-  free(info->modules);
 }
 
 bool rezygiskd_read_modules(struct zygisk_modules *modules) {
@@ -237,13 +232,11 @@ bool rezygiskd_read_modules(struct zygisk_modules *modules) {
 }
 
 void free_modules(struct zygisk_modules *modules) {
-  if (modules->modules) {
   for (size_t i = 0; i < modules->modules_count; i++) {
     free(modules->modules[i]);
   }
 
   free(modules->modules);
-  }
 }
 
 int rezygiskd_connect_companion(size_t index) {

loader/src/common/socket_utils.c

@@ -25,7 +25,7 @@ int read_fd(int fd) {
     .msg_controllen = sizeof(cmsgbuf)
   };
 
-  ssize_t ret = recvmsg(fd, &msg, MSG_WAITALL);
+  ssize_t ret = TEMP_FAILURE_RETRY(recvmsg(fd, &msg, MSG_WAITALL));
   if (ret == -1) {
     PLOGE("recvmsg");
 
@@ -47,14 +47,14 @@ int read_fd(int fd) {
 
 ssize_t write_string(int fd, const char *str) {
   size_t str_len = strlen(str);
-  ssize_t write_bytes = write(fd, &str_len, sizeof(size_t));
+  ssize_t write_bytes = TEMP_FAILURE_RETRY(write(fd, &str_len, sizeof(size_t)));
   if (write_bytes != (ssize_t)sizeof(size_t)) {
     LOGE("Failed to write string length: Not all bytes were written (%zd != %zu).\n", write_bytes, sizeof(size_t));
 
     return -1;
   }
 
-  write_bytes = write(fd, str, str_len);
+  write_bytes = TEMP_FAILURE_RETRY(write(fd, str, str_len));
   if (write_bytes != (ssize_t)str_len) {
     LOGE("Failed to write string: Promised bytes doesn't exist (%zd != %zu).\n", write_bytes, str_len);
 
@@ -66,7 +66,7 @@ ssize_t write_string(int fd, const char *str) {
 
 char *read_string(int fd) {
   size_t str_len = 0;
-  ssize_t read_bytes = read(fd, &str_len, sizeof(size_t));
+  ssize_t read_bytes = TEMP_FAILURE_RETRY(read(fd, &str_len, sizeof(size_t)));
   if (read_bytes != (ssize_t)sizeof(size_t)) {
     LOGE("Failed to read string length: Not all bytes were read (%zd != %zu).\n", read_bytes, sizeof(size_t));
 
@@ -80,7 +80,7 @@ char *read_string(int fd) {
     return NULL;
   }
 
-  read_bytes = read(fd, buf, str_len);
+  read_bytes = TEMP_FAILURE_RETRY(read(fd, buf, str_len));
   if (read_bytes != (ssize_t)str_len) {
     LOGE("Failed to read string: Promised bytes doesn't exist (%zd != %zu).\n", read_bytes, str_len);
 
@@ -96,12 +96,12 @@ char *read_string(int fd) {
 
 #define write_func(type)                                      \
   ssize_t write_## type(int fd, type val) {                   \
-    return write(fd, &val, sizeof(type));   \
+    return TEMP_FAILURE_RETRY(write(fd, &val, sizeof(type))); \
   }
 
 #define read_func(type)                                     \
   ssize_t read_## type(int fd, type *val) {                 \
-    return read(fd, val, sizeof(type));     \
+    return TEMP_FAILURE_RETRY(read(fd, val, sizeof(type))); \
   }
 
 write_func(uint8_t)

loader/src/include/daemon.h

@@ -42,7 +42,7 @@ enum root_impl {
 };
 
 struct rezygisk_info {
-  struct zygisk_modules *modules;
+  struct zygisk_modules modules;
   enum root_impl root_impl;
   pid_t pid;
   bool running;

loader/src/ptracer/main.c

@@ -77,11 +77,11 @@ int main(int argc, char **argv) {
       }
     }
 
-    if (info.modules->modules_count != 0) {
+    if (info.modules.modules_count != 0) {
-      printf("Modules: %zu\n", info.modules->modules_count);
+      printf("Modules: %zu\n", info.modules.modules_count);
 
-      for (size_t i = 0; i < info.modules->modules_count; i++) {
+      for (size_t i = 0; i < info.modules.modules_count; i++) {
-        printf(" - %s\n", info.modules->modules[i]);
+        printf(" - %s\n", info.modules.modules[i]);
       }
     } else {
       printf("Modules: N/A\n");

loader/src/ptracer/monitor.c

@@ -66,6 +66,7 @@ struct rezygiskd_status status32 = {
 
 int monitor_epoll_fd;
 bool monitor_events_running = true;
+typedef void (*monitor_event_callback_t)();
 
 bool monitor_events_init() {
   monitor_epoll_fd = epoll_create(1);
@@ -78,14 +79,9 @@ bool monitor_events_init() {
   return true;
 }
 
-struct monitor_event_cbs {
+bool monitor_events_register_event(monitor_event_callback_t event_cb, int fd, uint32_t events) {
-  void (*callback)();
-  void (*stop_callback)();
-};
-
-bool monitor_events_register_event(struct monitor_event_cbs *event_cbs, int fd, uint32_t events) {
   struct epoll_event ev = {
-    .data.ptr = event_cbs,
+    .data.ptr = (void *)event_cb,
     .events = events
   };
 
@@ -116,15 +112,16 @@ void monitor_events_loop() {
   struct epoll_event events[2];
   while (monitor_events_running) {
     int nfds = epoll_wait(monitor_epoll_fd, events, 2, -1);
-    if (nfds == -1) {
+    if (nfds == -1 && errno != EINTR) {
-      if (errno != EINTR) PLOGE("epoll_wait");
+      PLOGE("epoll_wait");
 
-      continue;
+      monitor_events_running = false;
+
+      break;
     }
 
     for (int i = 0; i < nfds; i++) { 
-      struct monitor_event_cbs *event_cbs = (struct monitor_event_cbs *)events[i].data.ptr;
+      ((monitor_event_callback_t)events[i].data.ptr)();
-      event_cbs->callback();
 
       if (!monitor_events_running) break;
     }
@@ -132,11 +129,6 @@ void monitor_events_loop() {
 
   if (monitor_epoll_fd >= 0) close(monitor_epoll_fd);
   monitor_epoll_fd = -1;
-
-  for (int i = 0; i < (int)(sizeof(events) / sizeof(events[0])); i++) {
-    struct monitor_event_cbs *event_cbs = (struct monitor_event_cbs *)events[i].data.ptr;
-    event_cbs->stop_callback();
-  }
 }
 
 int monitor_sock_fd;
@@ -272,15 +264,13 @@ void rezygiskd_listener_callback() {
           status64.daemon_info = NULL;
         }
 
-        status64.daemon_info = (char *)malloc(msg.length);
+        status64.daemon_info = strdup(msg_data);
         if (!status64.daemon_info) {
           PLOGE("malloc daemon64 info");
 
           break;
         }
 
-        strcpy(status64.daemon_info, msg_data);
-
         update_status(NULL);
 
         break;
@@ -293,15 +283,13 @@ void rezygiskd_listener_callback() {
           status32.daemon_info = NULL;
         }
 
-        status32.daemon_info = (char *)malloc(msg.length);
+        status32.daemon_info = strdup(msg_data);
         if (!status32.daemon_info) {
           PLOGE("malloc daemon32 info");
 
           break;
         }
 
-        strcpy(status32.daemon_info, msg_data);
-
         update_status(NULL);
 
         break;
@@ -316,15 +304,13 @@ void rezygiskd_listener_callback() {
           status64.daemon_error_info = NULL;
         }
 
-        status64.daemon_error_info = (char *)malloc(msg.length);
+        status64.daemon_error_info = strdup(msg_data);
         if (!status64.daemon_error_info) {
           PLOGE("malloc daemon64 error info");
 
           break;
         }
 
-        strcpy(status64.daemon_error_info, msg_data);
-
         update_status(NULL);
 
         break;
@@ -339,15 +325,13 @@ void rezygiskd_listener_callback() {
           status32.daemon_error_info = NULL;
         }
 
-        status32.daemon_error_info = (char *)malloc(msg.length);
+        status32.daemon_error_info = strdup(msg_data);
         if (!status32.daemon_error_info) {
           PLOGE("malloc daemon32 error info");
 
           break;
         }
 
-        strcpy(status32.daemon_error_info, msg_data);
-
         update_status(NULL);
 
         break;
@@ -847,10 +831,6 @@ void init_monitor() {
 
   monitor_events_init();
 
-  struct monitor_event_cbs listener_cbs = {
-    .callback = rezygiskd_listener_callback,
-    .stop_callback = rezygiskd_listener_stop
-  };
   if (!rezygiskd_listener_init()) {
     LOGE("failed to create socket");
 
@@ -859,12 +839,8 @@ void init_monitor() {
     exit(1);
   }
 
-  monitor_events_register_event(&listener_cbs, monitor_sock_fd, EPOLLIN | EPOLLET);
+  monitor_events_register_event(rezygiskd_listener_callback, monitor_sock_fd, EPOLLIN | EPOLLET);
 
-  struct monitor_event_cbs sigchld_cbs = {
-    .callback = sigchld_listener_callback,
-    .stop_callback = sigchld_listener_stop
-  };
   if (sigchld_listener_init() == false) {
     LOGE("failed to create signalfd");
 
@@ -874,10 +850,15 @@ void init_monitor() {
     exit(1);
   }
 
-  monitor_events_register_event(&sigchld_cbs, sigchld_signal_fd, EPOLLIN | EPOLLET);
+  monitor_events_register_event(sigchld_listener_callback, sigchld_signal_fd, EPOLLIN | EPOLLET);
 
   monitor_events_loop();
 
+  /* INFO: Once it stops the loop, we cannot access the epool data, so we
+             either manually call the stops or save to a structure. */
+  rezygiskd_listener_stop();
+  sigchld_listener_stop();
+
   if (status64.daemon_info) free(status64.daemon_info);
   if (status64.daemon_error_info) free(status64.daemon_error_info);
   if (status32.daemon_info) free(status32.daemon_info);

loader/src/ptracer/utils.c

@@ -119,8 +119,8 @@ struct maps *parse_maps(const char *filename) {
       path_offset++;
     }
 
-    maps->maps = (struct map *)realloc(maps->maps, (i + 1) * sizeof(struct map));
+    struct map *tmp_maps = (struct map *)realloc(maps->maps, (i + 1) * sizeof(struct map));
-    if (!maps->maps) {
+    if (!tmp_maps) {
       LOGE("Failed to allocate memory for maps->maps");
 
       maps->size = i;
@@ -130,6 +130,7 @@ struct maps *parse_maps(const char *filename) {
 
       return NULL;
     }
+    maps->maps = tmp_maps;
 
     maps->maps[i].start = addr_start;
     maps->maps[i].end = addr_end;

zygiskd/src/companion.c

@@ -153,7 +153,14 @@ void companion_entry(int fd) {
     LOGI("New companion request.\n - Module name: %s\n - Client fd: %d\n", name, client_fd);
 
     ret = write_uint8_t(client_fd, 1);
-    ASSURE_SIZE_WRITE("ZygiskdCompanion", "client_fd", ret, sizeof(uint8_t));
+    if (ret != sizeof(uint8_t)) {
+      LOGE("Failed to sent client_fd in ZygiskdCompanion: Expected %zu, got %zd\n", sizeof(uint8_t), ret);
+
+      free(args);
+      close(client_fd);
+
+      break;
+    }
 
     pthread_t thread;
     if (pthread_create(&thread, NULL, entry_thread, (void *)args) == 0)

zygiskd/src/root_impl/apatch.c

@@ -73,6 +73,14 @@ struct packages_config {
   size_t size;
 };
 
+void _apatch_free_package_config(struct packages_config *restrict config) {
+  for (size_t i = 0; i < config->size; i++) {
+    free(config->configs[i].process);
+  }
+
+  free(config->configs);
+}
+
 /* WARNING: Dynamic memory based */
 bool _apatch_get_package_config(struct packages_config *restrict config) {
   config->configs = NULL;
@@ -96,14 +104,16 @@ bool _apatch_get_package_config(struct packages_config *restrict config) {
   }
 
   while (fgets(line, sizeof(line), fp) != NULL) { 
-    config->configs = realloc(config->configs, (config->size + 1) * sizeof(struct package_config));
+    struct package_config *tmp_configs = realloc(config->configs, (config->size + 1) * sizeof(struct package_config));
-    if (config->configs == NULL) {
+    if (tmp_configs == NULL) {
       LOGE("Failed to realloc APatch config struct: %s\n", strerror(errno));
 
+      _apatch_free_package_config(config);
       fclose(fp);
 
       return false;
     }
+    config->configs = tmp_configs;
 
     config->configs[config->size].process = strdup(strtok(line, ","));
 
@@ -128,21 +138,9 @@ bool _apatch_get_package_config(struct packages_config *restrict config) {
   return true;
 }
 
-void _apatch_free_package_config(struct packages_config *restrict config) {
-  for (size_t i = 0; i < config->size; i++) {
-    free(config->configs[i].process);
-  }
-
-  free(config->configs);
-}
-
 bool apatch_uid_granted_root(uid_t uid) {
   struct packages_config config;
-  if (!_apatch_get_package_config(&config)) {
+  if (!_apatch_get_package_config(&config)) return false;
-    _apatch_free_package_config(&config);
-
-    return false;
-  }
 
   for (size_t i = 0; i < config.size; i++) {
     if (config.configs[i].uid != uid) continue;
@@ -162,11 +160,7 @@ bool apatch_uid_granted_root(uid_t uid) {
 
 bool apatch_uid_should_umount(uid_t uid, const char *const process) {
   struct packages_config config;
-  if (!_apatch_get_package_config(&config)) {
+  if (!_apatch_get_package_config(&config)) return false;
-    _apatch_free_package_config(&config);
-
-    return false;
-  }
 
   for (size_t i = 0; i < config.size; i++) {
     if (config.configs[i].uid != uid) continue;

zygiskd/src/root_impl/common.c

@@ -74,8 +74,7 @@ void root_impls_setup(void) {
 }
 
 void get_impl(struct root_impl *uimpl) {
-  uimpl->impl = impl.impl;
+  *uimpl = impl;
-  uimpl->variant = impl.variant;
 }
 
 bool uid_granted_root(uid_t uid) {

zygiskd/src/utils.c

@@ -109,12 +109,9 @@ static void get_current_attr(char *restrict output, size_t size) {
     return;
   }
 
-  if (fread(output, 1, size, current) == 0) {
+  if (fread(output, 1, size, current) == 0)
     LOGE("fread: %s\n", strerror(errno));
 
-    return;
-  }
-
   fclose(current);
 }
 
@@ -543,15 +540,13 @@ bool parse_mountinfo(const char *restrict pid, struct mountinfos *restrict mount
             &optional_start, &optional_end, &type_start, &type_end,
             &source_start, &source_end, &fs_option_start, &fs_option_end);
 
-    mounts->mounts = (struct mountinfo *)realloc(mounts->mounts, (i + 1) * sizeof(struct mountinfo));
+    struct mountinfo *tmp_mounts = (struct mountinfo *)realloc(mounts->mounts, (i + 1) * sizeof(struct mountinfo));
-    if (!mounts->mounts) {
+    if (!tmp_mounts) {
       LOGE("Failed to allocate memory for mounts->mounts");
 
-      fclose(mountinfo);
+      goto cleanup_mount_allocs;
-      free_mounts(mounts);
-
-      return false;
     }
+    mounts->mounts = tmp_mounts;
 
     unsigned int shared = 0;
     unsigned int master = 0;
@@ -572,16 +567,64 @@ bool parse_mountinfo(const char *restrict pid, struct mountinfos *restrict mount
     mounts->mounts[i].parent = parent;
     mounts->mounts[i].device = (dev_t)(makedev(maj, min));
     mounts->mounts[i].root = strndup(line + root_start, (size_t)(root_end - root_start));
+    if (mounts->mounts[i].root == NULL) {
+      LOGE("Failed to allocate memory for root\n");
+
+      goto cleanup_mount_allocs;
+    }
     mounts->mounts[i].target = strndup(line + target_start, (size_t)(target_end - target_start));
+    if (mounts->mounts[i].target == NULL) {
+      LOGE("Failed to allocate memory for target\n");
+
+      goto cleanup_root;
+    }
     mounts->mounts[i].vfs_option = strndup(line + vfs_option_start, (size_t)(vfs_option_end - vfs_option_start));
+    if (mounts->mounts[i].vfs_option == NULL) {
+      LOGE("Failed to allocate memory for vfs_option\n");
+
+      goto cleanup_target;
+    }
     mounts->mounts[i].optional.shared = shared;
     mounts->mounts[i].optional.master = master;
     mounts->mounts[i].optional.propagate_from = propagate_from;
     mounts->mounts[i].type = strndup(line + type_start, (size_t)(type_end - type_start));
+    if (mounts->mounts[i].type == NULL) {
+      LOGE("Failed to allocate memory for type\n");
+
+      goto cleanup_vfs_option;
+    }
     mounts->mounts[i].source = strndup(line + source_start, (size_t)(source_end - source_start));
+    if (mounts->mounts[i].source == NULL) {
+      LOGE("Failed to allocate memory for source\n");
+
+      goto cleanup_type;
+    }
     mounts->mounts[i].fs_option = strndup(line + fs_option_start, (size_t)(fs_option_end - fs_option_start));
+    if (mounts->mounts[i].fs_option == NULL) {
+      LOGE("Failed to allocate memory for fs_option\n");
+
+      goto cleanup_source;
+    }
 
     i++;
+
+    continue;
+
+    cleanup_source:
+      free((void *)mounts->mounts[i].source);
+    cleanup_type:
+      free((void *)mounts->mounts[i].type);
+    cleanup_vfs_option:
+      free((void *)mounts->mounts[i].vfs_option);
+    cleanup_target:
+      free((void *)mounts->mounts[i].target);
+    cleanup_root:
+      free((void *)mounts->mounts[i].root);
+    cleanup_mount_allocs:
+      fclose(mountinfo);
+      free_mounts(mounts);
+
+      return false;
   }
 
   fclose(mountinfo);

zygiskd/src/zygiskd.c

@@ -123,7 +123,6 @@ static void load_modules(enum Architecture arch, struct Context *restrict contex
       continue;
     }
 
-
     context->modules = realloc(context->modules, (size_t)((context->len + 1) * sizeof(struct Module)));
     if (context->modules == NULL) {
       LOGE("Failed reallocating memory for modules.\n");
@@ -143,8 +142,10 @@ static void load_modules(enum Architecture arch, struct Context *restrict contex
 static void free_modules(struct Context *restrict context) {
   for (size_t i = 0; i < context->len; i++) {
     free(context->modules[i].name);
-    if (context->modules[i].companion != -1) close(context->modules[i].companion);
+    if (context->modules[i].companion >= 0) close(context->modules[i].companion);
   }
+
+  free(context->modules);
 }
 
 static int create_daemon_socket(void) {
@@ -288,7 +289,7 @@ void zygiskd_start(char *restrict argv[]) {
     unix_datagram_sendto(CONTROLLER_SOCKET, &msg, sizeof(struct MsgHead));
     unix_datagram_sendto(CONTROLLER_SOCKET, msg_data, (size_t)msg.length);
 
-    free(msg_data);
+    exit(EXIT_FAILURE);
   } else {
     enum Architecture arch = get_arch();
     load_modules(arch, &context);
@@ -301,12 +302,24 @@ void zygiskd_start(char *restrict argv[]) {
     } else {
       for (size_t i = 0; i < context.len; i++) {
         if (i != context.len - 1) {
-          module_list = realloc(module_list, module_list_len + strlen(context.modules[i].name) + strlen(", ") + 1);
+          char *tmp_module_list = realloc(module_list, module_list_len + strlen(context.modules[i].name) + strlen(", ") + 1);
-          if (module_list == NULL) {
+          if (tmp_module_list == NULL) {
             LOGE("Failed reallocating memory for module list.\n");
 
-            return;
+            char *kmsg_failure = "Failed reallocating memory for module list";
+            struct MsgHead msg = {
+              .cmd = DAEMON_SET_ERROR_INFO,
+              .length = (int)strlen(kmsg_failure) + 1
+            };
+            unix_datagram_sendto(CONTROLLER_SOCKET, &msg, sizeof(struct MsgHead));
+            unix_datagram_sendto(CONTROLLER_SOCKET, kmsg_failure, (size_t)msg.length);
+
+            free(module_list);
+            free_modules(&context);
+
+            exit(EXIT_FAILURE);
           }
+          module_list = tmp_module_list;
 
           strcpy(module_list + module_list_len, context.modules[i].name);
 
@@ -316,12 +329,24 @@ void zygiskd_start(char *restrict argv[]) {
 
           module_list_len += strlen(", ");
         } else {
-          module_list = realloc(module_list, module_list_len + strlen(context.modules[i].name) + 1);
+          char *tmp_module_list = realloc(module_list, module_list_len + strlen(context.modules[i].name) + 1);
-          if (module_list == NULL) {
+          if (tmp_module_list == NULL) {
             LOGE("Failed reallocating memory for module list.\n");
 
-            return;
+            char *kmsg_failure = "Failed reallocating memory for module list";
+            struct MsgHead msg = {
+              .cmd = DAEMON_SET_ERROR_INFO,
+              .length = (int)strlen(kmsg_failure) + 1
+            };
+            unix_datagram_sendto(CONTROLLER_SOCKET, &msg, sizeof(struct MsgHead));
+            unix_datagram_sendto(CONTROLLER_SOCKET, kmsg_failure, (size_t)msg.length);
+
+            free(module_list);
+            free_modules(&context);
+
+            exit(EXIT_FAILURE);
           }
+          module_list = tmp_module_list;
 
           strcpy(module_list + module_list_len, context.modules[i].name);
 
@@ -344,7 +369,16 @@ void zygiskd_start(char *restrict argv[]) {
     if (msg_data == NULL) {
       LOGE("Failed allocating memory for message data.\n");
 
-      return;
+      char *kmsg_failure = "Failed allocating memory for message data";
+      msg.cmd = DAEMON_SET_ERROR_INFO;
+      msg.length = (int)strlen(kmsg_failure) + 1;
+      unix_datagram_sendto(CONTROLLER_SOCKET, &msg, sizeof(struct MsgHead));
+      unix_datagram_sendto(CONTROLLER_SOCKET, kmsg_failure, (size_t)msg.length);
+
+      free(module_list);
+      free_modules(&context);
+
+      exit(EXIT_FAILURE);
     }
 
     snprintf(msg_data, msg_length, "Root: %s, Modules: %s", impl_name, module_list);
@@ -360,6 +394,8 @@ void zygiskd_start(char *restrict argv[]) {
   if (socket_fd == -1) {
     LOGE("Failed creating daemon socket\n");
 
+    free_modules(&context);
+
     return;
   }
 
@@ -372,7 +408,7 @@ void zygiskd_start(char *restrict argv[]) {
     if (client_fd == -1) {
       LOGE("accept: %s\n", strerror(errno));
 
-      return;
+      break;
     }
 
     uint8_t action8 = 0;
@@ -380,11 +416,11 @@ void zygiskd_start(char *restrict argv[]) {
     if (len == -1) {
       LOGE("read: %s\n", strerror(errno));
 
-      return;
+      break;
     } else if (len == 0) {
       LOGI("Client disconnected\n");
 
-      return;
+      break;
     }
 
     enum DaemonSocketAction action = (enum DaemonSocketAction)action8;
@@ -402,11 +438,11 @@ void zygiskd_start(char *restrict argv[]) {
       }
       case ZygoteRestart: {
         for (size_t i = 0; i < context.len; i++) {
-          if (context.modules[i].companion != -1) {
+          if (context.modules[i].companion <= -1) continue;
+
           close(context.modules[i].companion);
           context.modules[i].companion = -1;
         }
-        }
 
         break;
       }
@@ -566,9 +602,19 @@ void zygiskd_start(char *restrict argv[]) {
         ssize_t ret = read_size_t(client_fd, &index);
         ASSURE_SIZE_READ_BREAK("RequestCompanionSocket", "index", ret, sizeof(index));
 
-        struct Module *module = &context.modules[index];
+        if (index >= context.len) {
+          LOGE("Invalid module index: %zu\n", index);
 
-        if (module->companion != -1) {
+          ret = write_uint8_t(client_fd, 0);
+          ASSURE_SIZE_WRITE_BREAK("RequestCompanionSocket", "response", ret, sizeof(int));
+
+          close(client_fd);
+
+          break;
+        }
+
+        struct Module *module = &context.modules[index];
+        if (module->companion >= 0) {
           if (!check_unix_socket(module->companion, false)) {
             LOGE(" - Companion for module \"%s\" crashed\n", module->name);
 
@@ -577,10 +623,10 @@ void zygiskd_start(char *restrict argv[]) {
           }
         }
 
-        if (module->companion == -1) {
+        if (module->companion <= -1) {
           module->companion = spawn_companion(argv, module->name, module->lib_fd);
 
-          if (module->companion > 0) {
+          if (module->companion >= 0) {
             LOGI(" - Spawned companion for \"%s\": %d\n", module->name, module->companion);
           } else {
             if (module->companion == -2) {
@@ -597,7 +643,7 @@ void zygiskd_start(char *restrict argv[]) {
                  so just sending the file descriptor of the client is
                  safe.
         */
-        if (module->companion != -1) {
+        if (module->companion >= 0) {
           LOGI(" - Sending companion fd socket of module \"%s\"\n", module->name);
 
           if (write_fd(module->companion, client_fd) == -1) {
@@ -629,6 +675,17 @@ void zygiskd_start(char *restrict argv[]) {
         ssize_t ret = read_size_t(client_fd, &index);
         ASSURE_SIZE_READ_BREAK("GetModuleDir", "index", ret, sizeof(index));
 
+        if (index >= context.len) {
+          LOGE("Invalid module index: %zu\n", index);
+
+          ret = write_uint8_t(client_fd, 0);
+          ASSURE_SIZE_WRITE_BREAK("GetModuleDir", "response", ret, sizeof(int));
+
+          close(client_fd);
+
+          break;
+        }
+
         char module_dir[PATH_MAX];
         snprintf(module_dir, PATH_MAX, "%s/%s", PATH_MODULES_DIR, context.modules[index].name);