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