Skip to content

Commit

Permalink
Add rename-command config (#272)
Browse files Browse the repository at this point in the history
* Add rename-command config
* Add rename-command config directive testcase
  • Loading branch information
karelrooted authored May 26, 2021
1 parent bc6446e commit fbd7372
Show file tree
Hide file tree
Showing 10 changed files with 153 additions and 97 deletions.
16 changes: 16 additions & 0 deletions kvrocks.conf
Original file line number Diff line number Diff line change
Expand Up @@ -270,6 +270,22 @@ compaction-checker-range 0-7
# e.g. bgsave-cron 0 3 * * * 0 4 * * *
# would bgsave the db at 3am and 4am everyday

# Command renaming.
#
# It is possible to change the name of dangerous commands in a shared
# environment. For instance the KEYS command may be renamed into something
# hard to guess so that it will still be available for internal-use tools
# but not available for general clients.
#
# Example:
#
# rename-command KEYS b840fc02d524045429941cc15f59e41cb7be6c52
#
# It is also possible to completely kill a command by renaming it into
# an empty string:
#
# rename-command KEYS ""

################################ ROCKSDB #####################################

# Specify the capacity of metadata column family block cache. Larger block cache
Expand Down
27 changes: 24 additions & 3 deletions src/config.cc
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,7 @@ Config::Config() {
{"profiling-sample-commands", false, new StringField(&profiling_sample_commands_, "")},
{"slowlog-max-len", false, new IntField(&slowlog_max_len, 128, 0, INT_MAX)},
{"purge-backup-on-fullsync", false, new YesNoField(&purge_backup_on_fullsync, false)},
{"rename-command", true, new StringField(&rename_command_, "")},
/* rocksdb options */
{"rocksdb.compression", false, new EnumField(&RocksDB.compression, compression_type_enum, 0)},
{"rocksdb.block_size", true, new IntField(&RocksDB.block_size, 4096, 0, INT_MAX)},
Expand Down Expand Up @@ -175,6 +176,27 @@ void Config::initFieldValidator() {
compaction_checker_range.Stop = stop;
return Status::OK();
}},
{"rename-command", [](const std::string &k, const std::string &v) -> Status {
std::vector<std::string> args;
Util::Split(v, " \t", &args);
if (args.size() != 2) {
return Status(Status::NotOK, "Invalid rename-command format");
}
auto commands = Redis::GetCommands();
auto cmd_iter = commands->find(Util::ToLower(args[0]));
if (cmd_iter == commands->end()) {
return Status(Status::NotOK, "No such command in rename-command");
}
if (args[1] != "\"\"") {
auto new_command_name = Util::ToLower(args[1]);
if (commands->find(new_command_name) != commands->end()) {
return Status(Status::NotOK, "Target command name already exists");
}
(*commands)[new_command_name] = cmd_iter->second;
}
commands->erase(cmd_iter);
return Status::OK();
}},
};
for (const auto& iter : validators) {
auto field_iter = fields_.find(iter.first);
Expand Down Expand Up @@ -230,8 +252,6 @@ void Config::initFieldCallback() {
return Status::OK();
}},
{"profiling-sample-commands", [this](Server* srv, const std::string &k, const std::string& v)->Status {
if (!srv) return Status::OK();

std::vector<std::string> cmds;
Util::Split(v, ",", &cmds);
profiling_sample_all_commands = false;
Expand All @@ -240,7 +260,8 @@ void Config::initFieldCallback() {
if (cmd == "*") {
profiling_sample_all_commands = true;
profiling_sample_commands.clear();
} else if (srv->IsCommandExists(cmd)) {
} else if (Redis::IsCommandExists(cmd)) {
// profiling_sample_commands use command's original name, regardless of rename-command directive
profiling_sample_commands.insert(cmd);
}
}
Expand Down
1 change: 1 addition & 0 deletions src/config.h
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,7 @@ struct Config{
std::string compaction_checker_range_;
std::string profiling_sample_commands_;
std::map<std::string, ConfigField*> fields_;
std::string rename_command_;

void initFieldValidator();
void initFieldCallback();
Expand Down
2 changes: 2 additions & 0 deletions src/main.cc
Original file line number Diff line number Diff line change
Expand Up @@ -275,6 +275,8 @@ int main(int argc, char* argv[]) {
auto opts = parseCommandLineOptions(argc, argv);
if (opts.show_usage) usage(argv[0]);

Redis::PopulateCommands();

Config config;
Status s = config.Load(opts.conf_file);
if (!s.IsOK()) {
Expand Down
92 changes: 85 additions & 7 deletions src/redis_cmd.cc
Original file line number Diff line number Diff line change
Expand Up @@ -3647,7 +3647,7 @@ class CommandCommand : public Commander {
public:
Status Execute(Server *svr, Connection *conn, std::string *output) override {
if (args_.size() == 1) {
svr->GetAllCommandsInfo(output);
GetAllCommandsInfo(output);
} else {
std::string sub_command = Util::ToLower(args_[1]);
if ((sub_command == "count" && args_.size() != 2) ||
Expand All @@ -3659,10 +3659,10 @@ class CommandCommand : public Commander {
if (sub_command == "count") {
*output = Redis::Integer(GetCommandNum());
} else if (sub_command == "info") {
svr->GetCommandsInfo(output, std::vector<std::string>(args_.begin() + 2, args_.end()));
GetCommandsInfo(output, std::vector<std::string>(args_.begin() + 2, args_.end()));
} else if (sub_command == "getkeys") {
std::vector<int> keys_indexes;
auto s = svr->GetKeysFromCommand(args_[2], args_.size() - 2, &keys_indexes);
auto s = GetKeysFromCommand(args_[2], args_.size() - 2, &keys_indexes);
if (!s.IsOK()) return s;
if (keys_indexes.size() == 0) {
*output = Redis::Error("Invalid arguments specified for command");
Expand Down Expand Up @@ -4230,12 +4230,90 @@ CommandAttributes redisCommandTable[] = {
ADD_CMD("_db_name", 1, false, 0, 0, 0, CommandDBName),
};

CommandAttributes * GetCommandTable() {
return redisCommandTable;
}
// Command table after rename-command directive
std::map<std::string, CommandAttributes *> commands;
// Original Command table before rename-command directive
std::map<std::string, CommandAttributes *> original_commands;

int GetCommandNum() {
return sizeof(redisCommandTable)/sizeof(struct CommandAttributes);
return sizeof(redisCommandTable) / sizeof(struct CommandAttributes);
}

std::map<std::string, CommandAttributes *> *GetCommands() {
return &commands;
}

std::map<std::string, CommandAttributes *> *GetOriginalCommands() {
return &original_commands;
}

void PopulateCommands() {
for (int i = 0; i < GetCommandNum(); i++) {
original_commands[redisCommandTable[i].name] = &redisCommandTable[i];
}
commands = original_commands;
}

std::string GetCommandInfo(const CommandAttributes *command_attributes) {
std::string command, command_flags;
command.append(Redis::MultiLen(6));
command.append(Redis::BulkString(command_attributes->name));
command.append(Redis::Integer(command_attributes->arity));
command_flags.append(Redis::MultiLen(1));
command_flags.append(Redis::BulkString(command_attributes->is_write ? "write" : "readonly"));
command.append(command_flags);
command.append(Redis::Integer(command_attributes->first_key));
command.append(Redis::Integer(command_attributes->last_key));
command.append(Redis::Integer(command_attributes->key_step));
return command;
}

void GetAllCommandsInfo(std::string *info) {
info->append(Redis::MultiLen(original_commands.size()));
for (const auto &iter : original_commands) {
auto command_attribute = iter.second;
auto command_info = GetCommandInfo(command_attribute);
info->append(command_info);
}
}

void GetCommandsInfo(std::string *info, const std::vector<std::string> &cmd_names) {
info->append(Redis::MultiLen(cmd_names.size()));
for (const auto &cmd_name : cmd_names) {
auto cmd_iter = original_commands.find(Util::ToLower(cmd_name));
if (cmd_iter == original_commands.end()) {
info->append(Redis::NilString());
} else {
auto command_attribute = cmd_iter->second;
auto command_info = GetCommandInfo(command_attribute);
info->append(command_info);
}
}
}

Status GetKeysFromCommand(const std::string &cmd_name, int argc, std::vector<int> *keys_indexes) {
auto cmd_iter = original_commands.find(Util::ToLower(cmd_name));
if (cmd_iter == original_commands.end()) {
return Status(Status::RedisUnknownCmd, "Invalid command specified");
}
auto command_attribute = cmd_iter->second;
if (command_attribute->first_key == 0) {
return Status(Status::NotOK, "The command has no key arguments");
}
if ((command_attribute->arity > 0 && command_attribute->arity != argc) || argc < -command_attribute->arity) {
return Status(Status::NotOK, "Invalid number of arguments specified for command");
}
auto last = command_attribute->last_key;
if (last < 0) last = argc + last;

for (int j = command_attribute->first_key; j <= last; j += command_attribute->key_step) {
keys_indexes->emplace_back(j);
}
return Status::OK();
}

bool IsCommandExists(const std::string &name) {
return original_commands.find(Util::ToLower(name)) != original_commands.end();
}

} // namespace Redis
9 changes: 8 additions & 1 deletion src/redis_cmd.h
Original file line number Diff line number Diff line change
Expand Up @@ -58,5 +58,12 @@ struct CommandAttributes {
};

int GetCommandNum();
CommandAttributes *GetCommandTable();
std::map<std::string, CommandAttributes *> *GetCommands();
std::map<std::string, CommandAttributes *> *GetOriginalCommands();
void PopulateCommands();
void GetAllCommandsInfo(std::string *info);
void GetCommandsInfo(std::string *info, const std::vector<std::string> &cmd_names);
std::string GetCommandInfo(const CommandAttributes *command_attributes);
Status GetKeysFromCommand(const std::string &name, int argc, std::vector<int> *keys_indexes);
bool IsCommandExists(const std::string &name);
} // namespace Redis
85 changes: 6 additions & 79 deletions src/server.cc
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,9 @@ std::atomic<int>Server::unix_time_ = {0};

Server::Server(Engine::Storage *storage, Config *config) :
storage_(storage), config_(config) {
populateCommands();
// init commands stats here to prevent concurrent insert, and cause core
for (const auto &iter : commands_) {
auto commands = Redis::GetOriginalCommands();
for (const auto &iter : *commands) {
stats_.commands_stats[iter.first].calls = 0;
stats_.commands_stats[iter.first].latency = 0;
}
Expand All @@ -48,9 +48,6 @@ Server::~Server() {
for (const auto &iter : conn_ctxs_) {
delete iter.first;
}
for (const auto &iter : commands_) {
delete iter.second;
}
}

Status Server::Start() {
Expand Down Expand Up @@ -1130,10 +1127,11 @@ ReplState Server::GetReplicationState() {
}

Status Server::LookupAndCreateCommand(const std::string &cmd_name,
std::unique_ptr<Redis::Commander> *cmd) {
std::unique_ptr<Redis::Commander> *cmd) {
auto commands = Redis::GetCommands();
if (cmd_name.empty()) return Status(Status::RedisUnknownCmd);
auto cmd_iter = commands_.find(Util::ToLower(cmd_name));
if (cmd_iter == commands_.end()) {
auto cmd_iter = commands->find(Util::ToLower(cmd_name));
if (cmd_iter == commands->end()) {
return Status(Status::RedisUnknownCmd);
}
auto redisCmd = cmd_iter->second;
Expand All @@ -1142,74 +1140,3 @@ Status Server::LookupAndCreateCommand(const std::string &cmd_name,
return Status::OK();
}

void Server::populateCommands() {
Redis::CommandAttributes* commandTable = Redis::GetCommandTable();
for (int i = 0; i < Redis::GetCommandNum(); i++) {
auto commandAttributes = new(Redis::CommandAttributes);
*commandAttributes = commandTable[i];
commands_[commandTable[i].name] = commandAttributes;
}
}

bool Server::IsCommandExists(const std::string &name) {
return commands_.find(name) != commands_.end();
}

std::string Server::GetCommandInfo(const Redis::CommandAttributes* command_attributes) {
std::string command, command_flags;
command.append(Redis::MultiLen(6));
command.append(Redis::BulkString(command_attributes->name));
command.append(Redis::Integer(command_attributes->arity));
command_flags.append(Redis::MultiLen(1));
command_flags.append(Redis::BulkString(command_attributes->is_write ? "write" : "readonly"));
command.append(command_flags);
command.append(Redis::Integer(command_attributes->first_key));
command.append(Redis::Integer(command_attributes->last_key));
command.append(Redis::Integer(command_attributes->key_step));
return command;
}

void Server::GetAllCommandsInfo(std::string *info) {
info->append(Redis::MultiLen(commands_.size()));
for (const auto &iter : commands_) {
auto command_attribute = iter.second;
auto command_info = GetCommandInfo(command_attribute);
info->append(command_info);
}
}

void Server::GetCommandsInfo(std::string *info, const std::vector<std::string> &cmd_names) {
info->append(Redis::MultiLen(cmd_names.size()));
for (const auto &cmd_name : cmd_names) {
auto cmd_iter = commands_.find(Util::ToLower(cmd_name));
if (cmd_iter == commands_.end()) {
info->append(Redis::NilString());
} else {
auto command_attribute = cmd_iter->second;
auto command_info = GetCommandInfo(command_attribute);
info->append(command_info);
}
}
}

Status Server::GetKeysFromCommand(const std::string &cmd_name, int argc, std::vector<int> *keys_indexes) {
auto cmd_iter = commands_.find(Util::ToLower(cmd_name));
if (cmd_iter == commands_.end()) {
return Status(Status::RedisUnknownCmd, "Invalid command specified");
}
auto command_attribute = cmd_iter->second;
if (command_attribute->first_key == 0) {
return Status(Status::NotOK, "The command has no key arguments");
}
if ((command_attribute->arity > 0 && command_attribute->arity != argc) || argc < -command_attribute->arity) {
return Status(Status::NotOK, "Invalid number of arguments specified for command");
}
auto last = command_attribute->last_key;
if (last < 0) last = argc + last;

for (int j = command_attribute->first_key; j <= last; j += command_attribute->key_step) {
keys_indexes->emplace_back(j);
}
return Status::OK();
}

8 changes: 1 addition & 7 deletions src/server.h
Original file line number Diff line number Diff line change
Expand Up @@ -50,13 +50,8 @@ class Server {
bool IsStopped() { return stop_; }
bool IsLoading() { return is_loading_; }
Config *GetConfig() { return config_; }
void populateCommands();
bool IsCommandExists(const std::string &name);
Status LookupAndCreateCommand(const std::string &cmd_name, std::unique_ptr<Redis::Commander> *cmd);
void GetAllCommandsInfo(std::string *info);
void GetCommandsInfo(std::string *info, const std::vector<std::string> &cmd_names);
std::string GetCommandInfo(const Redis::CommandAttributes* command_attributes);
Status GetKeysFromCommand(const std::string &name, int argc, std::vector<int> *keys_indexes);


Status AddMaster(std::string host, uint32_t port);
Status RemoveMaster();
Expand Down Expand Up @@ -141,7 +136,6 @@ class Server {
Config *config_ = nullptr;
std::string last_random_key_cursor_;
std::mutex last_random_key_cursor_mu_;
std::map<std::string, Redis::CommandAttributes*> commands_;

// client counters
std::atomic<uint64_t> client_id_{1};
Expand Down
1 change: 1 addition & 0 deletions tests/tcl/tests/test_helper.tcl
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ set ::all_tests {
unit/limits
unit/geo
unit/command
unit/config
integration/replication
}

Expand Down
9 changes: 9 additions & 0 deletions tests/tcl/tests/unit/config.tcl
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
start_server {tags {"config"} overrides {rename-command "KEYS KEYSNEW"}} {
test {Rename one command} {
catch {r KEYS *} e
assert_error "*invalid command name*" $e
assert_equal "" [r KEYSNEW *]
}
}


0 comments on commit fbd7372

Please sign in to comment.