git.m6w6.name Git - awesomized/libmemcached/blob - src/bin/common/options.hpp

   1 /*
   2     +--------------------------------------------------------------------+
   3     | libmemcached - C/C++ Client Library for memcached                  |
   4     +--------------------------------------------------------------------+
   5     | Redistribution and use in source and binary forms, with or without |
   6     | modification, are permitted under the terms of the BSD license.    |
   7     | You should have received a copy of the license in a bundled file   |
   8     | named LICENSE; in case you did not receive a copy you can review   |
   9     | the terms online at: https://opensource.org/licenses/BSD-3-Clause  |
  10     +--------------------------------------------------------------------+
  11     | Copyright (c) 2006-2014 Brian Aker   https://datadifferential.com/ |
  12     | Copyright (c) 2020 Michael Wallner   <mike@php.net>                |
  13     +--------------------------------------------------------------------+
  14 */
  15
  16 #pragma once
  17
  18 #include <cstdint>
  19 #include <climits>
  20 #include <getopt.h>
  21 #include <iostream>
  22 #include <string>
  23 #include <vector>
  24
  25 struct client_options {
  26   static const int first_flag = '/';
  27   static const int last_flag = 'z';
  28   static const int flag_count = last_flag - first_flag + 1;
  29
  30   enum class flag : int {
  31     _success = -1,
  32     _failure = '?',
  33
  34     absolute = '/',
  35     flush = '0',
  36     add_cmd = 'A',
  37     buffer = 'B',
  38     concurrency = 'C',
  39     flags = 'F',
  40     hash = 'H',
  41     load = 'L',
  42     tcp_nodelay = 'N',
  43     query = 'Q',
  44     replace_cmd = 'R',
  45     set_cmd = 'S',
  46     udp = 'U',
  47     version = 'V',
  48     analyze = 'a',
  49     binary = 'b',
  50     debug = 'd',
  51     expire = 'e',
  52     file = 'f',
  53     help = 'h',
  54     non_block = 'n',
  55     password = 'p',
  56     quiet = 'q',
  57     repeat = 'r',
  58     servers = 's',
  59     test = 't',
  60     username = 'u',
  61     verbose = 'v',
  62     zero = 'z',
  63   };
  64
  65   struct flags {
  66     int absolute;
  67     int flush;
  68     int add_cmd;
  69     int buffer;
  70     int flags;
  71     int tcp_nodelay;
  72     int replace_cmd;
  73     int set_cmd;
  74     int udp;
  75     int version;
  76     int analyze;
  77     int binary;
  78     int debug;
  79     int help;
  80     int non_block;
  81     int quiet;
  82     int verbose;
  83     int zero;
  84   } flags;
  85
  86   struct args {
  87     unsigned long concurrency;
  88     const char *hash;
  89     unsigned long load;
  90     const char *query;
  91     const char *file;
  92     unsigned long expire;
  93     const char *password;
  94     unsigned long repeat;
  95     const char *servers;
  96     const char *test;
  97     const char *username;
  98     const char * const *positional;
  99   } args;
 100
 101   option avail[flag_count];
 102   const char *help[flag_count];
 103
 104   std::string short_opts;
 105   std::vector<struct option> long_opts;
 106
 107   const char *progname;
 108   const char *progvers;
 109   const char *progdesc;
 110
 111   client_options(const char *prg, const char *ver, const char *dsc, const std::vector<client_options::flag> &f)
 112   : flags{}
 113   , args{}
 114   , avail{}
 115   , help{}
 116   , short_opts{}
 117   , long_opts{}
 118   , progname{prg}
 119   , progvers{ver}
 120   , progdesc{dsc}
 121   {
 122
 123 # define opt_flag(s, a, h) opt_flag_ex(#s, s, a, h)
 124 # define opt_flag_ex(n, v, a, h) \
 125   [static_cast<int>(flag::v)-first_flag] = {n, a, &flags.v, static_cast<int>(flag::v)}; \
 126     help[static_cast<int>(flag::v)-first_flag] = h
 127
 128 # define opt_data(s, a, h) opt_data_ex(#s, s, a, h)
 129 # define opt_data_ex(n, v, a, h) \
 130   [static_cast<int>(flag::v)-first_flag] = {n, a, nullptr, static_cast<int>(flag::v)}; \
 131     help[static_cast<int>(flag::v)-first_flag] = h
 132
 133     avail opt_flag(                 absolute,   no_argument,      "Use absolute path names.");
 134     avail opt_flag(                 flush,      no_argument,      "Flush the server(s).");
 135     avail opt_flag_ex("add",        add_cmd,    no_argument,      "Perform ADD operations.");
 136     avail opt_flag(                 buffer,     no_argument,      "Buffer requests.");
 137     avail opt_data(                 concurrency,required_argument,"Level of concurrency.");
 138     avail opt_flag(                 flags,      no_argument,      "Print or set associated flags for key(s).");
 139     avail opt_data(                 hash,       required_argument,"Select key hash.");
 140     avail opt_data(                 load,       required_argument,"Initial load.");
 141     avail opt_flag_ex("tcp-nodelay",tcp_nodelay,no_argument,      "Disable Nagle's algorithm.");
 142     avail opt_data(                 query,      no_argument,      "Query arguments.");
 143     avail opt_flag_ex("replace",    replace_cmd,no_argument,      "Perform REPLACE operations.");
 144     avail opt_flag_ex("set",        set_cmd,    no_argument,      "Perform SET operations.");
 145     avail opt_flag(                 udp,        no_argument,      "Use UDP.");
 146     avail opt_flag(                 version,    no_argument,      "Print program version.");
 147     avail opt_flag(                 analyze,    no_argument,      "Analyze the server's statistics.");
 148     avail opt_flag(                 binary,     no_argument,      "Use the binary memcached protocol.");
 149     avail opt_flag(                 debug,      no_argument,      "Print output useful only for debugging.");
 150     avail opt_data(                 expire,     required_argument,"Set associated expiry time for key(s).");
 151     avail opt_data(                 file,       required_argument,"File to save to or read from.");
 152     avail opt_flag(                 help,       no_argument,      "Print this help.");
 153     avail opt_flag_ex("non-block",  non_block,  no_argument,      "Use non blocking connections.");
 154     avail opt_data(                 password,   required_argument,"SASL password.");
 155     avail opt_flag(                 quiet,      no_argument,      "Print no output, not even errors.");
 156     avail opt_data(                 repeat,     required_argument,"Repeat operation n times.");
 157     avail opt_data(                 servers,    required_argument,"List of servers to connect to.");
 158     avail opt_data(                 test,       required_argument,"Test to perform.");
 159     avail opt_data(                 username,   required_argument,"SASL username.");
 160     avail opt_flag(                 verbose,    no_argument,      "Print more informational output.");
 161     avail opt_flag(                 zero,       no_argument,      "Zero.");
 162
 163     long_opts.reserve(f.size() + 1);
 164     short_opts.reserve(f.size() * 3);
 165
 166     for (auto o : f) {
 167       auto &opt = avail[static_cast<int>(o)-first_flag];
 168
 169       long_opts.push_back(opt);
 170       short_opts.push_back(opt.val);
 171
 172       for (int i = 0; i < opt.has_arg; ++i) {
 173         short_opts.push_back(':');
 174       }
 175     }
 176
 177     long_opts.push_back(option{});
 178   }
 179
 180   void printVersion() const {
 181     std::cout << progname << " v" << progvers
 182               << " (libmemcached v" << LIBMEMCACHED_VERSION_STRING << ")"
 183               << std::endl;
 184   }
 185
 186   void printHelp(const char *positional_args, const char *extra_doc = nullptr) const {
 187     printVersion();
 188     std::cout << "\n\t" << progdesc << "\n\n";
 189     std::cout << "Usage:\n\t" << progname << " -[";
 190     for (auto &opt : long_opts) {
 191       if (!opt.has_arg) {
 192         std::cout << (char) opt.val;
 193       }
 194     }
 195     std::cout << "] [--";
 196     for (auto &opt : long_opts) {
 197       if (opt.has_arg) {
 198         std::cout << (char) opt.val;
 199         if ((&opt)+1 != &*long_opts.rbegin())
 200         std::cout << '|';
 201       }
 202     }
 203     std::cout << " <arg>] ";
 204
 205     std::cout << positional_args << "\n\n";
 206     std::cout << "Options:\n";
 207     for (auto &opt : long_opts) {
 208       if (opt.name) {
 209         std::cout << "\t-" << (char) opt.val << "|--" << opt.name;
 210         if (opt.has_arg) {
 211           if (opt.has_arg == optional_argument) {
 212             std::cout << "[";
 213           }
 214           std::cout << "=arg";
 215           if (opt.has_arg == optional_argument) {
 216             std::cout << "]";
 217           }
 218         }
 219         std::cout << "\n\t\t" << help[opt.val - first_flag];
 220         std::cout << std::endl;
 221       }
 222     }
 223
 224     if (extra_doc) {
 225       std::cout << extra_doc << std::endl;
 226     }
 227   }
 228
 229   void printLastError(memcached_st *memc) const {
 230     if (!flags.quiet) {
 231       std::cerr << memcached_last_error_message(memc);
 232     }
 233   }
 234
 235   memcached_hash_t getHash() const {
 236     memcached_hash_t hash = MEMCACHED_HASH_DEFAULT;
 237     if (args.hash && *args.hash) {
 238       if (flags.verbose) {
 239         std::cerr << "Checking for hash '" << args.hash << "'.\n";
 240       }
 241       for (int h = hash; h < MEMCACHED_HASH_MAX; ++h) {
 242         auto hash_type = static_cast<memcached_hash_t>(h);
 243         std::string hash_string{libmemcached_string_hash(hash_type)};
 244         std::string hash_wanted{args.hash};
 245
 246         if (hash_wanted.length() == hash_string.length()) {
 247           auto ci = std::equal(hash_string.begin(), hash_string.end(), hash_wanted.begin(), [](int a, int b) {
 248             return std::tolower(a) == std::tolower(b);
 249           });
 250           if (ci) {
 251             hash = hash_type;
 252             break;
 253           }
 254         }
 255       }
 256       if (hash == MEMCACHED_HASH_DEFAULT) {
 257         if (!flags.quiet) {
 258           std::cerr << "Could not find hash '" << args.hash << "'.\n";
 259         }
 260       }
 261     }
 262     return hash;
 263   }
 264
 265   memcached_server_st *getServers() {
 266     if (!args.servers) {
 267       if (flags.verbose) {
 268         std::cerr << "Checking environment for a server list in MEMCACHED_SERVERS.\n";
 269       }
 270       args.servers = getenv("MEMCACHED_SERVERS");
 271       if (!args.servers || !*args.servers) {
 272         if (!flags.quiet) {
 273           std::cerr << "No servers provided.\n";
 274         }
 275         return nullptr;
 276       }
 277     }
 278
 279     auto servers = memcached_servers_parse(args.servers);
 280     if (!servers || !memcached_server_list_count(servers)) {
 281       if (!flags.quiet) {
 282         std::cerr << "Invalid server list provided: '" << args.servers << "'\n";
 283       }
 284       if (servers) {
 285         memcached_server_list_free(servers);
 286       }
 287       return nullptr;
 288     }
 289
 290     return servers;
 291   }
 292
 293   bool setupServers(memcached_st *memc) {
 294     auto servers = getServers();
 295     if (!servers) {
 296       return false;
 297     }
 298     if (MEMCACHED_SUCCESS != memcached_server_push(memc, servers)) {
 299       printLastError(memc);
 300       memcached_server_list_free(servers);
 301       return false;
 302     }
 303     memcached_server_list_free(servers);
 304     return true;
 305   }
 306
 307   bool setupBehavior(memcached_st *memc) const {
 308     if (MEMCACHED_SUCCESS != memcached_behavior_set_key_hash(memc, getHash())) {
 309       goto failure;
 310     }
 311     if (MEMCACHED_SUCCESS != memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_BINARY_PROTOCOL, flags.binary)) {
 312       goto failure;
 313     }
 314     if (MEMCACHED_SUCCESS != memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_BUFFER_REQUESTS, flags.buffer)) {
 315       goto failure;
 316     }
 317     if (MEMCACHED_SUCCESS != memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_TCP_NODELAY, flags.tcp_nodelay)) {
 318       goto failure;
 319     }
 320     if (MEMCACHED_SUCCESS != memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_NO_BLOCK, flags.non_block)) {
 321       goto failure;
 322     }
 323     if (MEMCACHED_SUCCESS != memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_USE_UDP, flags.udp)) {
 324       goto failure;
 325     }
 326     return true;
 327
 328   failure:
 329     printLastError(memc);
 330     return false;
 331   }
 332
 333   bool setupSASL(memcached_st *memc) const {
 334     if (args.username) {
 335       if (!LIBMEMCACHED_WITH_SASL_SUPPORT) {
 336         if (!flags.quiet) {
 337           std::cerr << "SASL username was supplied, but binary was not built with SASL support.\n";
 338           return false;
 339         }
 340       }
 341       if (MEMCACHED_SUCCESS != memcached_set_sasl_auth_data(memc, args.username, args.password)) {
 342         printLastError(memc);
 343         return false;
 344       }
 345     }
 346     return true;
 347   }
 348
 349   bool apply(memcached_st *memc) {
 350     if (!setupBehavior(memc)) {
 351       return false;
 352     }
 353     if (!setupServers(memc)) {
 354       return false;
 355     }
 356     if (!setupSASL(memc)) {
 357       return false;
 358     }
 359 #if defined(_WIN32)
 360     WSADATA wsaData;
 361     if (WSAStartup(MAKEWORD(2, 0), &wsaData)) {
 362       std::cerr << "Socket Initialization Error.\n";
 363       return false;
 364     }
 365 #endif // #if defined(_WIN32)
 366
 367     return true;
 368   }
 369
 370   bool parse(int argc, char *argv[]) {
 371     optind = 1;
 372
 373     while (true) {
 374       auto opt = getopt_long(argc, argv, short_opts.c_str(), long_opts.data(), nullptr);
 375
 376       if (flags.debug && opt != -1 && opt != '?') {
 377         std::cerr << "Processing option '" << (char) opt << "'\n";
 378       }
 379
 380       switch (static_cast<flag>(opt)) {
 381       case flag::_success:
 382         args.positional = &argv[optind];
 383         return true;
 384       case flag::_failure:
 385         return false;
 386       case flag::absolute:
 387         flags.absolute = true;
 388         break;
 389       case flag::flush:
 390         flags.flush = true;
 391         break;
 392       case flag::add_cmd:
 393         flags.add_cmd = true;
 394         flags.replace_cmd = false;
 395         flags.set_cmd = false;
 396         break;
 397       case flag::buffer:
 398         flags.buffer = true;
 399         break;
 400       case flag::concurrency:
 401         args.concurrency = std::stoul(optarg);
 402         break;
 403       case flag::flags:
 404         flags.flags = true;
 405         break;
 406       case flag::hash:
 407         args.hash = optarg;
 408         break;
 409       case flag::load:
 410         args.load = std::stoul(optarg);
 411         break;
 412       case flag::tcp_nodelay:
 413         flags.tcp_nodelay = true;
 414         break;
 415       case flag::query:
 416         args.query = optarg;
 417         break;
 418       case flag::replace_cmd:
 419         flags.add_cmd = false;
 420         flags.replace_cmd = true;
 421         flags.set_cmd = false;
 422         break;
 423       case flag::set_cmd:
 424         flags.add_cmd = false;
 425         flags.replace_cmd = false;
 426         flags.set_cmd = true;
 427         break;
 428       case flag::udp:
 429         flags.udp = true;
 430         break;
 431       case flag::version:
 432         flags.version = true;
 433         break;
 434       case flag::analyze:
 435         flags.analyze = true;
 436         break;
 437       case flag::binary:
 438         flags.binary = true;
 439         break;
 440       case flag::debug:
 441         flags.debug = true;
 442         flags.quiet = false;
 443         flags.verbose = true;
 444         break;
 445       case flag::expire:
 446         args.expire = std::stoul(optarg);
 447         break;
 448       case flag::file:
 449         args.file = optarg;
 450         break;
 451       case flag::help:
 452         flags.help = true;
 453         break;
 454       case flag::non_block:
 455         flags.non_block = true;
 456         break;
 457       case flag::password:
 458         args.password = optarg;
 459         break;
 460       case flag::quiet:
 461         flags.debug = false;
 462         flags.quiet = true;
 463         flags.verbose = false;
 464         break;
 465       case flag::repeat:
 466         args.repeat = std::stoul(optarg);
 467         break;
 468       case flag::servers:
 469         args.servers = optarg;
 470         break;
 471       case flag::test:
 472         args.test = optarg;
 473         break;
 474       case flag::username:
 475         args.username = optarg;
 476         break;
 477       case flag::verbose:
 478         flags.quiet = false;
 479         flags.verbose = true;
 480         break;
 481       case flag::zero:
 482         flags.zero = true;
 483         break;
 484       }
 485     }
 486   }
 487 };