git.m6w6.name Git - awesomized/libmemcached/blob - src/bin/memcp.cc

   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 #include "mem_config.h"
  17
  18 #include <cerrno>
  19 #include <climits>
  20 #include <cstdio>
  21 #include <cstdlib>
  22 #include <cstring>
  23 #include <fcntl.h>
  24 #include <getopt.h>
  25 #include <iostream>
  26 #ifdef HAVE_STRINGS_H
  27 #  include <strings.h>
  28 #endif
  29 #include <sys/stat.h>
  30 #include <sys/types.h>
  31 #include <unistd.h>
  32
  33 #include "libmemcached-1.0/memcached.h"
  34
  35 #include "client_options.h"
  36 #include "utilities.h"
  37
  38 #define PROGRAM_NAME        "memcp"
  39 #define PROGRAM_DESCRIPTION "Copy a set of files to a memcached cluster."
  40
  41 /* Prototypes */
  42 static void options_parse(int argc, char *argv[]);
  43
  44 static bool opt_binary = false;
  45 static bool opt_udp = false;
  46 static bool opt_buffer = false;
  47 static int opt_verbose = 0;
  48 static char *opt_servers = NULL;
  49 static char *opt_hash = NULL;
  50 static int opt_method = OPT_SET;
  51 static uint32_t opt_flags = 0;
  52 static time_t opt_expires = 0;
  53 static char *opt_username;
  54 static char *opt_passwd;
  55
  56 static long strtol_wrapper(const char *nptr, int base, bool *error) {
  57   long val;
  58   char *endptr;
  59
  60   errno = 0; /* To distinguish success/failure after call */
  61   val = strtol(nptr, &endptr, base);
  62
  63   /* Check for various possible errors */
  64
  65   if ((errno == ERANGE and (val == LONG_MAX or val == LONG_MIN)) or (errno && val == 0)) {
  66     *error = true;
  67     return 0;
  68   }
  69
  70   if (endptr == nptr) {
  71     *error = true;
  72     return 0;
  73   }
  74
  75   *error = false;
  76   return val;
  77 }
  78
  79 int main(int argc, char *argv[]) {
  80   options_parse(argc, argv);
  81
  82   if (optind >= argc) {
  83     fprintf(stderr, "Expected argument after options\n");
  84     exit(EXIT_FAILURE);
  85   }
  86
  87   initialize_sockets();
  88
  89   memcached_st *memc = memcached_create(NULL);
  90
  91   if (opt_udp) {
  92     if (opt_verbose) {
  93       std::cout << "Enabling UDP" << std::endl;
  94     }
  95
  96     if (memcached_failed(memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_USE_UDP, opt_udp))) {
  97       memcached_free(memc);
  98       std::cerr << "Could not enable UDP protocol." << std::endl;
  99       return EXIT_FAILURE;
 100     }
 101   }
 102
 103   if (opt_buffer) {
 104     if (opt_verbose) {
 105       std::cout << "Enabling MEMCACHED_BEHAVIOR_BUFFER_REQUESTS" << std::endl;
 106     }
 107
 108     if (memcached_failed(
 109             memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_BUFFER_REQUESTS, opt_buffer))) {
 110       memcached_free(memc);
 111       std::cerr << "Could not enable MEMCACHED_BEHAVIOR_BUFFER_REQUESTS." << std::endl;
 112       return EXIT_FAILURE;
 113     }
 114   }
 115
 116   process_hash_option(memc, opt_hash);
 117
 118   if (opt_servers == NULL) {
 119     char *temp;
 120
 121     if ((temp = getenv("MEMCACHED_SERVERS"))) {
 122       opt_servers = strdup(temp);
 123     }
 124 #if 0
 125     else if (argc >= 1 and argv[--argc])
 126     {
 127       opt_servers= strdup(argv[argc]);
 128     }
 129 #endif
 130
 131     if (opt_servers == NULL) {
 132       std::cerr << "No Servers provided" << std::endl;
 133       exit(EXIT_FAILURE);
 134     }
 135   }
 136
 137   memcached_server_st *servers = memcached_servers_parse(opt_servers);
 138   if (servers == NULL or memcached_server_list_count(servers) == 0) {
 139     std::cerr << "Invalid server list provided:" << opt_servers << std::endl;
 140     return EXIT_FAILURE;
 141   }
 142
 143   memcached_server_push(memc, servers);
 144   memcached_server_list_free(servers);
 145   memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_BINARY_PROTOCOL, opt_binary);
 146   if (opt_username and LIBMEMCACHED_WITH_SASL_SUPPORT == 0) {
 147     memcached_free(memc);
 148     std::cerr << "--username was supplied, but binary was not built with SASL support."
 149               << std::endl;
 150     return EXIT_FAILURE;
 151   }
 152
 153   if (opt_username) {
 154     memcached_return_t ret;
 155     if (memcached_failed(ret = memcached_set_sasl_auth_data(memc, opt_username, opt_passwd))) {
 156       std::cerr << memcached_last_error_message(memc) << std::endl;
 157       memcached_free(memc);
 158       return EXIT_FAILURE;
 159     }
 160   }
 161
 162   int exit_code = EXIT_SUCCESS;
 163   while (optind < argc) {
 164     int fd = open(argv[optind], O_RDONLY);
 165     if (fd < 0) {
 166       std::cerr << "memcp " << argv[optind] << " " << strerror(errno) << std::endl;
 167       optind++;
 168       exit_code = EXIT_FAILURE;
 169       continue;
 170     }
 171
 172     struct stat sbuf;
 173     if (fstat(fd, &sbuf) == -1) {
 174       std::cerr << "memcp " << argv[optind] << " " << strerror(errno) << std::endl;
 175       optind++;
 176       exit_code = EXIT_FAILURE;
 177       continue;
 178     }
 179
 180     char *ptr = rindex(argv[optind], '/');
 181     if (ptr) {
 182       ptr++;
 183     } else {
 184       ptr = argv[optind];
 185     }
 186
 187     if (opt_verbose) {
 188       static const char *opstr[] = {"set", "add", "replace"};
 189       printf("op: %s\nsource file: %s\nlength: %lu\n"
 190              "key: %s\nflags: %x\nexpires: %lu\n",
 191              opstr[opt_method - OPT_SET], argv[optind], (unsigned long) sbuf.st_size, ptr,
 192              opt_flags, (unsigned long) opt_expires);
 193     }
 194
 195     // The file may be empty
 196     char *file_buffer_ptr = NULL;
 197     if (sbuf.st_size > 0) {
 198       if ((file_buffer_ptr = (char *) malloc(sizeof(char) * (size_t) sbuf.st_size)) == NULL) {
 199         std::cerr << "Error allocating file buffer(" << strerror(errno) << ")" << std::endl;
 200         close(fd);
 201         exit(EXIT_FAILURE);
 202       }
 203
 204       ssize_t read_length;
 205       if ((read_length = ::read(fd, file_buffer_ptr, (size_t) sbuf.st_size)) == -1) {
 206         std::cerr << "Error while reading file " << file_buffer_ptr << " (" << strerror(errno)
 207                   << ")" << std::endl;
 208         close(fd);
 209         free(file_buffer_ptr);
 210         exit(EXIT_FAILURE);
 211       }
 212
 213       if (read_length != sbuf.st_size) {
 214         std::cerr << "Failure while reading file. Read length was not equal to stat() length"
 215                   << std::endl;
 216         close(fd);
 217         free(file_buffer_ptr);
 218         exit(EXIT_FAILURE);
 219       }
 220     }
 221
 222     memcached_return_t rc;
 223     if (opt_method == OPT_ADD) {
 224       rc = memcached_add(memc, ptr, strlen(ptr), file_buffer_ptr, (size_t) sbuf.st_size,
 225                          opt_expires, opt_flags);
 226     } else if (opt_method == OPT_REPLACE) {
 227       rc = memcached_replace(memc, ptr, strlen(ptr), file_buffer_ptr, (size_t) sbuf.st_size,
 228                              opt_expires, opt_flags);
 229     } else {
 230       rc = memcached_set(memc, ptr, strlen(ptr), file_buffer_ptr, (size_t) sbuf.st_size,
 231                          opt_expires, opt_flags);
 232     }
 233
 234     if (memcached_failed(rc)) {
 235       std::cerr << "Error occrrured during memcached_set(): " << memcached_last_error_message(memc)
 236                 << std::endl;
 237       exit_code = EXIT_FAILURE;
 238     }
 239
 240     ::free(file_buffer_ptr);
 241     ::close(fd);
 242     optind++;
 243   }
 244
 245   if (opt_verbose) {
 246     std::cout << "Calling memcached_free()" << std::endl;
 247   }
 248
 249   memcached_free(memc);
 250
 251   if (opt_servers) {
 252     free(opt_servers);
 253   }
 254
 255   if (opt_hash) {
 256     free(opt_hash);
 257   }
 258
 259   return exit_code;
 260 }
 261
 262 static void options_parse(int argc, char *argv[]) {
 263   memcached_programs_help_st help_options[] = {
 264       {0},
 265   };
 266
 267   static struct option long_options[] = {
 268       {(OPTIONSTRING) "version", no_argument, NULL, OPT_VERSION},
 269       {(OPTIONSTRING) "help", no_argument, NULL, OPT_HELP},
 270       {(OPTIONSTRING) "quiet", no_argument, NULL, OPT_QUIET},
 271       {(OPTIONSTRING) "udp", no_argument, NULL, OPT_UDP},
 272       {(OPTIONSTRING) "buffer", no_argument, NULL, OPT_BUFFER},
 273       {(OPTIONSTRING) "verbose", no_argument, &opt_verbose, OPT_VERBOSE},
 274       {(OPTIONSTRING) "debug", no_argument, &opt_verbose, OPT_DEBUG},
 275       {(OPTIONSTRING) "servers", required_argument, NULL, OPT_SERVERS},
 276       {(OPTIONSTRING) "flag", required_argument, NULL, OPT_FLAG},
 277       {(OPTIONSTRING) "expire", required_argument, NULL, OPT_EXPIRE},
 278       {(OPTIONSTRING) "set", no_argument, NULL, OPT_SET},
 279       {(OPTIONSTRING) "add", no_argument, NULL, OPT_ADD},
 280       {(OPTIONSTRING) "replace", no_argument, NULL, OPT_REPLACE},
 281       {(OPTIONSTRING) "hash", required_argument, NULL, OPT_HASH},
 282       {(OPTIONSTRING) "binary", no_argument, NULL, OPT_BINARY},
 283       {(OPTIONSTRING) "username", required_argument, NULL, OPT_USERNAME},
 284       {(OPTIONSTRING) "password", required_argument, NULL, OPT_PASSWD},
 285       {0, 0, 0, 0},
 286   };
 287
 288   bool opt_version = false;
 289   bool opt_help = false;
 290   int option_index = 0;
 291
 292   while (1) {
 293     int option_rv = getopt_long(argc, argv, "Vhvds:", long_options, &option_index);
 294
 295     if (option_rv == -1)
 296       break;
 297
 298     switch (option_rv) {
 299     case 0:
 300       break;
 301
 302     case OPT_BINARY:
 303       opt_binary = true;
 304       break;
 305
 306     case OPT_VERBOSE: /* --verbose or -v */
 307       opt_verbose = OPT_VERBOSE;
 308       break;
 309
 310     case OPT_DEBUG: /* --debug or -d */
 311       opt_verbose = OPT_DEBUG;
 312       break;
 313
 314     case OPT_VERSION: /* --version or -V */
 315       opt_version = true;
 316       break;
 317
 318     case OPT_HELP: /* --help or -h */
 319       opt_help = true;
 320       break;
 321
 322     case OPT_SERVERS: /* --servers or -s */
 323       opt_servers = strdup(optarg);
 324       break;
 325
 326     case OPT_FLAG: /* --flag */
 327     {
 328       bool strtol_error;
 329       opt_flags = (uint32_t) strtol_wrapper(optarg, 16, &strtol_error);
 330       if (strtol_error == true) {
 331         fprintf(stderr, "Bad value passed via --flag\n");
 332         exit(1);
 333       }
 334     } break;
 335
 336     case OPT_EXPIRE: /* --expire */
 337     {
 338       bool strtol_error;
 339       opt_expires = (time_t) strtol_wrapper(optarg, 10, &strtol_error);
 340       if (strtol_error == true) {
 341         fprintf(stderr, "Bad value passed via --expire\n");
 342         exit(1);
 343       }
 344     } break;
 345
 346     case OPT_SET:
 347       opt_method = OPT_SET;
 348       break;
 349
 350     case OPT_REPLACE:
 351       opt_method = OPT_REPLACE;
 352       break;
 353
 354     case OPT_ADD:
 355       opt_method = OPT_ADD;
 356       break;
 357
 358     case OPT_HASH:
 359       opt_hash = strdup(optarg);
 360       break;
 361
 362     case OPT_USERNAME:
 363       opt_username = optarg;
 364       break;
 365
 366     case OPT_PASSWD:
 367       opt_passwd = optarg;
 368       break;
 369
 370     case OPT_QUIET:
 371       close_stdio();
 372       break;
 373
 374     case OPT_UDP:
 375       opt_udp = true;
 376       break;
 377
 378     case OPT_BUFFER:
 379       opt_buffer = true;
 380       break;
 381
 382     case '?':
 383       /* getopt_long already printed an error message. */
 384       exit(1);
 385     default:
 386       abort();
 387     }
 388   }
 389
 390   if (opt_version) {
 391     version_command(PROGRAM_NAME);
 392     exit(EXIT_SUCCESS);
 393   }
 394
 395   if (opt_help) {
 396     help_command(PROGRAM_NAME, PROGRAM_DESCRIPTION, long_options, help_options);
 397     exit(EXIT_SUCCESS);
 398   }
 399 }