:: commit 61c764f6f8242b22eba00ac4e8379751fc23cfd6

Kamila Szewczyk <kspalaiologos@gmail.com> — 2024-12-17 11:08

parents: c5eb4c2234

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);
tab: 248 wrap: offon