From 9b0e362f4bf6efa125dbb1a4a70c3de7298c2ca8 Mon Sep 17 00:00:00 2001 From: Michael Wallner Date: Fri, 13 Nov 2020 16:06:24 +0100 Subject: [PATCH] bin: consolidate memcat --- src/bin/common/options.hpp | 487 +++++++++++++++++++++++++++++++++++++ src/bin/memcat.cc | 290 +++++++--------------- 2 files changed, 579 insertions(+), 198 deletions(-) create mode 100644 src/bin/common/options.hpp diff --git a/src/bin/common/options.hpp b/src/bin/common/options.hpp new file mode 100644 index 00000000..a5d6bdf7 --- /dev/null +++ b/src/bin/common/options.hpp @@ -0,0 +1,487 @@ +/* + +--------------------------------------------------------------------+ + | libmemcached - C/C++ Client Library for memcached | + +--------------------------------------------------------------------+ + | Redistribution and use in source and binary forms, with or without | + | modification, are permitted under the terms of the BSD license. | + | You should have received a copy of the license in a bundled file | + | named LICENSE; in case you did not receive a copy you can review | + | the terms online at: https://opensource.org/licenses/BSD-3-Clause | + +--------------------------------------------------------------------+ + | Copyright (c) 2006-2014 Brian Aker https://datadifferential.com/ | + | Copyright (c) 2020 Michael Wallner | + +--------------------------------------------------------------------+ +*/ + +#pragma once + +#include +#include +#include +#include +#include +#include + +struct client_options { + static const int first_flag = '/'; + static const int last_flag = 'z'; + static const int flag_count = last_flag - first_flag + 1; + + enum class flag : int { + _success = -1, + _failure = '?', + + absolute = '/', + flush = '0', + add_cmd = 'A', + buffer = 'B', + concurrency = 'C', + flags = 'F', + hash = 'H', + load = 'L', + tcp_nodelay = 'N', + query = 'Q', + replace_cmd = 'R', + set_cmd = 'S', + udp = 'U', + version = 'V', + analyze = 'a', + binary = 'b', + debug = 'd', + expire = 'e', + file = 'f', + help = 'h', + non_block = 'n', + password = 'p', + quiet = 'q', + repeat = 'r', + servers = 's', + test = 't', + username = 'u', + verbose = 'v', + zero = 'z', + }; + + struct flags { + int absolute; + int flush; + int add_cmd; + int buffer; + int flags; + int tcp_nodelay; + int replace_cmd; + int set_cmd; + int udp; + int version; + int analyze; + int binary; + int debug; + int help; + int non_block; + int quiet; + int verbose; + int zero; + } flags; + + struct args { + unsigned long concurrency; + const char *hash; + unsigned long load; + const char *query; + const char *file; + unsigned long expire; + const char *password; + unsigned long repeat; + const char *servers; + const char *test; + const char *username; + const char * const *positional; + } args; + + option avail[flag_count]; + const char *help[flag_count]; + + std::string short_opts; + std::vector long_opts; + + const char *progname; + const char *progvers; + const char *progdesc; + + client_options(const char *prg, const char *ver, const char *dsc, const std::vector &f) + : flags{} + , args{} + , avail{} + , help{} + , short_opts{} + , long_opts{} + , progname{prg} + , progvers{ver} + , progdesc{dsc} + { + +# define opt_flag(s, a, h) opt_flag_ex(#s, s, a, h) +# define opt_flag_ex(n, v, a, h) \ + [static_cast(flag::v)-first_flag] = {n, a, &flags.v, static_cast(flag::v)}; \ + help[static_cast(flag::v)-first_flag] = h + +# define opt_data(s, a, h) opt_data_ex(#s, s, a, h) +# define opt_data_ex(n, v, a, h) \ + [static_cast(flag::v)-first_flag] = {n, a, nullptr, static_cast(flag::v)}; \ + help[static_cast(flag::v)-first_flag] = h + + avail opt_flag( absolute, no_argument, "Use absolute path names."); + avail opt_flag( flush, no_argument, "Flush the server(s)."); + avail opt_flag_ex("add", add_cmd, no_argument, "Perform ADD operations."); + avail opt_flag( buffer, no_argument, "Buffer requests."); + avail opt_data( concurrency,required_argument,"Level of concurrency."); + avail opt_flag( flags, no_argument, "Print or set associated flags for key(s)."); + avail opt_data( hash, required_argument,"Select key hash."); + avail opt_data( load, required_argument,"Initial load."); + avail opt_flag_ex("tcp-nodelay",tcp_nodelay,no_argument, "Disable Nagle's algorithm."); + avail opt_data( query, no_argument, "Query arguments."); + avail opt_flag_ex("replace", replace_cmd,no_argument, "Perform REPLACE operations."); + avail opt_flag_ex("set", set_cmd, no_argument, "Perform SET operations."); + avail opt_flag( udp, no_argument, "Use UDP."); + avail opt_flag( version, no_argument, "Print program version."); + avail opt_flag( analyze, no_argument, "Analyze the server's statistics."); + avail opt_flag( binary, no_argument, "Use the binary memcached protocol."); + avail opt_flag( debug, no_argument, "Print output useful only for debugging."); + avail opt_data( expire, required_argument,"Set associated expiry time for key(s)."); + avail opt_data( file, required_argument,"File to save to or read from."); + avail opt_flag( help, no_argument, "Print this help."); + avail opt_flag_ex("non-block", non_block, no_argument, "Use non blocking connections."); + avail opt_data( password, required_argument,"SASL password."); + avail opt_flag( quiet, no_argument, "Print no output, not even errors."); + avail opt_data( repeat, required_argument,"Repeat operation n times."); + avail opt_data( servers, required_argument,"List of servers to connect to."); + avail opt_data( test, required_argument,"Test to perform."); + avail opt_data( username, required_argument,"SASL username."); + avail opt_flag( verbose, no_argument, "Print more informational output."); + avail opt_flag( zero, no_argument, "Zero."); + + long_opts.reserve(f.size() + 1); + short_opts.reserve(f.size() * 3); + + for (auto o : f) { + auto &opt = avail[static_cast(o)-first_flag]; + + long_opts.push_back(opt); + short_opts.push_back(opt.val); + + for (int i = 0; i < opt.has_arg; ++i) { + short_opts.push_back(':'); + } + } + + long_opts.push_back(option{}); + } + + void printVersion() const { + std::cout << progname << " v" << progvers + << " (libmemcached v" << LIBMEMCACHED_VERSION_STRING << ")" + << std::endl; + } + + void printHelp(const char *positional_args, const char *extra_doc = nullptr) const { + printVersion(); + std::cout << "\n\t" << progdesc << "\n\n"; + std::cout << "Usage:\n\t" << progname << " -["; + for (auto &opt : long_opts) { + if (!opt.has_arg) { + std::cout << (char) opt.val; + } + } + std::cout << "] [--"; + for (auto &opt : long_opts) { + if (opt.has_arg) { + std::cout << (char) opt.val; + if ((&opt)+1 != &*long_opts.rbegin()) + std::cout << '|'; + } + } + std::cout << " ] "; + + std::cout << positional_args << "\n\n"; + std::cout << "Options:\n"; + for (auto &opt : long_opts) { + if (opt.name) { + std::cout << "\t-" << (char) opt.val << "|--" << opt.name; + if (opt.has_arg) { + if (opt.has_arg == optional_argument) { + std::cout << "["; + } + std::cout << "=arg"; + if (opt.has_arg == optional_argument) { + std::cout << "]"; + } + } + std::cout << "\n\t\t" << help[opt.val - first_flag]; + std::cout << std::endl; + } + } + + if (extra_doc) { + std::cout << extra_doc << std::endl; + } + } + + void printLastError(memcached_st *memc) const { + if (!flags.quiet) { + std::cerr << memcached_last_error_message(memc); + } + } + + memcached_hash_t getHash() const { + memcached_hash_t hash = MEMCACHED_HASH_DEFAULT; + if (args.hash && *args.hash) { + if (flags.verbose) { + std::cerr << "Checking for hash '" << args.hash << "'.\n"; + } + for (int h = hash; h < MEMCACHED_HASH_MAX; ++h) { + auto hash_type = static_cast(h); + std::string hash_string{libmemcached_string_hash(hash_type)}; + std::string hash_wanted{args.hash}; + + if (hash_wanted.length() == hash_string.length()) { + auto ci = std::equal(hash_string.begin(), hash_string.end(), hash_wanted.begin(), [](int a, int b) { + return std::tolower(a) == std::tolower(b); + }); + if (ci) { + hash = hash_type; + break; + } + } + } + if (hash == MEMCACHED_HASH_DEFAULT) { + if (!flags.quiet) { + std::cerr << "Could not find hash '" << args.hash << "'.\n"; + } + } + } + return hash; + } + + memcached_server_st *getServers() { + if (!args.servers) { + if (flags.verbose) { + std::cerr << "Checking environment for a server list in MEMCACHED_SERVERS.\n"; + } + args.servers = getenv("MEMCACHED_SERVERS"); + if (!args.servers || !*args.servers) { + if (!flags.quiet) { + std::cerr << "No servers provided.\n"; + } + return nullptr; + } + } + + auto servers = memcached_servers_parse(args.servers); + if (!servers || !memcached_server_list_count(servers)) { + if (!flags.quiet) { + std::cerr << "Invalid server list provided: '" << args.servers << "'\n"; + } + if (servers) { + memcached_server_list_free(servers); + } + return nullptr; + } + + return servers; + } + + bool setupServers(memcached_st *memc) { + auto servers = getServers(); + if (!servers) { + return false; + } + if (MEMCACHED_SUCCESS != memcached_server_push(memc, servers)) { + printLastError(memc); + memcached_server_list_free(servers); + return false; + } + memcached_server_list_free(servers); + return true; + } + + bool setupBehavior(memcached_st *memc) const { + if (MEMCACHED_SUCCESS != memcached_behavior_set_key_hash(memc, getHash())) { + goto failure; + } + if (MEMCACHED_SUCCESS != memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_BINARY_PROTOCOL, flags.binary)) { + goto failure; + } + if (MEMCACHED_SUCCESS != memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_BUFFER_REQUESTS, flags.buffer)) { + goto failure; + } + if (MEMCACHED_SUCCESS != memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_TCP_NODELAY, flags.tcp_nodelay)) { + goto failure; + } + if (MEMCACHED_SUCCESS != memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_NO_BLOCK, flags.non_block)) { + goto failure; + } + if (MEMCACHED_SUCCESS != memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_USE_UDP, flags.udp)) { + goto failure; + } + return true; + + failure: + printLastError(memc); + return false; + } + + bool setupSASL(memcached_st *memc) const { + if (args.username) { + if (!LIBMEMCACHED_WITH_SASL_SUPPORT) { + if (!flags.quiet) { + std::cerr << "SASL username was supplied, but binary was not built with SASL support.\n"; + return false; + } + } + if (MEMCACHED_SUCCESS != memcached_set_sasl_auth_data(memc, args.username, args.password)) { + printLastError(memc); + return false; + } + } + return true; + } + + bool apply(memcached_st *memc) { + if (!setupBehavior(memc)) { + return false; + } + if (!setupServers(memc)) { + return false; + } + if (!setupSASL(memc)) { + return false; + } +#if defined(_WIN32) + WSADATA wsaData; + if (WSAStartup(MAKEWORD(2, 0), &wsaData)) { + std::cerr << "Socket Initialization Error.\n"; + return false; + } +#endif // #if defined(_WIN32) + + return true; + } + + bool parse(int argc, char *argv[]) { + optind = 1; + + while (true) { + auto opt = getopt_long(argc, argv, short_opts.c_str(), long_opts.data(), nullptr); + + if (flags.debug && opt != -1 && opt != '?') { + std::cerr << "Processing option '" << (char) opt << "'\n"; + } + + switch (static_cast(opt)) { + case flag::_success: + args.positional = &argv[optind]; + return true; + case flag::_failure: + return false; + case flag::absolute: + flags.absolute = true; + break; + case flag::flush: + flags.flush = true; + break; + case flag::add_cmd: + flags.add_cmd = true; + flags.replace_cmd = false; + flags.set_cmd = false; + break; + case flag::buffer: + flags.buffer = true; + break; + case flag::concurrency: + args.concurrency = std::stoul(optarg); + break; + case flag::flags: + flags.flags = true; + break; + case flag::hash: + args.hash = optarg; + break; + case flag::load: + args.load = std::stoul(optarg); + break; + case flag::tcp_nodelay: + flags.tcp_nodelay = true; + break; + case flag::query: + args.query = optarg; + break; + case flag::replace_cmd: + flags.add_cmd = false; + flags.replace_cmd = true; + flags.set_cmd = false; + break; + case flag::set_cmd: + flags.add_cmd = false; + flags.replace_cmd = false; + flags.set_cmd = true; + break; + case flag::udp: + flags.udp = true; + break; + case flag::version: + flags.version = true; + break; + case flag::analyze: + flags.analyze = true; + break; + case flag::binary: + flags.binary = true; + break; + case flag::debug: + flags.debug = true; + flags.quiet = false; + flags.verbose = true; + break; + case flag::expire: + args.expire = std::stoul(optarg); + break; + case flag::file: + args.file = optarg; + break; + case flag::help: + flags.help = true; + break; + case flag::non_block: + flags.non_block = true; + break; + case flag::password: + args.password = optarg; + break; + case flag::quiet: + flags.debug = false; + flags.quiet = true; + flags.verbose = false; + break; + case flag::repeat: + args.repeat = std::stoul(optarg); + break; + case flag::servers: + args.servers = optarg; + break; + case flag::test: + args.test = optarg; + break; + case flag::username: + args.username = optarg; + break; + case flag::verbose: + flags.quiet = false; + flags.verbose = true; + break; + case flag::zero: + flags.zero = true; + break; + } + } + } +}; diff --git a/src/bin/memcat.cc b/src/bin/memcat.cc index ca69d6f9..00519b52 100644 --- a/src/bin/memcat.cc +++ b/src/bin/memcat.cc @@ -15,231 +15,125 @@ #include "mem_config.h" -#include -#include -#include -#include -#include -#include "libmemcached-1.0/memcached.h" - -#include "utilities.h" - #define PROGRAM_NAME "memcat" #define PROGRAM_DESCRIPTION "Cat a set of key values to stdout." +#define PROGRAM_VERSION "1.1" -/* Prototypes */ -void options_parse(int argc, char *argv[]); +#include "libmemcached/common.h" +#include "common/options.hpp" -static int opt_binary = 0; -static int opt_verbose = 0; -static int opt_displayflag = 0; -static char *opt_servers = NULL; -static char *opt_hash = NULL; -static char *opt_username; -static char *opt_passwd; -static char *opt_file; +#include +#include -int main(int argc, char *argv[]) { - char *string; - size_t string_length; - uint32_t flags; +memcached_return_t memcat(const client_options &opt, memcached_st *memc, const char *key, std::ostream &ref) { memcached_return_t rc; + uint32_t flags; + size_t len; + auto val = memcached_get(memc, key, strlen(key), &len, &flags, &rc); - int return_code = EXIT_SUCCESS; - - options_parse(argc, argv); - initialize_sockets(); - - if (opt_servers == NULL) { - char *temp; - - if ((temp = getenv("MEMCACHED_SERVERS"))) { - opt_servers = strdup(temp); + if (MEMCACHED_SUCCESS == rc) { + if (opt.flags.verbose) { + ref << "key: " << key << "\n"; } - - if (opt_servers == NULL) { - std::cerr << "No servers provided" << std::endl; - exit(EXIT_FAILURE); + if (opt.flags.flags) { + ref << "flags: " << flags << "\n"; + } + if (opt.flags.verbose) { + ref << "value: "; } + ref.write(val, len); + ref << std::endl; } - memcached_server_st *servers = memcached_servers_parse(opt_servers); - if (servers == NULL or memcached_server_list_count(servers) == 0) { - std::cerr << "Invalid server list provided:" << opt_servers << std::endl; - return EXIT_FAILURE; + if (val) { + free(val); } + return rc; +} - memcached_st *memc = memcached_create(NULL); - process_hash_option(memc, opt_hash); - - memcached_server_push(memc, servers); - memcached_server_list_free(servers); - memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_BINARY_PROTOCOL, (uint64_t) opt_binary); +int main(int argc, char *argv[]) { + memcached_st memc; + client_options opt{PROGRAM_NAME, PROGRAM_VERSION, PROGRAM_DESCRIPTION, { + client_options::flag::help, + client_options::flag::version, + client_options::flag::verbose, + client_options::flag::debug, + client_options::flag::quiet, + client_options::flag::servers, + client_options::flag::binary, + client_options::flag::username, + client_options::flag::password, + client_options::flag::hash, + client_options::flag::flags, + client_options::flag::file, + }}; + + if (!opt.parse(argc, argv)) { + exit(EXIT_FAILURE); + } + if (opt.flags.help) { + opt.printHelp("key [ key ... ]"); + exit(EXIT_SUCCESS); + } + if (opt.flags.version) { + opt.printVersion(); + exit(EXIT_SUCCESS); + } - if (opt_username and LIBMEMCACHED_WITH_SASL_SUPPORT == 0) { - memcached_free(memc); - std::cerr << "--username was supplied, but binary was not built with SASL support." - << std::endl; - return EXIT_FAILURE; + if (opt.flags.quiet && !opt.args.file) { + std::cerr << "--quiet operation was requested, but --file was not set.\n"; + exit(EXIT_FAILURE); } - if (opt_username) { - memcached_return_t ret; - if (memcached_failed(ret = memcached_set_sasl_auth_data(memc, opt_username, opt_passwd))) { - std::cerr << memcached_last_error_message(memc) << std::endl; - memcached_free(memc); - return EXIT_FAILURE; + if (!memcached_create(&memc)) { + if (!opt.flags.quiet) { + std::cerr << "Failed to initialize memcached client.\n"; } + exit(EXIT_FAILURE); } - while (optind < argc) { - string = memcached_get(memc, argv[optind], strlen(argv[optind]), &string_length, &flags, &rc); - if (rc == MEMCACHED_SUCCESS) { - if (opt_displayflag) { - if (opt_verbose) { - std::cout << "key: " << argv[optind] << std::endl << "flags: "; - } - std::cout << flags << std::endl; - } else { - if (opt_verbose) { - std::cout << "key: " << argv[optind] << std::endl - << "flags: " << flags << std::endl - << "length: " << string_length << std::endl - << "value: "; - } + if (!opt.apply(&memc)) { + memcached_free(&memc); + exit(EXIT_FAILURE); + } - if (opt_file) { - FILE *fp = fopen(opt_file, "w"); - if (fp == NULL) { - perror("fopen"); - return_code = EXIT_FAILURE; - break; - } + if (!*opt.args.positional) { + if (!opt.flags.quiet) { + std::cerr << "No key(s) provided.\n"; + memcached_free(&memc); + exit(EXIT_FAILURE); + } + } - size_t written = fwrite(string, 1, string_length, fp); - if (written != string_length) { - std::cerr << "error writing file to file " << opt_file << " wrote " << written - << ", should have written" << string_length << std::endl; - return_code = EXIT_FAILURE; - break; - } + auto exit_code = EXIT_SUCCESS; + for (auto arg = opt.args.positional; *arg; ++arg) { + auto key = *arg; + if (*key) { + memcached_return_t rc; + if (opt.args.file && *opt.args.file) { + std::ofstream file{opt.args.file, std::ios::binary}; + rc = memcat(opt, &memc, key, file); + } else { + rc = memcat(opt, &memc, key, std::cout); + } + if (MEMCACHED_SUCCESS != rc) { + exit_code = EXIT_FAILURE; - if (fclose(fp)) { - std::cerr << "error closing " << opt_file << std::endl; - return_code = EXIT_FAILURE; - break; + if (MEMCACHED_NOTFOUND == rc) { + if (opt.flags.verbose) { + std::cerr << "not found: " << key << "\n"; } + // continue; } else { - std::cout.write(string, string_length); - std::cout << std::endl; + if (!opt.flags.quiet) { + std::cerr << memcached_last_error_message(&memc) << "\n"; + } + break; } } - } else if (rc != MEMCACHED_NOTFOUND) { - std::cerr << "error on " << argv[optind] << "(" << memcached_strerror(memc, rc) << ")"; - if (memcached_last_error_errno(memc)) { - std::cerr << " system error (" << strerror(memcached_last_error_errno(memc)) << ")" - << std::endl; - } - std::cerr << std::endl; - - return_code = EXIT_FAILURE; - break; - } else // Unknown Issue - { - std::cerr << "error on " << argv[optind] << "(" << memcached_strerror(NULL, rc) << ")" - << std::endl; - return_code = EXIT_FAILURE; } - optind++; - free(string); - string = nullptr; - } - if (string) { - free(string); - } - - memcached_free(memc); - - if (opt_servers) { - free(opt_servers); - } - if (opt_hash) { - free(opt_hash); } - return return_code; -} - -void options_parse(int argc, char *argv[]) { - int option_index = 0; - - memcached_programs_help_st help_options[] = { - {0}, - }; - - static struct option long_options[] = { - {(OPTIONSTRING) "version", no_argument, NULL, OPT_VERSION}, - {(OPTIONSTRING) "help", no_argument, NULL, OPT_HELP}, - {(OPTIONSTRING) "quiet", no_argument, NULL, OPT_QUIET}, - {(OPTIONSTRING) "verbose", no_argument, &opt_verbose, OPT_VERBOSE}, - {(OPTIONSTRING) "debug", no_argument, &opt_verbose, OPT_DEBUG}, - {(OPTIONSTRING) "servers", required_argument, NULL, OPT_SERVERS}, - {(OPTIONSTRING) "flag", no_argument, &opt_displayflag, OPT_FLAG}, - {(OPTIONSTRING) "hash", required_argument, NULL, OPT_HASH}, - {(OPTIONSTRING) "binary", no_argument, NULL, OPT_BINARY}, - {(OPTIONSTRING) "username", required_argument, NULL, OPT_USERNAME}, - {(OPTIONSTRING) "password", required_argument, NULL, OPT_PASSWD}, - {(OPTIONSTRING) "file", required_argument, NULL, OPT_FILE}, - {0, 0, 0, 0}, - }; - - while (1) { - int option_rv = getopt_long(argc, argv, "Vhvds:", long_options, &option_index); - if (option_rv == -1) - break; - switch (option_rv) { - case 0: - break; - case OPT_BINARY: - opt_binary = 1; - break; - case OPT_VERBOSE: /* --verbose or -v */ - opt_verbose = OPT_VERBOSE; - break; - case OPT_DEBUG: /* --debug or -d */ - opt_verbose = OPT_DEBUG; - break; - case OPT_VERSION: /* --version or -V */ - version_command(PROGRAM_NAME); - break; - case OPT_HELP: /* --help or -h */ - help_command(PROGRAM_NAME, PROGRAM_DESCRIPTION, long_options, help_options); - break; - case OPT_SERVERS: /* --servers or -s */ - opt_servers = strdup(optarg); - break; - case OPT_HASH: - opt_hash = strdup(optarg); - break; - case OPT_USERNAME: - opt_username = optarg; - break; - case OPT_PASSWD: - opt_passwd = optarg; - break; - case OPT_FILE: - opt_file = optarg; - break; - - case OPT_QUIET: - close_stdio(); - break; - - case '?': - /* getopt_long already printed an error message. */ - exit(EXIT_FAILURE); - default: - abort(); - } - } + memcached_free(&memc); + exit(exit_code); } -- 2.30.2