Home | History | Annotate | Download | only in util
      1 /*
      2  *
      3  * Copyright 2015 gRPC authors.
      4  *
      5  * Licensed under the Apache License, Version 2.0 (the "License");
      6  * you may not use this file except in compliance with the License.
      7  * You may obtain a copy of the License at
      8  *
      9  *     http://www.apache.org/licenses/LICENSE-2.0
     10  *
     11  * Unless required by applicable law or agreed to in writing, software
     12  * distributed under the License is distributed on an "AS IS" BASIS,
     13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     14  * See the License for the specific language governing permissions and
     15  * limitations under the License.
     16  *
     17  */
     18 
     19 #include "test/core/util/cmdline.h"
     20 
     21 #include <limits.h>
     22 #include <stdio.h>
     23 #include <string.h>
     24 
     25 #include <grpc/support/alloc.h>
     26 #include <grpc/support/log.h>
     27 #include <grpc/support/string_util.h>
     28 #include "src/core/lib/gpr/string.h"
     29 
     30 typedef enum { ARGTYPE_INT, ARGTYPE_BOOL, ARGTYPE_STRING } argtype;
     31 
     32 typedef struct arg {
     33   const char* name;
     34   const char* help;
     35   argtype type;
     36   void* value;
     37   struct arg* next;
     38 } arg;
     39 
     40 struct gpr_cmdline {
     41   const char* description;
     42   arg* args;
     43   const char* argv0;
     44 
     45   const char* extra_arg_name;
     46   const char* extra_arg_help;
     47   void (*extra_arg)(void* user_data, const char* arg);
     48   void* extra_arg_user_data;
     49 
     50   int (*state)(gpr_cmdline* cl, char* arg);
     51   arg* cur_arg;
     52 
     53   int survive_failure;
     54 };
     55 
     56 static int normal_state(gpr_cmdline* cl, char* arg);
     57 
     58 gpr_cmdline* gpr_cmdline_create(const char* description) {
     59   gpr_cmdline* cl = static_cast<gpr_cmdline*>(gpr_zalloc(sizeof(gpr_cmdline)));
     60 
     61   cl->description = description;
     62   cl->state = normal_state;
     63 
     64   return cl;
     65 }
     66 
     67 void gpr_cmdline_set_survive_failure(gpr_cmdline* cl) {
     68   cl->survive_failure = 1;
     69 }
     70 
     71 void gpr_cmdline_destroy(gpr_cmdline* cl) {
     72   while (cl->args) {
     73     arg* a = cl->args;
     74     cl->args = a->next;
     75     gpr_free(a);
     76   }
     77   gpr_free(cl);
     78 }
     79 
     80 static void add_arg(gpr_cmdline* cl, const char* name, const char* help,
     81                     argtype type, void* value) {
     82   arg* a;
     83 
     84   for (a = cl->args; a; a = a->next) {
     85     GPR_ASSERT(0 != strcmp(a->name, name));
     86   }
     87 
     88   a = static_cast<arg*>(gpr_zalloc(sizeof(arg)));
     89   a->name = name;
     90   a->help = help;
     91   a->type = type;
     92   a->value = value;
     93   a->next = cl->args;
     94   cl->args = a;
     95 }
     96 
     97 void gpr_cmdline_add_int(gpr_cmdline* cl, const char* name, const char* help,
     98                          int* value) {
     99   add_arg(cl, name, help, ARGTYPE_INT, value);
    100 }
    101 
    102 void gpr_cmdline_add_flag(gpr_cmdline* cl, const char* name, const char* help,
    103                           int* value) {
    104   add_arg(cl, name, help, ARGTYPE_BOOL, value);
    105 }
    106 
    107 void gpr_cmdline_add_string(gpr_cmdline* cl, const char* name, const char* help,
    108                             const char** value) {
    109   add_arg(cl, name, help, ARGTYPE_STRING, value);
    110 }
    111 
    112 void gpr_cmdline_on_extra_arg(
    113     gpr_cmdline* cl, const char* name, const char* help,
    114     void (*on_extra_arg)(void* user_data, const char* arg), void* user_data) {
    115   GPR_ASSERT(!cl->extra_arg);
    116   GPR_ASSERT(on_extra_arg);
    117 
    118   cl->extra_arg = on_extra_arg;
    119   cl->extra_arg_user_data = user_data;
    120   cl->extra_arg_name = name;
    121   cl->extra_arg_help = help;
    122 }
    123 
    124 /* recursively descend argument list, adding the last element
    125    to s first - so that arguments are added in the order they were
    126    added to the list by api calls */
    127 static void add_args_to_usage(gpr_strvec* s, arg* a) {
    128   char* tmp;
    129 
    130   if (!a) return;
    131   add_args_to_usage(s, a->next);
    132 
    133   switch (a->type) {
    134     case ARGTYPE_BOOL:
    135       gpr_asprintf(&tmp, " [--%s|--no-%s]", a->name, a->name);
    136       gpr_strvec_add(s, tmp);
    137       break;
    138     case ARGTYPE_STRING:
    139       gpr_asprintf(&tmp, " [--%s=string]", a->name);
    140       gpr_strvec_add(s, tmp);
    141       break;
    142     case ARGTYPE_INT:
    143       gpr_asprintf(&tmp, " [--%s=int]", a->name);
    144       gpr_strvec_add(s, tmp);
    145       break;
    146   }
    147 }
    148 
    149 char* gpr_cmdline_usage_string(gpr_cmdline* cl, const char* argv0) {
    150   /* TODO(ctiller): make this prettier */
    151   gpr_strvec s;
    152   char* tmp;
    153   const char* name = strrchr(argv0, '/');
    154 
    155   if (name) {
    156     name++;
    157   } else {
    158     name = argv0;
    159   }
    160 
    161   gpr_strvec_init(&s);
    162 
    163   gpr_asprintf(&tmp, "Usage: %s", name);
    164   gpr_strvec_add(&s, tmp);
    165   add_args_to_usage(&s, cl->args);
    166   if (cl->extra_arg) {
    167     gpr_asprintf(&tmp, " [%s...]", cl->extra_arg_name);
    168     gpr_strvec_add(&s, tmp);
    169   }
    170   gpr_strvec_add(&s, gpr_strdup("\n"));
    171 
    172   tmp = gpr_strvec_flatten(&s, nullptr);
    173   gpr_strvec_destroy(&s);
    174   return tmp;
    175 }
    176 
    177 static int print_usage_and_die(gpr_cmdline* cl) {
    178   char* usage = gpr_cmdline_usage_string(cl, cl->argv0);
    179   fprintf(stderr, "%s", usage);
    180   gpr_free(usage);
    181   if (!cl->survive_failure) {
    182     exit(1);
    183   }
    184   return 0;
    185 }
    186 
    187 static int extra_state(gpr_cmdline* cl, char* str) {
    188   if (!cl->extra_arg) {
    189     return print_usage_and_die(cl);
    190   }
    191   cl->extra_arg(cl->extra_arg_user_data, str);
    192   return 1;
    193 }
    194 
    195 static arg* find_arg(gpr_cmdline* cl, char* name) {
    196   arg* a;
    197 
    198   for (a = cl->args; a; a = a->next) {
    199     if (0 == strcmp(a->name, name)) {
    200       break;
    201     }
    202   }
    203 
    204   if (!a) {
    205     fprintf(stderr, "Unknown argument: %s\n", name);
    206     return nullptr;
    207   }
    208 
    209   return a;
    210 }
    211 
    212 static int value_state(gpr_cmdline* cl, char* str) {
    213   long intval;
    214   char* end;
    215 
    216   GPR_ASSERT(cl->cur_arg);
    217 
    218   switch (cl->cur_arg->type) {
    219     case ARGTYPE_INT:
    220       intval = strtol(str, &end, 0);
    221       if (*end || intval < INT_MIN || intval > INT_MAX) {
    222         fprintf(stderr, "expected integer, got '%s' for %s\n", str,
    223                 cl->cur_arg->name);
    224         return print_usage_and_die(cl);
    225       }
    226       *static_cast<int*>(cl->cur_arg->value) = static_cast<int>(intval);
    227       break;
    228     case ARGTYPE_BOOL:
    229       if (0 == strcmp(str, "1") || 0 == strcmp(str, "true")) {
    230         *static_cast<int*>(cl->cur_arg->value) = 1;
    231       } else if (0 == strcmp(str, "0") || 0 == strcmp(str, "false")) {
    232         *static_cast<int*>(cl->cur_arg->value) = 0;
    233       } else {
    234         fprintf(stderr, "expected boolean, got '%s' for %s\n", str,
    235                 cl->cur_arg->name);
    236         return print_usage_and_die(cl);
    237       }
    238       break;
    239     case ARGTYPE_STRING:
    240       *static_cast<char**>(cl->cur_arg->value) = str;
    241       break;
    242   }
    243 
    244   cl->state = normal_state;
    245   return 1;
    246 }
    247 
    248 static int normal_state(gpr_cmdline* cl, char* str) {
    249   char* eq = nullptr;
    250   char* tmp = nullptr;
    251   char* arg_name = nullptr;
    252   int r = 1;
    253 
    254   if (0 == strcmp(str, "-help") || 0 == strcmp(str, "--help") ||
    255       0 == strcmp(str, "-h")) {
    256     return print_usage_and_die(cl);
    257   }
    258 
    259   cl->cur_arg = nullptr;
    260 
    261   if (str[0] == '-') {
    262     if (str[1] == '-') {
    263       if (str[2] == 0) {
    264         /* handle '--' to move to just extra args */
    265         cl->state = extra_state;
    266         return 1;
    267       }
    268       str += 2;
    269     } else {
    270       str += 1;
    271     }
    272     /* first byte of str is now past the leading '-' or '--' */
    273     if (str[0] == 'n' && str[1] == 'o' && str[2] == '-') {
    274       /* str is of the form '--no-foo' - it's a flag disable */
    275       str += 3;
    276       cl->cur_arg = find_arg(cl, str);
    277       if (cl->cur_arg == nullptr) {
    278         return print_usage_and_die(cl);
    279       }
    280       if (cl->cur_arg->type != ARGTYPE_BOOL) {
    281         fprintf(stderr, "%s is not a flag argument\n", str);
    282         return print_usage_and_die(cl);
    283       }
    284       *static_cast<int*>(cl->cur_arg->value) = 0;
    285       return 1; /* early out */
    286     }
    287     eq = strchr(str, '=');
    288     if (eq != nullptr) {
    289       /* copy the string into a temp buffer and extract the name */
    290       tmp = arg_name =
    291           static_cast<char*>(gpr_malloc(static_cast<size_t>(eq - str + 1)));
    292       memcpy(arg_name, str, static_cast<size_t>(eq - str));
    293       arg_name[eq - str] = 0;
    294     } else {
    295       arg_name = str;
    296     }
    297     cl->cur_arg = find_arg(cl, arg_name);
    298     if (cl->cur_arg == nullptr) {
    299       return print_usage_and_die(cl);
    300     }
    301     if (eq != nullptr) {
    302       /* str was of the type --foo=value, parse the value */
    303       r = value_state(cl, eq + 1);
    304     } else if (cl->cur_arg->type != ARGTYPE_BOOL) {
    305       /* flag types don't have a '--foo value' variant, other types do */
    306       cl->state = value_state;
    307     } else {
    308       /* flag parameter: just set the value */
    309       *static_cast<int*>(cl->cur_arg->value) = 1;
    310     }
    311   } else {
    312     r = extra_state(cl, str);
    313   }
    314 
    315   gpr_free(tmp);
    316   return r;
    317 }
    318 
    319 int gpr_cmdline_parse(gpr_cmdline* cl, int argc, char** argv) {
    320   int i;
    321 
    322   GPR_ASSERT(argc >= 1);
    323   cl->argv0 = argv[0];
    324 
    325   for (i = 1; i < argc; i++) {
    326     if (!cl->state(cl, argv[i])) {
    327       return 0;
    328     }
    329   }
    330   return 1;
    331 }
    332