yarg: oom handling
diff --git a/include/yarg.h b/include/yarg.h
index 81c0c56..ae00491 100644
--- a/include/yarg.h
+++ b/include/yarg.h
@@ -47,20 +47,21 @@ typedef struct {
char * error;
} yarg_result;
-static void * yarg_alloc(size_t size) {
- void * ptr = calloc(size, 1);
- if (!ptr) { perror("calloc"); exit(1); }
- return ptr;
-}
-
+static const char yarg_oom[] = "Out of memory";
static int yarg_asprintf(char ** strp, const char * fmt, ...) {
va_list ap;
va_start(ap, fmt);
int len = vsnprintf(NULL, 0, fmt, ap);
va_end(ap);
- if (len < 0) return -1;
+ if (len < 0) {
+ memcpy(*strp, yarg_oom, sizeof(yarg_oom));
+ return sizeof(yarg_oom);
+ }
*strp = (char *) malloc(len + 1);
- if (!*strp) return -1;
+ if (!*strp) {
+ memcpy(*strp, yarg_oom, sizeof(yarg_oom));
+ return sizeof(yarg_oom);
+ }
va_start(ap, fmt);
len = vsnprintf(*strp, len + 1, fmt, ap);
va_end(ap);
@@ -68,13 +69,14 @@ static int yarg_asprintf(char ** strp, const char * fmt, ...) {
}
static char * yarg_strdup(const char * str) {
- char * new_str = (char *) yarg_alloc(strlen(str) + 1);
+ char * new_str = (char *) calloc(strlen(str) + 1, 1);
+ if (!new_str) return NULL;
strcpy(new_str, str);
return new_str;
}
-static void yarg_parse_unix(int argc, char * argv[], yarg_options opt[],
- yarg_result * res, bool dash_dash) {
+static int yarg_parse_unix(int argc, char * argv[], yarg_options opt[],
+ yarg_result * res, bool dash_dash) {
int no_args = 0, no_pos_args = 0;
for (int i = 1; i < argc; i++) {
if (argv[i][0] == '-') {
@@ -88,7 +90,7 @@ static void yarg_parse_unix(int argc, char * argv[], yarg_options opt[],
{ o = &opt[j]; break; }
if (!o) {
yarg_asprintf(&res->error, "--%.*s -- unknown option\n", len, long_opt);
- return;
+ return 0;
}
if (o->type == required_argument) {
if (long_opt[len] == '=') {
@@ -97,7 +99,7 @@ static void yarg_parse_unix(int argc, char * argv[], yarg_options opt[],
i++;
} else {
yarg_asprintf(&res->error, "--%s -- missing argument\n", o->long_opt);
- return;
+ return 0;
}
} else if (o->type == optional_argument) {
if (long_opt[len] == '=') {
@@ -114,7 +116,7 @@ static void yarg_parse_unix(int argc, char * argv[], yarg_options opt[],
{ o = &opt[k]; break; }
if (!o) {
yarg_asprintf(&res->error, "-%c -- unknown option\n", c);
- return;
+ return 0;
}
if (o->type == required_argument) {
if (argv[i][j + 1]) {
@@ -123,7 +125,7 @@ static void yarg_parse_unix(int argc, char * argv[], yarg_options opt[],
i++;
} else {
yarg_asprintf(&res->error, "-%c -- missing argument\n", c);
- return;
+ return 0;
}
no_args++;
break;
@@ -144,15 +146,22 @@ static void yarg_parse_unix(int argc, char * argv[], yarg_options opt[],
} else no_pos_args++;
}
- res->args = (yarg_option *) yarg_alloc((no_args + 1) * sizeof(yarg_option));
- res->pos_args = (char **) yarg_alloc((no_pos_args + 1) * sizeof(char *));
+ res->args = (yarg_option *) calloc((no_args + 1) * sizeof(yarg_option), 1);
+ res->pos_args = (char **) calloc((no_pos_args + 1) * sizeof(char *), 1);
+ if(!res->args || !res->pos_args) {
+ yarg_asprintf(&res->error, yarg_oom);
+ return 0;
+ }
for (int i = 1; i < argc; i++) {
if (argv[i][0] == '-') {
if (argv[i][1] == '-') {
if (dash_dash && argv[i][2] == '\0') {
for (int j = i + 1; j < argc; j++)
- res->pos_args[res->pos_argc++] = yarg_strdup(argv[j]);
+ if(!(res->pos_args[res->pos_argc++] = yarg_strdup(argv[j]))) {
+ yarg_asprintf(&res->error, yarg_oom);
+ return 0;
+ }
break;
}
char * long_opt = argv[i] + 2; yarg_options * o = NULL;
@@ -164,9 +173,15 @@ static void yarg_parse_unix(int argc, char * argv[], yarg_options opt[],
res->args[res->argc].long_opt = o->long_opt;
if (o->type == required_argument || o->type == optional_argument) {
if (long_opt[len] == '=') {
- res->args[res->argc].arg = yarg_strdup(long_opt + len + 1);
+ if(!(res->args[res->argc].arg = yarg_strdup(long_opt + len + 1))) {
+ yarg_asprintf(&res->error, yarg_oom);
+ return 0;
+ }
} else if (argv[i + 1] && argv[i + 1][0] != '-') {
- res->args[res->argc].arg = yarg_strdup(argv[++i]);
+ if(!(res->args[res->argc].arg = yarg_strdup(argv[++i]))) {
+ yarg_asprintf(&res->error, yarg_oom);
+ return 0;
+ }
}
}
res->argc++;
@@ -178,28 +193,39 @@ static void yarg_parse_unix(int argc, char * argv[], yarg_options opt[],
{ o = &opt[k]; break; }
if (!o) {
yarg_asprintf(&res->error, "-%c -- unknown option\n", c);
- return;
+ return 0;
}
res->args[res->argc].opt = c;
res->args[res->argc].long_opt = o->long_opt;
if (o->type == required_argument || o->type == optional_argument) {
if (argv[i][j + 1]) {
- res->args[res->argc++].arg = yarg_strdup(argv[i] + j + 1);
+ if(!(res->args[res->argc++].arg = yarg_strdup(argv[i] + j + 1))) {
+ yarg_asprintf(&res->error, yarg_oom);
+ return 0;
+ }
break;
} else if (argv[i + 1] && argv[i + 1][0] != '-') {
- res->args[res->argc++].arg = yarg_strdup(argv[++i]);
+ if(!(res->args[res->argc++].arg = yarg_strdup(argv[++i]))) {
+ yarg_asprintf(&res->error, yarg_oom);
+ return 0;
+ }
break;
}
}
res->argc++;
}
}
- } else res->pos_args[res->pos_argc++] = yarg_strdup(argv[i]);
+ } else if(!(res->pos_args[res->pos_argc++] = yarg_strdup(argv[i]))) {
+ yarg_asprintf(&res->error, yarg_oom);
+ return 0;
+ }
}
+
+ return 1;
}
-static void yarg_parse_unix_short(int argc, char * argv[], yarg_options opt[],
- yarg_result * res, bool dash_dash, char opt_char) {
+static int yarg_parse_unix_short(int argc, char * argv[], yarg_options opt[],
+ yarg_result * res, bool dash_dash, char opt_char) {
int no_args = 0, no_pos_args = 0;
for (int i = 1; i < argc; i++) {
if (argv[i][0] == opt_char) {
@@ -214,7 +240,7 @@ static void yarg_parse_unix_short(int argc, char * argv[], yarg_options opt[],
{ o = &opt[j]; break; }
if (!o) {
yarg_asprintf(&res->error, "%c%.*s -- unknown option\n", opt_char, len, long_opt);
- return;
+ return 0;
}
if (o->type == required_argument) {
if (long_opt[len] == '=') {
@@ -223,7 +249,7 @@ static void yarg_parse_unix_short(int argc, char * argv[], yarg_options opt[],
i++;
} else {
yarg_asprintf(&res->error, "%c%s -- missing argument\n", opt_char, o->long_opt);
- return;
+ return 0;
}
} else if (o->type == optional_argument) {
if (long_opt[len] == '=') {
@@ -236,14 +262,21 @@ static void yarg_parse_unix_short(int argc, char * argv[], yarg_options opt[],
} else no_pos_args++;
}
- res->args = (yarg_option *) yarg_alloc((no_args + 1) * sizeof(yarg_option));
- res->pos_args = (char **) yarg_alloc((no_pos_args + 1) * sizeof(char *));
+ res->args = (yarg_option *) calloc((no_args + 1) * sizeof(yarg_option), 1);
+ res->pos_args = (char **) calloc((no_pos_args + 1) * sizeof(char *), 1);
+ if (!res->args || !res->pos_args) {
+ yarg_asprintf(&res->error, yarg_oom);
+ return 0;
+ }
for (int i = 1; i < argc; i++) {
if (argv[i][0] == opt_char) {
if (dash_dash && argv[i][1] == '\0') {
for (int j = i + 1; j < argc; j++)
- res->pos_args[res->pos_argc++] = yarg_strdup(argv[j]);
+ if(!(res->pos_args[res->pos_argc++] = yarg_strdup(argv[j]))) {
+ yarg_asprintf(&res->error, yarg_oom);
+ return 0;
+ }
break;
}
char * long_opt = argv[i] + 1; yarg_options * o = NULL;
@@ -255,31 +288,50 @@ static void yarg_parse_unix_short(int argc, char * argv[], yarg_options opt[],
res->args[res->argc].long_opt = o->long_opt;
if (o->type == required_argument || o->type == optional_argument) {
if (long_opt[len] == '=') {
- res->args[res->argc].arg = yarg_strdup(long_opt + len + 1);
+ if(!(res->args[res->argc].arg = yarg_strdup(long_opt + len + 1))) {
+ yarg_asprintf(&res->error, yarg_oom);
+ return 0;
+ }
} else if (argv[i + 1] && argv[i + 1][0] != opt_char) {
- res->args[res->argc].arg = yarg_strdup(argv[++i]);
+ if(!(res->args[res->argc].arg = yarg_strdup(argv[++i]))) {
+ yarg_asprintf(&res->error, yarg_oom);
+ return 0;
+ }
}
}
res->argc++;
- } else res->pos_args[res->pos_argc++] = yarg_strdup(argv[i]);
+ } else if(!(res->pos_args[res->pos_argc++] = yarg_strdup(argv[i]))) {
+ yarg_asprintf(&res->error, yarg_oom);
+ return 0;
+ }
}
+
+ return 1;
}
void yarg_destroy(yarg_result * r) {
- for (int i = 0; i < r->argc; i++) {
- free(r->args[i].arg);
- }
- free(r->args);
- for (int i = 0; i < r->pos_argc; i++) {
- free(r->pos_args[i]);
+ if(r) {
+ if(r->args) {
+ for (int i = 0; i < r->argc; i++) {
+ free(r->args[i].arg);
+ }
+ }
+ free(r->args);
+ if(r->pos_args) {
+ for (int i = 0; i < r->pos_argc; i++) {
+ free(r->pos_args[i]);
+ }
+ }
+ free(r->pos_args);
+ if (r->error != yarg_oom)
+ free(r->error);
}
- free(r->pos_args);
- free(r->error);
free(r);
}
yarg_result * yarg_parse(int argc, char * argv[], yarg_options opt[], yarg_settings settings) {
- yarg_result * res = (yarg_result *) yarg_alloc(sizeof(yarg_result));
+ yarg_result * res = (yarg_result *) calloc(sizeof(yarg_result), 1);
+ if (!res) return NULL;
switch (settings.style) {
case YARG_STYLE_WINDOWS:
yarg_parse_unix_short(argc, argv, opt, res, false, '/');
diff --git a/src/main.c b/src/main.c
index 17116f3..643e188 100644
--- a/src/main.c
+++ b/src/main.c
@@ -583,6 +583,10 @@ int main(int argc, char * argv[]) {
.style = YARG_STYLE_UNIX,
};
yarg_result * res = yarg_parse(argc, argv, opt, settings);
+ if (!res) {
+ fprintf(stderr, "bzip3: out of memory.\n");
+ return 1;
+ }
if (res->error) {
fputs(res->error, stderr);
fputs("Try 'bzip3 --help' for more information.\n", stderr);
