btrfs-progs: handle memory allocation failures in traverse_directory

Signed-off-by: David Sterba <dsterba@suse.com>

1 files changed, 6 insertions, 0 deletions

diff --git a/mkfs.c b/mkfs.c
index 7316cb1..c09cbd3 100644
--- a/mkfs.c
+++ b/mkfs.c
@@ -786,6 +786,8 @@ static int traverse_directory(struct btrfs_trans_handle *trans,
 
 	/* Add list for source directory */
 	dir_entry = malloc(sizeof(struct directory_name_entry));
+	if (!dir_entry)
+		return -ENOMEM;
 	dir_entry->dir_name = dir_name;
 	dir_entry->path = realpath(dir_name, real_path);
 	if (!dir_entry->path) {
@@ -887,6 +889,10 @@ static int traverse_directory(struct btrfs_trans_handle *trans,
 
 			if (S_ISDIR(st.st_mode)) {
 				dir_entry = malloc(sizeof(struct directory_name_entry));
+				if (!dir_entry) {
+					ret = -ENOMEM;
+					goto fail;
+				}
 				dir_entry->dir_name = cur_file->d_name;
 				dir_entry->path = make_path(parent_dir_entry->path,
 							    cur_file->d_name);