git.m6w6.name Git - awesomized/libmemcached/blob - clients/utilities.cc

   1 /* LibMemcached
   2  * Copyright (C) 2006-2009 Brian Aker
   3  * All rights reserved.
   4  *
   5  * Use and distribution licensed under the BSD license.  See
   6  * the COPYING file in the parent directory for full text.
   7  *
   8  * Summary:
   9  *
  10  */
  11 #include "config.h"
  12
  13 #include <clients/utilities.h>
  14 #include <cstdio>
  15 #include <cstring>
  16 #include <ctype.h>
  17
  18
  19 long int timedif(struct timeval a, struct timeval b)
  20 {
  21   long us, s;
  22
  23   us = (int)(a.tv_usec - b.tv_usec);
  24   us /= 1000;
  25   s = (int)(a.tv_sec - b.tv_sec);
  26   s *= 1000;
  27   return s + us;
  28 }
  29
  30 void version_command(const char *command_name)
  31 {
  32   printf("%s v%u.%u\n", command_name, 1U, 0U);
  33   exit(0);
  34 }
  35
  36 static const char *lookup_help(memcached_options option)
  37 {
  38   switch (option)
  39   {
  40   case OPT_SERVERS: return("List which servers you wish to connect to.");
  41   case OPT_VERSION: return("Display the version of the application and then exit.");
  42   case OPT_HELP: return("Display this message and then exit.");
  43   case OPT_VERBOSE: return("Give more details on the progression of the application.");
  44   case OPT_DEBUG: return("Provide output only useful for debugging.");
  45   case OPT_FLAG: return("Provide flag information for storage operation.");
  46   case OPT_EXPIRE: return("Set the expire option for the object.");
  47   case OPT_SET: return("Use set command with memcached when storing.");
  48   case OPT_REPLACE: return("Use replace command with memcached when storing.");
  49   case OPT_ADD: return("Use add command with memcached when storing.");
  50   case OPT_SLAP_EXECUTE_NUMBER: return("Number of times to execute the given test.");
  51   case OPT_SLAP_INITIAL_LOAD: return("Number of key pairs to load before executing tests.");
  52   case OPT_SLAP_TEST: return("Test to run (currently \"get\" or \"set\").");
  53   case OPT_SLAP_CONCURRENCY: return("Number of users to simulate with load.");
  54   case OPT_SLAP_NON_BLOCK: return("Set TCP up to use non-blocking IO.");
  55   case OPT_SLAP_TCP_NODELAY: return("Set TCP socket up to use nodelay.");
  56   case OPT_FLUSH: return("Flush servers before running tests.");
  57   case OPT_HASH: return("Select hash type.");
  58   case OPT_BINARY: return("Switch to binary protocol.");
  59   case OPT_ANALYZE: return("Analyze the provided servers.");
  60   case OPT_UDP: return("Use UDP protocol when communicating with server.");
  61   case OPT_USERNAME: return "Username to use for SASL authentication";
  62   case OPT_PASSWD: return "Password to use for SASL authentication";
  63   case OPT_FILE: return "Path to file in which to save result";
  64   case OPT_STAT_ARGS: return "Argument for statistics";
  65   default: WATCHPOINT_ASSERT(0);
  66   };
  67
  68   WATCHPOINT_ASSERT(0);
  69   return "forgot to document this function :)";
  70 }
  71
  72 void help_command(const char *command_name, const char *description,
  73                   const struct option *long_options,
  74                   memcached_programs_help_st *options)
  75 {
  76   unsigned int x;
  77   (void)options;
  78
  79   printf("%s v%u.%u\n\n", command_name, 1U, 0U);
  80   printf("\t%s\n\n", description);
  81   printf("Current options. A '=' means the option takes a value.\n\n");
  82
  83   for (x= 0; long_options[x].name; x++)
  84   {
  85     const char *help_message;
  86
  87     printf("\t --%s%c\n", long_options[x].name,
  88            long_options[x].has_arg ? '=' : ' ');
  89     if ((help_message= lookup_help(memcached_options(long_options[x].val))))
  90       printf("\t\t%s\n", help_message);
  91   }
  92
  93   printf("\n");
  94   exit(0);
  95 }
  96
  97 void process_hash_option(memcached_st *memc, char *opt_hash)
  98 {
  99   uint64_t set;
 100   memcached_return_t rc;
 101
 102   if (opt_hash == NULL)
 103     return;
 104
 105   set= MEMCACHED_HASH_DEFAULT; /* Just here to solve warning */
 106   if (!strcasecmp(opt_hash, "CRC"))
 107     set= MEMCACHED_HASH_CRC;
 108   else if (!strcasecmp(opt_hash, "FNV1_64"))
 109     set= MEMCACHED_HASH_FNV1_64;
 110   else if (!strcasecmp(opt_hash, "FNV1A_64"))
 111     set= MEMCACHED_HASH_FNV1A_64;
 112   else if (!strcasecmp(opt_hash, "FNV1_32"))
 113     set= MEMCACHED_HASH_FNV1_32;
 114   else if (!strcasecmp(opt_hash, "FNV1A_32"))
 115     set= MEMCACHED_HASH_FNV1A_32;
 116   else
 117   {
 118     fprintf(stderr, "hash: type not recognized %s\n", opt_hash);
 119     exit(1);
 120   }
 121
 122   rc= memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_HASH, set);
 123   if (rc != MEMCACHED_SUCCESS)
 124   {
 125     fprintf(stderr, "hash: memcache error %s\n", memcached_strerror(memc, rc));
 126     exit(1);
 127   }
 128 }
 129
 130
 131 static char *username;
 132 static char *passwd;
 133
 134 #if defined(LIBMEMCACHED_WITH_SASL_SUPPORT) && LIBMEMCACHED_WITH_SASL_SUPPORT
 135
 136 static int get_username(void *context, int id, const char **result, unsigned int *len)
 137 {
 138   (void)context;
 139   if (!result || (id != SASL_CB_USER && id != SASL_CB_AUTHNAME))
 140     return SASL_BADPARAM;
 141
 142   *result= username;
 143   if (len)
 144   {
 145      *len= (username == NULL) ? 0 : (unsigned int)strlen(username);
 146   }
 147
 148   return SASL_OK;
 149 }
 150
 151 static int get_password(sasl_conn_t *conn, void *context, int id,
 152                         sasl_secret_t **psecret)
 153 {
 154   (void)context;
 155   static sasl_secret_t* ptr;
 156
 157   if (!conn || ! psecret || id != SASL_CB_PASS)
 158     return SASL_BADPARAM;
 159
 160   if (passwd == NULL)
 161   {
 162      *psecret= NULL;
 163      return SASL_OK;
 164   }
 165
 166   size_t len= strlen(passwd);
 167   ptr= (sasl_secret_t *)malloc(sizeof(sasl_secret_t) + len +1);
 168   if (not ptr)
 169     return SASL_NOMEM;
 170
 171   ptr->len= len;
 172   memcpy(ptr->data, passwd, len);
 173   ptr->data[len]= 0;
 174
 175   *psecret= ptr;
 176   return SASL_OK;
 177 }
 178
 179 typedef int (*local_sasl_fn)(void);
 180
 181 /* callbacks we support */
 182 static sasl_callback_t sasl_callbacks[] = {
 183   { SASL_CB_USER, (local_sasl_fn)get_username, NULL },
 184   { SASL_CB_AUTHNAME, (local_sasl_fn)get_username, NULL },
 185   { SASL_CB_PASS, (local_sasl_fn)get_password, NULL },
 186   { SASL_CB_LIST_END, NULL, NULL }
 187 };
 188
 189 #endif
 190
 191 bool initialize_sasl(memcached_st *memc, char *user, char *password)
 192 {
 193   if (LIBMEMCACHED_WITH_SASL_SUPPORT == 0)
 194   {
 195     return false;
 196   }
 197
 198   if (memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_BINARY_PROTOCOL, (uint64_t)true) == 0)
 199   {
 200     return false;
 201   }
 202
 203   if (user != NULL && password != NULL)
 204   {
 205     username= user;
 206     passwd= password;
 207
 208 #if defined(LIBMEMCACHED_WITH_SASL_SUPPORT) && LIBMEMCACHED_WITH_SASL_SUPPORT
 209     if (sasl_client_init(NULL) != SASL_OK)
 210     {
 211       fprintf(stderr, "Failed to initialize sasl library!\n");
 212       return false;
 213     }
 214     memcached_set_sasl_callbacks(memc, sasl_callbacks);
 215 #else
 216     (void)memc;
 217 #endif
 218   }
 219
 220   return true;
 221 }
 222
 223 void shutdown_sasl(void)
 224 {
 225   if (username or passwd)
 226   {
 227 #if defined(LIBMEMCACHED_WITH_SASL_SUPPORT) && LIBMEMCACHED_WITH_SASL_SUPPORT
 228     sasl_done();
 229 #endif
 230   }
 231 }
 232
 233 void initialize_sockets(void)
 234 {
 235   /* Define the function for all platforms to avoid #ifdefs in each program */
 236 #if defined(WIN32) && WIN32
 237   WSADATA wsaData;
 238   if (WSAStartup(MAKEWORD(2,0), &wsaData) != 0)
 239   {
 240     fprintf(stderr, "Socket Initialization Error. Program aborted\n");
 241     exit(EXIT_FAILURE);
 242   }
 243 #endif
 244 }