better checking of malloc

Author: jwdeitch

src/diff.c

@@ -2,10 +2,15 @@
 #include "fs.h"
 #include <stdlib.h>
 #include <string.h>
+#include <stdio.h>
 
 char **missing_from_db(char **dbList, struct fs_discovered_migrations *fsList) {
 	int fsi = 0;
 	char **list = malloc(1000 * sizeof(char *));
+	if (list == NULL) {
+		printf("malloc failed to allocate\n");
+		exit(1);
+	}
 	int filesLocated = 0;
 
 	while (strcmp(fsList[fsi].name, "\0")) {
@@ -20,12 +25,20 @@ char **missing_from_db(char **dbList, struct fs_discovered_migrations *fsList) {
 		}
 		if (!located) {
 			list[filesLocated] = (char *) malloc(PATH_MAX + 1);
+			if (list[filesLocated] == NULL) {
+				printf("malloc failed to allocate\n");
+				exit(1);
+			}
 			strcpy(list[filesLocated], fsList[fsi].name);
 			filesLocated++;
 		}
 		fsi++;
 	}
 	list[filesLocated] = (char *)malloc(PATH_MAX + 1);
+	if (list[filesLocated] == NULL) {
+		printf("malloc failed to allocate\n");
+		exit(1);
+	}
 	strcpy(list[filesLocated],"\0");
 	return list;
 

src/fs.c

@@ -58,7 +58,15 @@ void listdir(const char *name, int level) {
 
 struct fs_discovered_migrations *getMigrationsFromFs(const char *dir) {
 	file_names_arr = calloc(1000, sizeof(*file_names_arr));
+	if (file_names_arr == NULL) {
+		printf("failed to allocate memory\n");
+		exit(1);
+	}
 	itr = calloc(1000, sizeof(int));
+	if (itr == NULL) {
+		printf("failed to allocate memory\n");
+		exit(1);
+	}
 	*itr = 0;
 	listdir(".", 15);
 	strcpy(file_names_arr[*itr].name, "\0");

src/main.c

@@ -8,7 +8,7 @@
 #include "diff.h"
 
 void printHelp();
-const char* version = "1.0.0";
+const char *version = "1.0.0";
 
 int main(int argc, char *argv[]) {
 	extern char *optarg;
@@ -18,10 +18,14 @@ int main(int argc, char *argv[]) {
 	 * s-status, H-Host, u-up, d-down, v-version, p-soft, g-provision, h-help
 	 */
 	int c = 0, u = 0, d = 0, s = 0, g = 0, p = 0, H = 0;
-	char *connStr =  (char*)malloc(PATH_MAX*sizeof(connStr));
+	char *connStr = (char *) malloc(PATH_MAX * sizeof(connStr));
+	if (connStr == NULL) {
+		printf("malloc failed to allocate\n");
+		exit(1);
+	}
 
-	if( connStr == NULL ) {
-		printf("malloc failed to dimension connStr");
+	if (connStr == NULL) {
+		printf("malloc failed to dimension connStr\n");
 		exit(1);
 	}
 
@@ -123,10 +127,14 @@ int main(int argc, char *argv[]) {
 	}
 
 	if (u) {
+		if (argv[optind] == NULL) {
+			fprintf(stderr, "No directory provided\n");
+			exit(1);
+		}
 		char *file = argv[optind];
 		char path[PATH_MAX];
 		realpath(file, path);
-		DIR* dir = opendir(path);
+		DIR *dir = opendir(path);
 		if (!dir) {
 			fprintf(stderr, "No valid directory provided %s\n", path);
 			exit(1);

src/pg.c

@@ -95,15 +95,26 @@ char **getMigrationsFromDb(PGconn *connection) {
 	char **list = malloc(1000 * sizeof(char *));
 	if (list == NULL) {
 		PQfinish(connection);
+		printf("malloc failed\n");
 		exit(1);
 	}
 
 	int i = 0;
 	for (; i < rows; i++) {
 		list[i] = (char *) malloc(PATH_MAX + 1);
+		if (list[i] == NULL) {
+			PQfinish(connection);
+			printf("malloc failed\n");
+			exit(1);
+		}
 		strcpy(list[i], strdup(PQgetvalue(res, i, 0)));
 	}
 	list[i + 1] = (char *) malloc(PATH_MAX + 1);
+	if (list[i + 1] == NULL) {
+		PQfinish(connection);
+		printf("malloc failed\n");
+		exit(1);
+	}
 	strcpy(list[i + 1], "\0");
 
 	PQclear(res);
@@ -143,6 +154,10 @@ void runMigrations(PGconn *connection, char **migrationsToBeRan, int should_simu
 			}
 
 			char *fileContents = malloc(fsize + 1);
+			if (fileContents == NULL) {
+				printf("malloc failed to allocate\n");
+				exit(1);
+			}
 			fread(fileContents, fsize, 1, f);
 			fclose(f);
 
@@ -236,6 +251,10 @@ void rollbackMigrations(PGconn *connection, int should_simulate) {
 
 		if (should_simulate == 0) {
 			char *fileContents = malloc(fsize + 1);
+			if (fileContents == NULL) {
+				printf("malloc failed to allocate\n");
+				exit(1);
+			}
 			fread(fileContents, fsize, 1, f);
 			fclose(f);