git.m6w6.name Git - m6w6/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 #define PROGRAM_NAME        "memcp"
  19 #define PROGRAM_DESCRIPTION "Copy a set of files to a memcached cluster."
  20 #define PROGRAM_VERSION     "1.1"
  21
  22 #include "common/options.hpp"
  23 #include "common/checks.hpp"
  24
  25 #include <cerrno>
  26 #include <climits>
  27 #include <cstdlib>
  28 #if HAVE_LIBGEN_H
  29 #  include <libgen.h>
  30 #endif
  31 #include <fstream>
  32 #include <sstream>
  33
  34 #ifndef PATH_MAX
  35 #  ifdef MAX_PATH
  36 #    define PATH_MAX MAX_PATH
  37 #  else
  38 #    define PATH_MAX 256
  39 #  endif
  40 #endif
  41
  42 struct memcp_file {
  43   enum class type {
  44     basename,
  45     relative,
  46     absolute
  47   } key;
  48   enum class mode {
  49     SET,
  50     ADD,
  51     REPLACE
  52   } op;
  53   char *path;
  54   uint32_t flags;
  55   time_t expire;
  56 };
  57
  58 static inline std::string stream2string(const std::istream &istream) {
  59   return dynamic_cast<std::ostringstream &&>(std::ostringstream{} << istream.rdbuf()).str();
  60 }
  61
  62 static memcached_return_t memcp(const client_options &opt, memcached_st &memc, const char *key,
  63                   const memcp_file &file) {
  64   std::ifstream fstream{};
  65   std::istream *istream = check_istream(opt, file.path, fstream);
  66
  67   if (!istream){
  68     return MEMCACHED_ERROR;
  69   }
  70
  71   const char *mode;
  72   memcached_return_t rc;
  73   auto data = stream2string(*istream);
  74   if (file.op == memcp_file::mode::REPLACE) {
  75     mode = "replace";
  76     rc = memcached_replace(&memc, key, strlen(key), data.c_str(), data.length(),
  77         file.expire, file.flags);
  78   } else if (file.op == memcp_file::mode::ADD) {
  79     mode = "add";
  80     rc = memcached_add(&memc, key, strlen(key), data.c_str(), data.length(),
  81         file.expire, file.flags);
  82   } else {
  83     mode = "set";
  84     rc = memcached_set(&memc, key, strlen(key), data.c_str(), data.length(),
  85         file.expire, file.flags);
  86   }
  87
  88   if (!memcached_success(rc)) {
  89     auto error = memcached_last_error(&memc)
  90                  ? memcached_last_error_message(&memc)
  91                  : memcached_strerror(&memc, rc);
  92     std::cerr << "Error occurred during memcached_" << mode <<"('" << key << "'): " << error << "\n";
  93   }
  94   return rc;
  95 }
  96
  97 static void add_file(std::vector<memcp_file> &files, const client_options &opt, char *file) {
  98   memcp_file::type type = memcp_file::type::basename;
  99   memcp_file::mode mode = memcp_file::mode::SET;
 100   uint32_t flags = 0;
 101   time_t expire = 0;
 102
 103   if (opt.isset("absolute")) {
 104     type = memcp_file::type::absolute;
 105   } else if (opt.isset("relative")) {
 106     type = memcp_file::type::relative;
 107   }
 108
 109   if (opt.isset("replace")) {
 110     mode = memcp_file::mode::REPLACE;
 111   } else if (opt.isset("add")) {
 112     mode = memcp_file::mode::ADD;
 113   }
 114
 115   if (auto flags_str = opt.argof("flags")) {
 116     flags = std::stoul(flags_str);
 117   }
 118   if (auto expire_str = opt.argof("expire")) {
 119     expire = std::stoul(expire_str);
 120   }
 121
 122   if (opt.isset("debug")) {
 123     auto mode_str = mode == memcp_file::mode::REPLACE ? "REPLACE" : mode == memcp_file::mode::ADD ? "ADD" : "SET";
 124     std::cerr << "Scheduling " << mode_str << " '" << file << "' (expire=" << expire << ", flags=" << flags << ").\n";
 125   }
 126
 127   files.emplace_back(memcp_file{type, mode, file, flags, expire});
 128 }
 129
 130 static bool path2key(const client_options &opt, memcp_file &file, char **path) {
 131   static char rpath[PATH_MAX + 1];
 132   if (file.key == memcp_file::type::absolute) {
 133     *path = realpath(file.path, rpath);
 134     if (!*path) {
 135       if (!opt.isset("quiet")) {
 136         perror(file.path);
 137       }
 138       return false;
 139     }
 140   } else if (file.key == memcp_file::type::relative) {
 141     *path = file.path;
 142   } else {
 143     *path = basename(file.path);
 144   }
 145   return true;
 146 }
 147
 148 int main(int argc, char *argv[]) {
 149   std::vector<memcp_file> files{};
 150   client_options opt{PROGRAM_NAME, PROGRAM_VERSION, PROGRAM_DESCRIPTION,
 151                      "file [file ...]"
 152                      "\n\t\t\t# NOTE: order of flags and positional"
 153                      "\n\t\t\t#       arguments matters on GNU systems)"};
 154
 155   opt.add(nullptr, '-', no_argument, "GNU argv extension")
 156       .parse = [&files](client_options &opt_, client_options::extended_option &ext) {
 157     add_file(files, opt_, ext.arg);
 158     return true;
 159   };
 160
 161   for (const auto &def : opt.defaults) {
 162     opt.add(def);
 163   }
 164
 165   opt.add("set", 'S', no_argument, "Perform SET operations.")
 166       .parse = [](client_options &opt_, client_options::extended_option &) {
 167     opt_.unset("add");
 168     opt_.unset("replace");
 169     return true;
 170   };
 171   opt.add("add", 'A', no_argument, "Perform ADD operations.")
 172       .parse = [](client_options &opt_, client_options::extended_option &) {
 173     opt_.unset("set");
 174     opt_.unset("replace");
 175     return true;
 176   };
 177   opt.add("replace", 'R', no_argument, "Perform REPLACE operations.")
 178       .parse = [](client_options &opt_, client_options::extended_option &) {
 179     opt_.unset("set");
 180     opt_.unset("add");
 181     return true;
 182   };
 183
 184   opt.add("udp", 'U', no_argument, "Use UDP.")
 185       .apply = [](const client_options &opt_, const client_options::extended_option &ext, memcached_st *memc) {
 186     if (MEMCACHED_SUCCESS != memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_USE_UDP, ext.set)) {
 187       if (!opt_.isset("quiet")) {
 188         std::cerr << memcached_last_error_message(memc) << "\n";
 189       }
 190       return false;
 191     }
 192     return true;
 193   };
 194   opt.add("flags", 'F', required_argument, "Set key flags, too.");
 195   opt.add("expire", 'e', required_argument, "Set expire time, too.");
 196
 197   opt.add("basename", '.', no_argument, "Use basename of path as key (default).");
 198   opt.add("relative", '+', no_argument, "Use relative path (as passed), instead of basename only.");
 199   opt.add("absolute", '/', no_argument, "Use absolute path (real path), instead of basename only.");
 200
 201   // defaults
 202   opt.set("set");
 203   opt.set("basename");
 204
 205   char **argp = nullptr;
 206   if (!opt.parse(argc, argv, &argp)) {
 207     exit(EXIT_FAILURE);
 208   }
 209
 210   memcached_st memc;
 211   if (!check_memcached(opt, memc)) {
 212     exit(EXIT_FAILURE);
 213   }
 214
 215   if (!opt.apply(&memc)) {
 216     exit(EXIT_FAILURE);
 217   }
 218
 219   if (files.empty()) {
 220     if (!check_argp(opt, argp, "No file(s) provided.")) {
 221       memcached_free(&memc);
 222       exit(EXIT_FAILURE);
 223     }
 224     for (auto arg = argp; *arg; ++arg) {
 225       add_file(files, opt, *arg);
 226     }
 227   }
 228
 229   auto exit_code = EXIT_SUCCESS;
 230   for (auto &file : files) {
 231     char *path = nullptr;
 232     if (!path2key(opt, file, &path)) {
 233       exit_code = EXIT_FAILURE;
 234       continue;
 235     }
 236
 237     auto rc = memcp(opt, memc, path, file);
 238     if (memcached_success(rc)) {
 239       if (opt.isset("verbose")) {
 240         std::cout << path << "\n";
 241       }
 242     } else {
 243       exit_code = EXIT_FAILURE;
 244     }
 245   }
 246
 247   if (!check_buffering(opt, memc)) {
 248     exit_code = EXIT_FAILURE;
 249   }
 250
 251   memcached_free(&memc);
 252   exit(exit_code);
 253 }