Skip to content

Commit

Permalink
feat(server): add eval_ro and evalsha_ro
Browse files Browse the repository at this point in the history
  • Loading branch information
@andydunstall
andydunstall committed Nov 9, 2024
1 parent c756832 commit 6bfdacd
Show file tree
Hide file tree
Showing 5 changed files with 76 additions and 13 deletions.
1 change: 1 addition & 0 deletions src/server/conn_context.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ struct ConnectionState {
size_t UsedMemory() const;

absl::flat_hash_set<LockTag> lock_tags; // declared tags
bool read_only = false;

size_t async_cmds_heap_mem = 0; // bytes used by async_cmds
size_t async_cmds_heap_limit = 0; // max bytes allowed for async_cmds
Expand Down
37 changes: 30 additions & 7 deletions src/server/main_service.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1225,6 +1225,10 @@ std::optional<ErrorReply> Service::VerifyCommandState(const CommandId* cid, CmdA
return err.value();
}
}

if (dfly_cntx.conn_state.script_info->read_only && is_write_cmd) {
return ErrorReply{"Write commands are not allowed from read-only scripts"};
}
}

return VerifyConnectionAclStatus(cid, &dfly_cntx, "has no ACL permissions", tail_args);
Expand Down Expand Up @@ -1865,7 +1869,7 @@ void Service::CallFromScript(ConnectionContext* cntx, Interpreter::CallArgs& ca)
}

void Service::Eval(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder,
ConnectionContext* cntx) {
ConnectionContext* cntx, bool read_only) {
string_view body = ArgS(args, 0);

auto* rb = static_cast<RedisReplyBuilder*>(builder);
Expand All @@ -1880,19 +1884,29 @@ void Service::Eval(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder,

string sha{std::move(res.value())};

CallSHA(args, sha, interpreter, builder, cntx);
CallSHA(args, sha, interpreter, builder, cntx, read_only);
}

void Service::EvalRo(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder,
ConnectionContext* cntx) {
Eval(args, tx, builder, cntx, true);
}

void Service::EvalSha(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder,
ConnectionContext* cntx) {
ConnectionContext* cntx, bool read_only) {
string sha = absl::AsciiStrToLower(ArgS(args, 0));

BorrowedInterpreter interpreter{cntx->transaction, cntx};
CallSHA(args, sha, interpreter, builder, cntx);
CallSHA(args, sha, interpreter, builder, cntx, read_only);
}

void Service::EvalShaRo(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder,
ConnectionContext* cntx) {
EvalSha(args, tx, builder, cntx, true);
}

void Service::CallSHA(CmdArgList args, string_view sha, Interpreter* interpreter,
SinkReplyBuilder* builder, ConnectionContext* cntx) {
SinkReplyBuilder* builder, ConnectionContext* cntx, bool read_only) {
uint32_t num_keys;
CHECK(absl::SimpleAtoi(ArgS(args, 1), &num_keys)); // we already validated this

Expand All @@ -1902,7 +1916,7 @@ void Service::CallSHA(CmdArgList args, string_view sha, Interpreter* interpreter
ev_args.args = args.subspan(2 + num_keys);

uint64_t start = absl::GetCurrentTimeNanos();
EvalInternal(args, ev_args, interpreter, builder, cntx);
EvalInternal(args, ev_args, interpreter, builder, cntx, read_only);

uint64_t end = absl::GetCurrentTimeNanos();
ServerState::tlocal()->RecordCallLatency(sha, (end - start) / 1000);
Expand Down Expand Up @@ -1988,7 +2002,7 @@ static bool CanRunSingleShardMulti(optional<ShardId> sid, const ScriptMgr::Scrip
}

void Service::EvalInternal(CmdArgList args, const EvalArgs& eval_args, Interpreter* interpreter,
SinkReplyBuilder* builder, ConnectionContext* cntx) {
SinkReplyBuilder* builder, ConnectionContext* cntx, bool read_only) {
DCHECK(!eval_args.sha.empty());

// Sanitizing the input to avoid code injection.
Expand All @@ -2010,6 +2024,7 @@ void Service::EvalInternal(CmdArgList args, const EvalArgs& eval_args, Interpret
auto& sinfo = cntx->conn_state.script_info;
sinfo = make_unique<ConnectionState::ScriptInfo>();
sinfo->lock_tags.reserve(eval_args.keys.size());
sinfo->read_only = read_only;

optional<ShardId> sid;

Expand Down Expand Up @@ -2683,7 +2698,9 @@ constexpr uint32_t kWatch = FAST | TRANSACTION;
constexpr uint32_t kUnwatch = FAST | TRANSACTION;
constexpr uint32_t kDiscard = FAST | TRANSACTION;
constexpr uint32_t kEval = SLOW | SCRIPTING;
constexpr uint32_t kEvalRo = SLOW | SCRIPTING;
constexpr uint32_t kEvalSha = SLOW | SCRIPTING;
constexpr uint32_t kEvalShaRo = SLOW | SCRIPTING;
constexpr uint32_t kExec = SLOW | TRANSACTION;
constexpr uint32_t kPublish = PUBSUB | FAST;
constexpr uint32_t kSubscribe = PUBSUB | SLOW;
Expand All @@ -2708,9 +2725,15 @@ void Service::Register(CommandRegistry* registry) {
<< CI{"EVAL", CO::NOSCRIPT | CO::VARIADIC_KEYS, -3, 3, 3, acl::kEval}
.MFUNC(Eval)
.SetValidator(&EvalValidator)
<< CI{"EVAL_RO", CO::NOSCRIPT | CO::VARIADIC_KEYS, -3, 3, 3, acl::kEvalRo}
.MFUNC(EvalRo)
.SetValidator(&EvalValidator)
<< CI{"EVALSHA", CO::NOSCRIPT | CO::VARIADIC_KEYS, -3, 3, 3, acl::kEvalSha}
.MFUNC(EvalSha)
.SetValidator(&EvalValidator)
<< CI{"EVALSHA_RO", CO::NOSCRIPT | CO::VARIADIC_KEYS, -3, 3, 3, acl::kEvalShaRo}
.MFUNC(EvalShaRo)
.SetValidator(&EvalValidator)
<< CI{"EXEC", CO::LOADING | CO::NOSCRIPT, 1, 0, 0, acl::kExec}.MFUNC(Exec)
<< CI{"PUBLISH", CO::LOADING | CO::FAST, 3, 0, 0, acl::kPublish}.MFUNC(Publish)
<< CI{"SUBSCRIBE", CO::NOSCRIPT | CO::LOADING, -2, 0, 0, acl::kSubscribe}.MFUNC(Subscribe)
Expand Down
14 changes: 9 additions & 5 deletions src/server/main_service.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -134,9 +134,13 @@ class Service : public facade::ServiceInterface {

void Discard(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder,
ConnectionContext* cntx);
void Eval(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder, ConnectionContext* cntx);
void EvalSha(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder,
ConnectionContext* cntx);
void Eval(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder, ConnectionContext* cntx,
bool read_only = false);
void EvalRo(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder, ConnectionContext* cntx);
void EvalSha(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder, ConnectionContext* cntx,
bool read_only = false);
void EvalShaRo(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder,
ConnectionContext* cntx);
void Exec(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder, ConnectionContext* cntx);
void Publish(CmdArgList args, Transaction* tx, SinkReplyBuilder* builder,
ConnectionContext* cntx);
Expand Down Expand Up @@ -170,9 +174,9 @@ class Service : public facade::ServiceInterface {
const ConnectionContext& dfly_cntx);

void EvalInternal(CmdArgList args, const EvalArgs& eval_args, Interpreter* interpreter,
SinkReplyBuilder* builder, ConnectionContext* cntx);
SinkReplyBuilder* builder, ConnectionContext* cntx, bool read_only);
void CallSHA(CmdArgList args, std::string_view sha, Interpreter* interpreter,
SinkReplyBuilder* builder, ConnectionContext* cntx);
SinkReplyBuilder* builder, ConnectionContext* cntx, bool read_only);

// Return optional payload - first received error that occured when executing commands.
std::optional<facade::CapturingReplyBuilder::Payload> FlushEvalAsyncCmds(ConnectionContext* cntx,
Expand Down
34 changes: 34 additions & 0 deletions src/server/multi_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1162,4 +1162,38 @@ TEST_F(MultiTest, MultiTypes) {
RespArray(ElementsAre("none", "none", "none", "none", "none", "none")));
}

TEST_F(MultiTest, EvalRo) {
RespExpr resp;

resp = Run({"set", "foo", "bar"});
EXPECT_THAT(resp, "OK");

resp = Run({"eval_ro", "return redis.call('get', KEYS[1])", "1", "foo"});
EXPECT_THAT(resp, "bar");

resp = Run({"eval_ro", "return redis.call('set', KEYS[1], 'car')", "1", "foo"});
EXPECT_THAT(resp, ErrArg("Write commands are not allowed from read-only scripts"));
}

TEST_F(MultiTest, EvalShaRo) {
RespExpr resp;

const char* read_script = "return redis.call('get', KEYS[1]);";
const char* write_script = "return redis.call('set', KEYS[1], 'car');";

auto sha_resp = Run({"script", "load", read_script});
auto read_sha = facade::ToSV(sha_resp.GetBuf());
sha_resp = Run({"script", "load", write_script});
auto write_sha = facade::ToSV(sha_resp.GetBuf());

resp = Run({"set", "foo", "bar"});
EXPECT_THAT(resp, "OK");

resp = Run({"evalsha_ro", read_sha, "1", "foo"});
EXPECT_THAT(resp, "bar");

resp = Run({"evalsha_ro", write_sha, "1", "foo"});
EXPECT_THAT(resp, ErrArg("Write commands are not allowed from read-only scripts"));
}

} // namespace dfly
3 changes: 2 additions & 1 deletion src/server/transaction.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,8 @@ void Transaction::Shutdown() {
Transaction::Transaction(const CommandId* cid) : cid_{cid} {
InitTxTime();
string_view cmd_name(cid_->name());
if (cmd_name == "EXEC" || cmd_name == "EVAL" || cmd_name == "EVALSHA") {
if (cmd_name == "EXEC" || cmd_name == "EVAL" || cmd_name == "EVAL_RO" || cmd_name == "EVALSHA" ||
cmd_name == "EVALSHA_RO") {
multi_.reset(new MultiData);
multi_->mode = NOT_DETERMINED;
multi_->role = DEFAULT;
Expand Down

0 comments on commit 6bfdacd

Please sign in to comment.