From 04f6d975ca70fd326c8a6af64467db7282e632e5 Mon Sep 17 00:00:00 2001
From: Nick Chan <towinchenmi@gmail.com>
Date: Thu, 23 Nov 2023 12:35:01 +0800
Subject: [PATCH] check malloc returns

---
 Makefile              |  2 +-
 src/exec_checkra1n.c  | 12 ++++++++++++
 src/lockdown_helper.c | 12 ++++++++++++
 src/optparse.c        |  6 +++++-
 4 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/Makefile b/Makefile
index 0667217..27a2f26 100644
--- a/Makefile
+++ b/Makefile
@@ -3,7 +3,7 @@ DEP = $(SRC)/dep_root
 STRIP = strip
 CC ?= cc
 CFLAGS += -isystem $(DEP)/include -I$(SRC)/include -I$(SRC) -D_XOPEN_SOURCE=500
-CFLAGS += -Wall -Wextra -Wno-unused-parameter -DPALERAIN_VERSION=\"2.0.0\"
+CFLAGS += -Wall -Wextra -Wno-unused-parameter -DPALERAIN_VERSION=\"2.0.0\" -DHAVE_LIBIMOBILEDEVICE
 CFLAGS += -Wno-unused-variable -I$(SRC)/src -std=c99 -pedantic-errors -D_C99_SOURCE -D_POSIX_C_SOURCE=200112L
 LIBS += $(DEP)/lib/libimobiledevice-1.0.a $(DEP)/lib/libirecovery-1.0.a $(DEP)/lib/libusbmuxd-2.0.a
 LIBS += $(DEP)/lib/libimobiledevice-glue-1.0.a $(DEP)/lib/libplist-2.0.a -pthread -lm
diff --git a/src/exec_checkra1n.c b/src/exec_checkra1n.c
index 850b251..9cfe6d3 100644
--- a/src/exec_checkra1n.c
+++ b/src/exec_checkra1n.c
@@ -43,6 +43,10 @@ int exec_checkra1n(void) {
 #ifndef NO_CUSTOM_PONGO
 	if (pongo_path == NULL) {
 		pongo_path = malloc(strlen(tmpdir) + 20);
+		if (pongo_path == NULL) {
+			LOG(LOG_FATAL, "memory allocation failed\n");
+			return -1;
+		}
 		snprintf(pongo_path, strlen(tmpdir) + 20, "%s/Pongo.bin.XXXXXX", tmpdir);
 		fd_pongo = mkstemp(pongo_path);
 		if (fd_pongo == -1) {
@@ -62,6 +66,10 @@ int exec_checkra1n(void) {
 	} else external_pongo = true;
 #endif
 	checkra1n_path = malloc(strlen(tmpdir) + 20);
+	if (checkra1n_path == NULL) {
+		LOG(LOG_FATAL, "memory allocation failed\n");
+		return -1;
+	}
 	snprintf(checkra1n_path, strlen(tmpdir) + 20, "%s/checkra1n.XXXXXX", tmpdir);
 	fd_checkra1n = mkstemp(checkra1n_path);
 	if (fd_checkra1n == -1) {
@@ -99,6 +107,10 @@ int exec_checkra1n(void) {
 		if (darwinMajor < 20) {
 #endif
 			libcheckra1nhelper_dylib_path = malloc(strlen(tmpdir) + 40);
+			if (libcheckra1nhelper_dylib_path == NULL) {
+				LOG(LOG_FATAL, "memory allocation failed\n");
+				return -1;
+			}
 			snprintf(libcheckra1nhelper_dylib_path, strlen(tmpdir) + 40, "%s/libcheckra1nhelper.dylib.XXXXXX", tmpdir);
 			int helper_fd = mkstemp(libcheckra1nhelper_dylib_path);
 			if (helper_fd == -1) {
diff --git a/src/lockdown_helper.c b/src/lockdown_helper.c
index b7b941c..f332149 100644
--- a/src/lockdown_helper.c
+++ b/src/lockdown_helper.c
@@ -80,6 +80,10 @@ static char* _send_request_and_get_reply(int fd, const char* request, uint32_t l
 
 	uint32_t resp_size = ntohl(besize);
 	char* resp = (char*)malloc(resp_size+1);
+	if (resp == NULL) {
+		log_error("memory allocation failed");
+		return -1;
+	}
 
 	recv_bytes = 0;
 	usbmuxd_recv(fd, resp, resp_size, &recv_bytes);
@@ -196,6 +200,11 @@ int lockdown_get_string_value(int fd, const char* domain, const char* key, char*
 	}
 	size_t str_len = p - str_start;
 	*value = (char*)malloc(str_len + 1);
+	if (*value == NULL) {
+		log_error("memory allocation failed");
+		free(resp);
+		return -1;
+	}
 	snprintf(*value, str_len, "%s", str_start);
 	(*value)[str_len] = '\0';
 	free(resp);
@@ -242,4 +251,7 @@ int lockdown_enter_recovery(int fd) {
 	return result;
 }
 
+#else
+extern char** environ; /* ISO C requires a translation unit to contain at least one declaration */
+
 #endif
diff --git a/src/optparse.c b/src/optparse.c
index 272009b..77da1e8 100644
--- a/src/optparse.c
+++ b/src/optparse.c
@@ -200,7 +200,11 @@ int optparse(int argc, char* argv[]) {
 				return -1;
 			}
 			pongo_path = malloc(strlen(optarg) + 1);
-			strcpy(pongo_path, optarg);
+			if (pongo_path == NULL) {
+				LOG(LOG_FATAL, "memory allocation failed");
+				return -1;
+			}
+			snprintf(pongo_path, strlen(optarg) + 1, "%s", optarg);
 			break;
 		case 'o':
 			if (override_file(&override_overlay, overlay_to_upload, &binpack_dmg_len, optarg))