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