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