Fixed Memory Leak

1 files changed, 19 insertions, 2 deletions

diff --git a/lib/dir.c b/lib/dir.c
index 361775d..f0776ca 100644
--- a/lib/dir.c
+++ b/lib/dir.c
@@ -22,7 +22,7 @@ char **list_directory(char *dirname) {
   struct dirent *dir;
   d = opendir(dirname);
   char **file_names = (char **)g_malloc(sizeof(char *));
-  if (!file_names) {
+  if (file_names == NULL) {
     return NULL;
   }
   file_names[0] = NULL;
@@ -32,10 +32,24 @@ char **list_directory(char *dirname) {
       if (dir->d_type == DT_REG) {
         // When a regular file is reached
         /// Create space for it in the list
-        file_names = realloc(file_names, (file_count + 1 + 1) * sizeof(char *));
+        char **temp =
+            g_realloc(file_names, (file_count + 1 + 1) * sizeof(char *));
+        if (temp == NULL) {
+          for (size_t file_idx = 0; file_idx < file_count; file_idx++) {
+	    g_free(file_names[file_idx]);
+          }
+	  return NULL;
+        }
+        file_names = temp;
         /// Create space for the name
         file_names[file_count] =
             g_calloc(strlen(dir->d_name) + 1, sizeof(char));
+        if (file_names[file_count] == NULL) {
+          for (size_t file_idx = 0; file_idx < file_count; file_idx++) {
+	    g_free(file_names[file_idx]);
+          }
+	  return NULL;
+        }
         /// Copy the name
         strcpy(file_names[file_count], dir->d_name);
         /// Mark the end of the file list with a NULL entry
@@ -53,6 +67,9 @@ char *full_path(char *dir, char *file) {
   size_t dir_len = strlen(dir);
   size_t file_len = strlen(file);
   fpath = (char *)g_calloc(dir_len + file_len + 2, sizeof(char));
+  if (fpath == NULL) {
+    return NULL;
+  }
   strcpy(fpath, dir);
   strcpy(fpath + dir_len + 1, file);
   fpath[dir_len] = '/';