From a4cde11f6b1e9bab24567a9a4fbefd707dbaa977 Mon Sep 17 00:00:00 2001
From: ubq323 <ubq323@ubq323.website>
Date: Thu, 28 Nov 2024 12:06:37 +0000
Subject: check malloc results

---
 run.c | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

(limited to 'run.c')

diff --git a/run.c b/run.c
index c1b2ced..ac25d23 100644
--- a/run.c
+++ b/run.c
@@ -8,6 +8,7 @@
 #include <lauxlib.h>
 
 #define CKE(val, name) do { if (val == -1) { perror(name); exit(1); } } while (0)
+static inline void *CKN(void *p) { if (p == NULL) { fprintf(stderr,"out of memory!"); exit(1); } return p; }
 
 #ifdef DEBUG
 #define DP(...) fprintf(stderr, __VA_ARGS__);
@@ -61,8 +62,8 @@ void empipe(int pipefd[2]) {
 
 char *run(int count, char **progs[], size_t inputlen, char *inputstr, size_t *outlen) {
 	if (inputstr) count++;
-	int (*pipefds)[2] = malloc(count * sizeof(int[2]));
-	pid_t *pids = malloc(count * sizeof(pid_t));
+	int (*pipefds)[2] = CKN(malloc(count * sizeof(int[2])));
+	pid_t *pids = CKN(malloc(count * sizeof(pid_t)));
 	for (int i = 0; i < count; i++) empipe(pipefds[i]);
 	for (int i = 0; i < count; i++) {
 		int input  = (i == 0) ? -1 : pipefds[i-1][0];
@@ -85,7 +86,7 @@ char *run(int count, char **progs[], size_t inputlen, char *inputstr, size_t *ou
 	// if the program produces too much data, you will be killed
 	size_t readed = 0;
 	size_t cap = 1024*1024;
-	char *buf = malloc(cap*sizeof(char));
+	char *buf = CKN(malloc(cap*sizeof(char)));
 	size_t remaining = cap;
 	size_t amt = 0;
 	do {
@@ -93,7 +94,7 @@ char *run(int count, char **progs[], size_t inputlen, char *inputstr, size_t *ou
 		readed += amt;
 		remaining -= amt;
 		if (remaining == 0) {
-			buf = realloc(buf, cap + cap);
+			buf = CKN(realloc(buf, cap + cap));
 			remaining += cap;
 			cap += cap;
 		}
@@ -123,18 +124,18 @@ int do_the_thing(lua_State *L) {
 	if (lua_type(L, 1) != LUA_TTABLE) FAIL("run needs table argumnt");
 	nprogs = lua_rawlen(L, 1);
 	if (nprogs == 0) FAIL("need at least one program to run");
-	argvs = calloc(nprogs, sizeof(char**));
+	argvs = CKN(calloc(nprogs, sizeof(char**)));
 	for (int i = 0; i < nprogs; i++) {
 		int ty = lua_rawgeti(L, 1, i + 1);
 		if (ty != LUA_TTABLE) FAIL("program arglist must be table");
 		size_t argc = lua_rawlen(L, 2);
 		if (argc == 0) FAIL("need at least one argument to program");
-		argvs[i] = calloc(argc + 1, sizeof(char*));
+		argvs[i] = CKN(calloc(argc + 1, sizeof(char*)));
 		for (int j = 0; j < argc; j++) {
 			int ty = lua_rawgeti(L, 2, j + 1);
 			if (ty != LUA_TSTRING) FAIL("program argument must be string");
 			size_t slen = 0; char *str = lua_tolstring(L, 3, &slen);
-			argvs[i][j] = calloc(slen + 1, sizeof(char));
+			argvs[i][j] = CKN(calloc(slen + 1, sizeof(char)));
 			memcpy(argvs[i][j], str, slen);
 			argvs[i][j][slen] = '\0'; // just to make sure
 			lua_pop(L, 1);
-- 
cgit v1.2.3