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