wehub-resource-sync
1634 lines
60 KiB
C++
1634 lines
60 KiB
C++
// Licensed to the Apache Software Foundation (ASF) under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you under the Apache License, Version 2.0 (the
|
|
// "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing,
|
|
// software distributed under the License is distributed on an
|
|
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
// KIND, either express or implied. See the License for the
|
|
// specific language governing permissions and limitations
|
|
// under the License.
|
|
|
|
|
|
#include <iostream>
|
|
#include <memory>
|
|
#include <unordered_map>
|
|
#include <butil/time.h>
|
|
#include <butil/logging.h>
|
|
#include <brpc/redis.h>
|
|
#include <brpc/channel.h>
|
|
#include <brpc/policy/redis_authenticator.h>
|
|
#include <brpc/server.h>
|
|
#include <brpc/redis_command.h>
|
|
#include <gtest/gtest.h>
|
|
#include <gflags/gflags.h>
|
|
|
|
namespace brpc {
|
|
DECLARE_int32(idle_timeout_second);
|
|
DECLARE_int32(redis_max_allocation_size);
|
|
DECLARE_int32(redis_max_reply_depth);
|
|
}
|
|
|
|
int main(int argc, char* argv[]) {
|
|
brpc::FLAGS_idle_timeout_second = 0;
|
|
testing::InitGoogleTest(&argc, argv);
|
|
return RUN_ALL_TESTS();
|
|
}
|
|
|
|
namespace {
|
|
static pthread_once_t download_redis_server_once = PTHREAD_ONCE_INIT;
|
|
|
|
static pid_t g_redis_pid = -1;
|
|
|
|
static void RemoveRedisServer() {
|
|
if (g_redis_pid > 0) {
|
|
puts("[Stopping redis-server]");
|
|
char cmd[256];
|
|
#if defined(BAIDU_INTERNAL)
|
|
snprintf(cmd, sizeof(cmd), "kill %d; rm -rf redis_server_for_test", g_redis_pid);
|
|
#else
|
|
snprintf(cmd, sizeof(cmd), "kill %d", g_redis_pid);
|
|
#endif
|
|
CHECK(0 == system(cmd));
|
|
// Wait for redis to stop
|
|
usleep(50000);
|
|
}
|
|
}
|
|
|
|
#define REDIS_SERVER_BIN "redis-server"
|
|
#define REDIS_SERVER_PORT "6479"
|
|
|
|
static void RunRedisServer() {
|
|
#if defined(BAIDU_INTERNAL)
|
|
puts("Downloading redis-server...");
|
|
if (system("mkdir -p redis_server_for_test && cd redis_server_for_test && svn co https://svn.baidu.com/third-64/tags/redis/redis_2-6-14-100_PD_BL/bin") != 0) {
|
|
puts("Fail to get redis-server from svn");
|
|
return;
|
|
}
|
|
# undef REDIS_SERVER_BIN
|
|
# define REDIS_SERVER_BIN "redis_server_for_test/bin/redis-server";
|
|
#else
|
|
if (system("which " REDIS_SERVER_BIN) != 0) {
|
|
puts("Fail to find " REDIS_SERVER_BIN ", following tests will be skipped");
|
|
return;
|
|
}
|
|
#endif
|
|
atexit(RemoveRedisServer);
|
|
|
|
g_redis_pid = fork();
|
|
if (g_redis_pid < 0) {
|
|
puts("Fail to fork");
|
|
exit(1);
|
|
} else if (g_redis_pid == 0) {
|
|
puts("[Starting redis-server]");
|
|
char* const argv[] = { (char*)REDIS_SERVER_BIN,
|
|
(char*)"--port", (char*)REDIS_SERVER_PORT,
|
|
NULL };
|
|
unlink("dump.rdb");
|
|
if (execvp(REDIS_SERVER_BIN, argv) < 0) {
|
|
puts("Fail to run " REDIS_SERVER_BIN);
|
|
exit(1);
|
|
}
|
|
}
|
|
// Wait for redis to start.
|
|
usleep(50000);
|
|
}
|
|
|
|
class ScopedRedisMaxReplyDepth {
|
|
public:
|
|
explicit ScopedRedisMaxReplyDepth(int32_t depth)
|
|
: _old_depth(brpc::FLAGS_redis_max_reply_depth) {
|
|
brpc::FLAGS_redis_max_reply_depth = depth;
|
|
}
|
|
~ScopedRedisMaxReplyDepth() {
|
|
brpc::FLAGS_redis_max_reply_depth = _old_depth;
|
|
}
|
|
|
|
private:
|
|
int32_t _old_depth;
|
|
};
|
|
|
|
class RedisTest : public testing::Test {
|
|
protected:
|
|
RedisTest() {}
|
|
void SetUp() {
|
|
pthread_once(&download_redis_server_once, RunRedisServer);
|
|
}
|
|
void TearDown() {}
|
|
};
|
|
|
|
void AssertReplyEqual(const brpc::RedisReply& reply1,
|
|
const brpc::RedisReply& reply2) {
|
|
if (&reply1 == &reply2) {
|
|
return;
|
|
}
|
|
CHECK_EQ(reply1.type(), reply2.type());
|
|
switch (reply1.type()) {
|
|
case brpc::REDIS_REPLY_ARRAY:
|
|
ASSERT_EQ(reply1.size(), reply2.size());
|
|
for (size_t j = 0; j < reply1.size(); ++j) {
|
|
ASSERT_NE(&reply1[j], &reply2[j]); // from different arena
|
|
AssertReplyEqual(reply1[j], reply2[j]);
|
|
}
|
|
break;
|
|
case brpc::REDIS_REPLY_INTEGER:
|
|
ASSERT_EQ(reply1.integer(), reply2.integer());
|
|
break;
|
|
case brpc::REDIS_REPLY_NIL:
|
|
break;
|
|
case brpc::REDIS_REPLY_STRING:
|
|
// fall through
|
|
case brpc::REDIS_REPLY_STATUS:
|
|
ASSERT_NE(reply1.c_str(), reply2.c_str()); // from different arena
|
|
ASSERT_EQ(reply1.data(), reply2.data());
|
|
break;
|
|
case brpc::REDIS_REPLY_ERROR:
|
|
ASSERT_NE(reply1.error_message(), reply2.error_message()); // from different arena
|
|
ASSERT_STREQ(reply1.error_message(), reply2.error_message());
|
|
break;
|
|
}
|
|
}
|
|
|
|
void AssertResponseEqual(const brpc::RedisResponse& r1,
|
|
const brpc::RedisResponse& r2,
|
|
int repeated_times = 1) {
|
|
if (&r1 == &r2) {
|
|
ASSERT_EQ(repeated_times, 1);
|
|
return;
|
|
}
|
|
ASSERT_EQ(r2.reply_size()* repeated_times, r1.reply_size());
|
|
for (int j = 0; j < repeated_times; ++j) {
|
|
for (int i = 0; i < r2.reply_size(); ++i) {
|
|
ASSERT_NE(&r2.reply(i), &r1.reply(j * r2.reply_size() + i));
|
|
AssertReplyEqual(r2.reply(i), r1.reply(j * r2.reply_size() + i));
|
|
}
|
|
}
|
|
}
|
|
|
|
TEST_F(RedisTest, sanity) {
|
|
if (g_redis_pid < 0) {
|
|
puts("Skipped due to absence of redis-server");
|
|
return;
|
|
}
|
|
brpc::ChannelOptions options;
|
|
options.protocol = brpc::PROTOCOL_REDIS;
|
|
brpc::Channel channel;
|
|
ASSERT_EQ(0, channel.Init("0.0.0.0:" REDIS_SERVER_PORT, &options));
|
|
brpc::RedisRequest request;
|
|
brpc::RedisResponse response;
|
|
brpc::Controller cntl;
|
|
|
|
ASSERT_TRUE(request.AddCommand("get hello"));
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(1, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_NIL, response.reply(0).type())
|
|
<< response;
|
|
|
|
cntl.Reset();
|
|
request.Clear();
|
|
response.Clear();
|
|
request.AddCommand("set hello world");
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(1, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(0).type());
|
|
ASSERT_EQ("OK", response.reply(0).data());
|
|
|
|
cntl.Reset();
|
|
request.Clear();
|
|
response.Clear();
|
|
ASSERT_TRUE(request.AddCommand("get hello"));
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed());
|
|
ASSERT_EQ(1, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(0).type());
|
|
ASSERT_EQ("world", response.reply(0).data());
|
|
|
|
cntl.Reset();
|
|
request.Clear();
|
|
response.Clear();
|
|
request.AddCommand("set hello world2");
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(1, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(0).type());
|
|
ASSERT_EQ("OK", response.reply(0).data());
|
|
|
|
cntl.Reset();
|
|
request.Clear();
|
|
response.Clear();
|
|
ASSERT_TRUE(request.AddCommand("get hello"));
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed());
|
|
ASSERT_EQ(1, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(0).type());
|
|
ASSERT_EQ("world2", response.reply(0).data());
|
|
|
|
cntl.Reset();
|
|
request.Clear();
|
|
response.Clear();
|
|
ASSERT_TRUE(request.AddCommand("del hello"));
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_INTEGER, response.reply(0).type());
|
|
ASSERT_EQ(1, response.reply(0).integer());
|
|
|
|
cntl.Reset();
|
|
request.Clear();
|
|
response.Clear();
|
|
ASSERT_TRUE(request.AddCommand("get %s", "hello"));
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(1, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_NIL, response.reply(0).type());
|
|
}
|
|
|
|
TEST_F(RedisTest, keys_with_spaces) {
|
|
if (g_redis_pid < 0) {
|
|
puts("Skipped due to absence of redis-server");
|
|
return;
|
|
}
|
|
brpc::ChannelOptions options;
|
|
options.protocol = brpc::PROTOCOL_REDIS;
|
|
brpc::Channel channel;
|
|
ASSERT_EQ(0, channel.Init("0.0.0.0:" REDIS_SERVER_PORT, &options));
|
|
brpc::RedisRequest request;
|
|
brpc::RedisResponse response;
|
|
brpc::Controller cntl;
|
|
|
|
cntl.Reset();
|
|
request.Clear();
|
|
response.Clear();
|
|
ASSERT_TRUE(request.AddCommand("set %s 'he1 he1 da1'", "hello world"));
|
|
ASSERT_TRUE(request.AddCommand("set 'hello2 world2' 'he2 he2 da2'"));
|
|
ASSERT_TRUE(request.AddCommand("set \"hello3 world3\" \"he3 he3 da3\""));
|
|
ASSERT_TRUE(request.AddCommand("get \"hello world\""));
|
|
ASSERT_TRUE(request.AddCommand("get 'hello world'"));
|
|
ASSERT_TRUE(request.AddCommand("get 'hello2 world2'"));
|
|
ASSERT_TRUE(request.AddCommand("get 'hello3 world3'"));
|
|
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(7, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(0).type());
|
|
ASSERT_EQ("OK", response.reply(0).data());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(1).type());
|
|
ASSERT_EQ("OK", response.reply(1).data());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(2).type());
|
|
ASSERT_EQ("OK", response.reply(2).data());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(3).type());
|
|
ASSERT_EQ("he1 he1 da1", response.reply(3).data());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(4).type());
|
|
ASSERT_EQ("he1 he1 da1", response.reply(4).data());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(5).type());
|
|
ASSERT_EQ("he2 he2 da2", response.reply(5).data());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(6).type());
|
|
ASSERT_EQ("he3 he3 da3", response.reply(6).data());
|
|
|
|
brpc::RedisResponse response2 = response;
|
|
AssertResponseEqual(response2, response);
|
|
response2.MergeFrom(response);
|
|
AssertResponseEqual(response2, response, 2);
|
|
}
|
|
|
|
TEST_F(RedisTest, incr_and_decr) {
|
|
if (g_redis_pid < 0) {
|
|
puts("Skipped due to absence of redis-server");
|
|
return;
|
|
}
|
|
brpc::ChannelOptions options;
|
|
options.protocol = brpc::PROTOCOL_REDIS;
|
|
brpc::Channel channel;
|
|
ASSERT_EQ(0, channel.Init("0.0.0.0:" REDIS_SERVER_PORT, &options));
|
|
brpc::RedisRequest request;
|
|
brpc::RedisResponse response;
|
|
brpc::Controller cntl;
|
|
|
|
request.AddCommand("incr counter1");
|
|
request.AddCommand("decr counter1");
|
|
request.AddCommand("incrby counter1 %d", 10);
|
|
request.AddCommand("decrby counter1 %d", 20);
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(4, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_INTEGER, response.reply(0).type());
|
|
ASSERT_EQ(1, response.reply(0).integer());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_INTEGER, response.reply(1).type());
|
|
ASSERT_EQ(0, response.reply(1).integer());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_INTEGER, response.reply(2).type());
|
|
ASSERT_EQ(10, response.reply(2).integer());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_INTEGER, response.reply(3).type());
|
|
ASSERT_EQ(-10, response.reply(3).integer());
|
|
|
|
brpc::RedisResponse response2 = response;
|
|
AssertResponseEqual(response2, response);
|
|
response2.MergeFrom(response);
|
|
AssertResponseEqual(response2, response, 2);
|
|
}
|
|
|
|
TEST_F(RedisTest, by_components) {
|
|
if (g_redis_pid < 0) {
|
|
puts("Skipped due to absence of redis-server");
|
|
return;
|
|
}
|
|
brpc::ChannelOptions options;
|
|
options.protocol = brpc::PROTOCOL_REDIS;
|
|
brpc::Channel channel;
|
|
ASSERT_EQ(0, channel.Init("0.0.0.0:" REDIS_SERVER_PORT, &options));
|
|
brpc::RedisRequest request;
|
|
brpc::RedisResponse response;
|
|
brpc::Controller cntl;
|
|
|
|
butil::StringPiece comp1[] = { "incr", "counter2" };
|
|
butil::StringPiece comp2[] = { "decr", "counter2" };
|
|
butil::StringPiece comp3[] = { "incrby", "counter2", "10" };
|
|
butil::StringPiece comp4[] = { "decrby", "counter2", "20" };
|
|
|
|
request.AddCommandByComponents(comp1, arraysize(comp1));
|
|
request.AddCommandByComponents(comp2, arraysize(comp2));
|
|
request.AddCommandByComponents(comp3, arraysize(comp3));
|
|
request.AddCommandByComponents(comp4, arraysize(comp4));
|
|
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(4, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_INTEGER, response.reply(0).type());
|
|
ASSERT_EQ(1, response.reply(0).integer());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_INTEGER, response.reply(1).type());
|
|
ASSERT_EQ(0, response.reply(1).integer());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_INTEGER, response.reply(2).type());
|
|
ASSERT_EQ(10, response.reply(2).integer());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_INTEGER, response.reply(3).type());
|
|
ASSERT_EQ(-10, response.reply(3).integer());
|
|
|
|
brpc::RedisResponse response2 = response;
|
|
AssertResponseEqual(response2, response);
|
|
response2.MergeFrom(response);
|
|
AssertResponseEqual(response2, response, 2);
|
|
}
|
|
|
|
static std::string GeneratePassword() {
|
|
std::string result;
|
|
result.reserve(12);
|
|
for (size_t i = 0; i < result.capacity(); ++i) {
|
|
result.push_back(butil::fast_rand_in('a', 'z'));
|
|
}
|
|
return result;
|
|
}
|
|
|
|
TEST_F(RedisTest, auth) {
|
|
if (g_redis_pid < 0) {
|
|
puts("Skipped due to absence of redis-server");
|
|
return;
|
|
}
|
|
// generate a random password
|
|
const std::string passwd1 = GeneratePassword();
|
|
const std::string passwd2 = GeneratePassword();
|
|
LOG(INFO) << "Generated passwd1=" << passwd1 << " passwd2=" << passwd2;
|
|
|
|
// config auth
|
|
{
|
|
brpc::ChannelOptions options;
|
|
options.protocol = brpc::PROTOCOL_REDIS;
|
|
brpc::Channel channel;
|
|
ASSERT_EQ(0, channel.Init("0.0.0.0:" REDIS_SERVER_PORT, &options));
|
|
brpc::RedisRequest request;
|
|
brpc::RedisResponse response;
|
|
brpc::Controller cntl;
|
|
|
|
request.AddCommand("set mykey %s", passwd1.c_str());
|
|
request.AddCommand("config set requirepass %s", passwd1.c_str());
|
|
request.AddCommand("auth %s", passwd1.c_str());
|
|
request.AddCommand("get mykey");
|
|
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(4, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(0).type());
|
|
ASSERT_STREQ("OK", response.reply(0).c_str());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(1).type());
|
|
ASSERT_STREQ("OK", response.reply(1).c_str());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(2).type());
|
|
ASSERT_STREQ("OK", response.reply(2).c_str());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(3).type());
|
|
ASSERT_STREQ(passwd1.c_str(), response.reply(3).c_str());
|
|
}
|
|
|
|
// Auth failed
|
|
{
|
|
brpc::ChannelOptions options;
|
|
options.protocol = brpc::PROTOCOL_REDIS;
|
|
brpc::Channel channel;
|
|
ASSERT_EQ(0, channel.Init("0.0.0.0:" REDIS_SERVER_PORT, &options));
|
|
brpc::RedisRequest request;
|
|
brpc::RedisResponse response;
|
|
brpc::Controller cntl;
|
|
|
|
request.AddCommand("get mykey");
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(1, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_ERROR, response.reply(0).type());
|
|
}
|
|
|
|
// Auth with RedisAuthenticator and change to passwd2 (setting to empty
|
|
// pass does not work on redis 6.0.6)
|
|
{
|
|
brpc::ChannelOptions options;
|
|
options.protocol = brpc::PROTOCOL_REDIS;
|
|
brpc::Channel channel;
|
|
brpc::policy::RedisAuthenticator* auth =
|
|
new brpc::policy::RedisAuthenticator(passwd1.c_str());
|
|
options.auth = auth;
|
|
ASSERT_EQ(0, channel.Init("0.0.0.0:" REDIS_SERVER_PORT, &options));
|
|
brpc::RedisRequest request;
|
|
brpc::RedisResponse response;
|
|
brpc::Controller cntl;
|
|
|
|
request.AddCommand("get mykey");
|
|
request.AddCommand("config set requirepass %s", passwd2.c_str());
|
|
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(2, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(0).type());
|
|
ASSERT_STREQ(passwd1.c_str(), response.reply(0).c_str());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(1).type());
|
|
ASSERT_STREQ("OK", response.reply(1).c_str());
|
|
}
|
|
|
|
// Auth with passwd2
|
|
{
|
|
brpc::ChannelOptions options;
|
|
options.protocol = brpc::PROTOCOL_REDIS;
|
|
brpc::policy::RedisAuthenticator* auth =
|
|
new brpc::policy::RedisAuthenticator(passwd2.c_str());
|
|
options.auth = auth;
|
|
brpc::Channel channel;
|
|
ASSERT_EQ(0, channel.Init("0.0.0.0:" REDIS_SERVER_PORT, &options));
|
|
brpc::RedisRequest request;
|
|
brpc::RedisResponse response;
|
|
brpc::Controller cntl;
|
|
|
|
request.AddCommand("get mykey");
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(1, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(0).type()) << response.reply(0);
|
|
ASSERT_STREQ(passwd1.c_str(), response.reply(0).c_str());
|
|
}
|
|
}
|
|
|
|
TEST_F(RedisTest, cmd_format) {
|
|
if (g_redis_pid < 0) {
|
|
puts("Skipped due to absence of redis-server");
|
|
return;
|
|
}
|
|
brpc::RedisRequest request;
|
|
// set empty string
|
|
request.AddCommand("set a ''");
|
|
ASSERT_STREQ("*3\r\n$3\r\nset\r\n$1\r\na\r\n$0\r\n\r\n",
|
|
request._buf.to_string().c_str());
|
|
request.Clear();
|
|
|
|
request.AddCommand("mset b '' c ''");
|
|
ASSERT_STREQ("*5\r\n$4\r\nmset\r\n$1\r\nb\r\n$0\r\n\r\n$1\r\nc\r\n$0\r\n\r\n",
|
|
request._buf.to_string().c_str());
|
|
request.Clear();
|
|
// set non-empty string
|
|
request.AddCommand("set a 123");
|
|
ASSERT_STREQ("*3\r\n$3\r\nset\r\n$1\r\na\r\n$3\r\n123\r\n",
|
|
request._buf.to_string().c_str());
|
|
request.Clear();
|
|
|
|
request.AddCommand("mset b '' c ccc");
|
|
ASSERT_STREQ("*5\r\n$4\r\nmset\r\n$1\r\nb\r\n$0\r\n\r\n$1\r\nc\r\n$3\r\nccc\r\n",
|
|
request._buf.to_string().c_str());
|
|
request.Clear();
|
|
|
|
request.AddCommand("get ''key value"); // == get <empty> key value
|
|
ASSERT_STREQ("*4\r\n$3\r\nget\r\n$0\r\n\r\n$3\r\nkey\r\n$5\r\nvalue\r\n", request._buf.to_string().c_str());
|
|
request.Clear();
|
|
|
|
request.AddCommand("get key'' value"); // == get key <empty> value
|
|
ASSERT_STREQ("*4\r\n$3\r\nget\r\n$3\r\nkey\r\n$0\r\n\r\n$5\r\nvalue\r\n", request._buf.to_string().c_str());
|
|
request.Clear();
|
|
|
|
request.AddCommand("get 'ext'key value "); // == get ext key value
|
|
ASSERT_STREQ("*4\r\n$3\r\nget\r\n$3\r\next\r\n$3\r\nkey\r\n$5\r\nvalue\r\n", request._buf.to_string().c_str());
|
|
request.Clear();
|
|
|
|
request.AddCommand(" get key'ext' value "); // == get key ext value
|
|
ASSERT_STREQ("*4\r\n$3\r\nget\r\n$3\r\nkey\r\n$3\r\next\r\n$5\r\nvalue\r\n", request._buf.to_string().c_str());
|
|
request.Clear();
|
|
}
|
|
|
|
TEST_F(RedisTest, quote_and_escape) {
|
|
if (g_redis_pid < 0) {
|
|
puts("Skipped due to absence of redis-server");
|
|
return;
|
|
}
|
|
brpc::RedisRequest request;
|
|
request.AddCommand("set a 'foo bar'");
|
|
ASSERT_STREQ("*3\r\n$3\r\nset\r\n$1\r\na\r\n$7\r\nfoo bar\r\n",
|
|
request._buf.to_string().c_str());
|
|
request.Clear();
|
|
|
|
request.AddCommand("set a 'foo \\'bar'");
|
|
ASSERT_STREQ("*3\r\n$3\r\nset\r\n$1\r\na\r\n$8\r\nfoo 'bar\r\n",
|
|
request._buf.to_string().c_str());
|
|
request.Clear();
|
|
|
|
request.AddCommand("set a 'foo \"bar'");
|
|
ASSERT_STREQ("*3\r\n$3\r\nset\r\n$1\r\na\r\n$8\r\nfoo \"bar\r\n",
|
|
request._buf.to_string().c_str());
|
|
request.Clear();
|
|
|
|
request.AddCommand("set a 'foo \\\"bar'");
|
|
ASSERT_STREQ("*3\r\n$3\r\nset\r\n$1\r\na\r\n$9\r\nfoo \\\"bar\r\n",
|
|
request._buf.to_string().c_str());
|
|
request.Clear();
|
|
|
|
request.AddCommand("set a \"foo 'bar\"");
|
|
ASSERT_STREQ("*3\r\n$3\r\nset\r\n$1\r\na\r\n$8\r\nfoo 'bar\r\n",
|
|
request._buf.to_string().c_str());
|
|
request.Clear();
|
|
|
|
request.AddCommand("set a \"foo \\'bar\"");
|
|
ASSERT_STREQ("*3\r\n$3\r\nset\r\n$1\r\na\r\n$9\r\nfoo \\'bar\r\n",
|
|
request._buf.to_string().c_str());
|
|
request.Clear();
|
|
|
|
request.AddCommand("set a \"foo \\\"bar\"");
|
|
ASSERT_STREQ("*3\r\n$3\r\nset\r\n$1\r\na\r\n$8\r\nfoo \"bar\r\n",
|
|
request._buf.to_string().c_str());
|
|
request.Clear();
|
|
}
|
|
|
|
std::string GetCompleteCommand(const std::vector<butil::StringPiece>& commands) {
|
|
std::string res;
|
|
for (int i = 0; i < (int)commands.size(); ++i) {
|
|
if (i != 0) {
|
|
res.push_back(' ');
|
|
}
|
|
res.append(commands[i].data(), commands[i].size());
|
|
}
|
|
return res;
|
|
}
|
|
|
|
|
|
TEST_F(RedisTest, command_parser) {
|
|
brpc::RedisCommandParser parser;
|
|
butil::IOBuf buf;
|
|
std::vector<butil::StringPiece> command_out;
|
|
butil::Arena arena;
|
|
{
|
|
// parse from whole command
|
|
std::string command = "set abc edc";
|
|
ASSERT_TRUE(brpc::RedisCommandNoFormat(&buf, command.c_str()).ok());
|
|
ASSERT_EQ(brpc::PARSE_OK, parser.Consume(buf, &command_out, &arena));
|
|
ASSERT_TRUE(buf.empty());
|
|
ASSERT_EQ(command, GetCompleteCommand(command_out));
|
|
}
|
|
{
|
|
// simulate parsing from network following RESP
|
|
int t = 100;
|
|
std::string raw_string("*3\r\n$3\r\nset\r\n$3\r\nabc\r\n$3\r\ndef\r\n");
|
|
int size = raw_string.size();
|
|
while (t--) {
|
|
for (int i = 0; i < size; ++i) {
|
|
buf.push_back(raw_string[i]);
|
|
if (i == size - 1) {
|
|
ASSERT_EQ(brpc::PARSE_OK, parser.Consume(buf, &command_out, &arena));
|
|
} else {
|
|
if (butil::fast_rand_less_than(2) == 0) {
|
|
ASSERT_EQ(brpc::PARSE_ERROR_NOT_ENOUGH_DATA,
|
|
parser.Consume(buf, &command_out, &arena));
|
|
}
|
|
}
|
|
}
|
|
ASSERT_TRUE(buf.empty());
|
|
ASSERT_EQ(GetCompleteCommand(command_out), "set abc def");
|
|
}
|
|
}
|
|
{
|
|
// simulate parsing from network under inline protocol
|
|
int t = 100;
|
|
std::string raw_string("set abc def\r\n");
|
|
int size = raw_string.size();
|
|
while (t--) {
|
|
for (int i = 0; i < size; ++i) {
|
|
buf.push_back(raw_string[i]);
|
|
if (i == size - 1) {
|
|
ASSERT_EQ(brpc::PARSE_OK, parser.Consume(buf, &command_out, &arena));
|
|
} else {
|
|
if (butil::fast_rand_less_than(2) == 0) {
|
|
ASSERT_EQ(brpc::PARSE_ERROR_NOT_ENOUGH_DATA,
|
|
parser.Consume(buf, &command_out, &arena));
|
|
}
|
|
}
|
|
}
|
|
ASSERT_TRUE(buf.empty());
|
|
ASSERT_EQ(GetCompleteCommand(command_out), "set abc def");
|
|
}
|
|
}
|
|
{
|
|
// there is a non-string message in command and parse should fail
|
|
buf.append("*3\r\n$3");
|
|
ASSERT_EQ(brpc::PARSE_ERROR_NOT_ENOUGH_DATA, parser.Consume(buf, &command_out, &arena));
|
|
ASSERT_EQ((int)buf.size(), 2); // left "$3"
|
|
buf.append("\r\nset\r\n:123\r\n$3\r\ndef\r\n");
|
|
ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, parser.Consume(buf, &command_out, &arena));
|
|
parser.Reset();
|
|
}
|
|
{
|
|
// not array
|
|
buf.append(":123456\r\n");
|
|
ASSERT_EQ(brpc::PARSE_ERROR_TRY_OTHERS, parser.Consume(buf, &command_out, &arena));
|
|
parser.Reset();
|
|
}
|
|
{
|
|
// not array
|
|
buf.append("+Error\r\n");
|
|
ASSERT_EQ(brpc::PARSE_ERROR_TRY_OTHERS, parser.Consume(buf, &command_out, &arena));
|
|
parser.Reset();
|
|
}
|
|
{
|
|
// not array
|
|
buf.append("+OK\r\n");
|
|
ASSERT_EQ(brpc::PARSE_ERROR_TRY_OTHERS, parser.Consume(buf, &command_out, &arena));
|
|
parser.Reset();
|
|
}
|
|
{
|
|
// not array
|
|
buf.append("$5\r\nhello\r\n");
|
|
ASSERT_EQ(brpc::PARSE_ERROR_TRY_OTHERS, parser.Consume(buf, &command_out, &arena));
|
|
parser.Reset();
|
|
}
|
|
}
|
|
|
|
TEST(RedisCommandParserTest, reject_empty_resp_array_command) {
|
|
brpc::RedisCommandParser parser;
|
|
butil::IOBuf buf;
|
|
std::vector<butil::StringPiece> command_out;
|
|
butil::Arena arena;
|
|
|
|
// Empty RESP arrays are not valid commands and must not leave the parser in
|
|
// a state that accepts following bulk strings.
|
|
buf.append("*0\r\n$1\r\nx\r\n");
|
|
ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG,
|
|
parser.Consume(buf, &command_out, &arena));
|
|
}
|
|
|
|
// Regression test for issue #3109: the inline redis protocol must not consume
|
|
// the HTTP/2 connection preface as a command, otherwise protocol auto-detection
|
|
// never falls through to HTTP/2 and gRPC clients fail with "connection closed
|
|
// before server preface received".
|
|
TEST_F(RedisTest, inline_does_not_eat_h2_preface) {
|
|
brpc::RedisCommandParser parser;
|
|
butil::IOBuf buf;
|
|
std::vector<butil::StringPiece> command_out;
|
|
butil::Arena arena;
|
|
{
|
|
// Full HTTP/2 client connection preface: must defer to other protocols.
|
|
buf.append("PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n");
|
|
ASSERT_EQ(brpc::PARSE_ERROR_TRY_OTHERS,
|
|
parser.Consume(buf, &command_out, &arena));
|
|
buf.clear();
|
|
parser.Reset();
|
|
}
|
|
{
|
|
// A not-yet-complete prefix of the preface must also defer, leaving the
|
|
// bytes intact for the HTTP/2 parser instead of being misparsed.
|
|
buf.append("PRI * HT");
|
|
ASSERT_EQ(brpc::PARSE_ERROR_TRY_OTHERS,
|
|
parser.Consume(buf, &command_out, &arena));
|
|
buf.clear();
|
|
parser.Reset();
|
|
}
|
|
{
|
|
// A genuine inline command sharing the leading 'P' must still parse.
|
|
buf.append("PING\r\n");
|
|
ASSERT_EQ(brpc::PARSE_OK, parser.Consume(buf, &command_out, &arena));
|
|
ASSERT_EQ("ping", GetCompleteCommand(command_out));
|
|
buf.clear();
|
|
parser.Reset();
|
|
}
|
|
}
|
|
|
|
TEST_F(RedisTest, redis_reply_codec) {
|
|
butil::Arena arena;
|
|
// status
|
|
{
|
|
brpc::RedisReply r(&arena);
|
|
butil::IOBuf buf;
|
|
butil::IOBufAppender appender;
|
|
r.SetStatus("OK");
|
|
ASSERT_TRUE(r.SerializeTo(&appender));
|
|
appender.move_to(buf);
|
|
ASSERT_STREQ(buf.to_string().c_str(), "+OK\r\n");
|
|
ASSERT_STREQ(r.c_str(), "OK");
|
|
|
|
brpc::RedisReply r2(&arena);
|
|
brpc::ParseError err = r2.ConsumePartialIOBuf(buf);
|
|
ASSERT_EQ(err, brpc::PARSE_OK);
|
|
ASSERT_TRUE(r2.is_string());
|
|
ASSERT_STREQ("OK", r2.c_str());
|
|
}
|
|
// error
|
|
{
|
|
brpc::RedisReply r(&arena);
|
|
butil::IOBuf buf;
|
|
butil::IOBufAppender appender;
|
|
r.SetError("not exist \'key\'");
|
|
ASSERT_TRUE(r.SerializeTo(&appender));
|
|
appender.move_to(buf);
|
|
ASSERT_STREQ(buf.to_string().c_str(), "-not exist \'key\'\r\n");
|
|
|
|
brpc::RedisReply r2(&arena);
|
|
brpc::ParseError err = r2.ConsumePartialIOBuf(buf);
|
|
ASSERT_EQ(err, brpc::PARSE_OK);
|
|
ASSERT_TRUE(r2.is_error());
|
|
ASSERT_STREQ("not exist \'key\'", r2.error_message());
|
|
}
|
|
// string
|
|
{
|
|
brpc::RedisReply r(&arena);
|
|
butil::IOBuf buf;
|
|
butil::IOBufAppender appender;
|
|
r.SetNullString();
|
|
ASSERT_TRUE(r.SerializeTo(&appender));
|
|
appender.move_to(buf);
|
|
ASSERT_STREQ(buf.to_string().c_str(), "$-1\r\n");
|
|
|
|
brpc::RedisReply r2(&arena);
|
|
brpc::ParseError err = r2.ConsumePartialIOBuf(buf);
|
|
ASSERT_EQ(err, brpc::PARSE_OK);
|
|
ASSERT_TRUE(r2.is_nil());
|
|
|
|
r.SetString("abcde'hello world");
|
|
ASSERT_TRUE(r.SerializeTo(&appender));
|
|
appender.move_to(buf);
|
|
ASSERT_STREQ(buf.to_string().c_str(), "$17\r\nabcde'hello world\r\n");
|
|
ASSERT_STREQ("abcde'hello world", r.c_str());
|
|
|
|
r.FormatString("int:%d str:%s fp:%.2f", 123, "foobar", 3.21);
|
|
ASSERT_TRUE(r.SerializeTo(&appender));
|
|
appender.move_to(buf);
|
|
ASSERT_STREQ(buf.to_string().c_str(), "$26\r\nint:123 str:foobar fp:3.21\r\n");
|
|
ASSERT_STREQ("int:123 str:foobar fp:3.21", r.c_str());
|
|
|
|
r.FormatString("verylongstring verylongstring verylongstring verylongstring int:%d str:%s fp:%.2f", 123, "foobar", 3.21);
|
|
ASSERT_TRUE(r.SerializeTo(&appender));
|
|
appender.move_to(buf);
|
|
ASSERT_STREQ(buf.to_string().c_str(), "$86\r\nverylongstring verylongstring verylongstring verylongstring int:123 str:foobar fp:3.21\r\n");
|
|
ASSERT_STREQ("verylongstring verylongstring verylongstring verylongstring int:123 str:foobar fp:3.21", r.c_str());
|
|
|
|
brpc::RedisReply r3(&arena);
|
|
err = r3.ConsumePartialIOBuf(buf);
|
|
ASSERT_EQ(err, brpc::PARSE_OK);
|
|
ASSERT_TRUE(r3.is_string());
|
|
ASSERT_STREQ(r.c_str(), r3.c_str());
|
|
}
|
|
// integer
|
|
{
|
|
brpc::RedisReply r(&arena);
|
|
butil::IOBuf buf;
|
|
butil::IOBufAppender appender;
|
|
int t = 2;
|
|
int input[] = { -1, 1234567 };
|
|
const char* output[] = { ":-1\r\n", ":1234567\r\n" };
|
|
for (int i = 0; i < t; ++i) {
|
|
r.SetInteger(input[i]);
|
|
ASSERT_TRUE(r.SerializeTo(&appender));
|
|
appender.move_to(buf);
|
|
ASSERT_STREQ(buf.to_string().c_str(), output[i]);
|
|
|
|
brpc::RedisReply r2(&arena);
|
|
brpc::ParseError err = r2.ConsumePartialIOBuf(buf);
|
|
ASSERT_EQ(err, brpc::PARSE_OK);
|
|
ASSERT_TRUE(r2.is_integer());
|
|
ASSERT_EQ(r2.integer(), input[i]);
|
|
}
|
|
}
|
|
// array
|
|
{
|
|
brpc::RedisReply r(&arena);
|
|
butil::IOBuf buf;
|
|
butil::IOBufAppender appender;
|
|
r.SetArray(3);
|
|
brpc::RedisReply& sub_reply = r[0];
|
|
sub_reply.SetArray(2);
|
|
sub_reply[0].SetString("hello, it's me");
|
|
sub_reply[1].SetInteger(422);
|
|
r[1].SetString("To go over everything");
|
|
r[2].SetInteger(1);
|
|
ASSERT_TRUE(r[3].is_nil());
|
|
ASSERT_TRUE(r.SerializeTo(&appender));
|
|
appender.move_to(buf);
|
|
ASSERT_STREQ(buf.to_string().c_str(),
|
|
"*3\r\n*2\r\n$14\r\nhello, it's me\r\n:422\r\n$21\r\n"
|
|
"To go over everything\r\n:1\r\n");
|
|
|
|
brpc::RedisReply r2(&arena);
|
|
ASSERT_EQ(r2.ConsumePartialIOBuf(buf), brpc::PARSE_OK);
|
|
ASSERT_TRUE(r2.is_array());
|
|
ASSERT_EQ(3ul, r2.size());
|
|
ASSERT_TRUE(r2[0].is_array());
|
|
ASSERT_EQ(2ul, r2[0].size());
|
|
ASSERT_TRUE(r2[0][0].is_string());
|
|
ASSERT_STREQ(r2[0][0].c_str(), "hello, it's me");
|
|
ASSERT_TRUE(r2[0][1].is_integer());
|
|
ASSERT_EQ(r2[0][1].integer(), 422);
|
|
ASSERT_TRUE(r2[1].is_string());
|
|
ASSERT_STREQ(r2[1].c_str(), "To go over everything");
|
|
ASSERT_TRUE(r2[2].is_integer());
|
|
ASSERT_EQ(1, r2[2].integer());
|
|
|
|
// null array
|
|
r.SetNullArray();
|
|
ASSERT_TRUE(r.SerializeTo(&appender));
|
|
appender.move_to(buf);
|
|
ASSERT_STREQ(buf.to_string().c_str(), "*-1\r\n");
|
|
ASSERT_EQ(r.ConsumePartialIOBuf(buf), brpc::PARSE_OK);
|
|
ASSERT_TRUE(r.is_nil());
|
|
}
|
|
|
|
// CopyFromDifferentArena
|
|
{
|
|
brpc::RedisReply r(&arena);
|
|
r.SetArray(1);
|
|
brpc::RedisReply& sub_reply = r[0];
|
|
sub_reply.SetArray(2);
|
|
sub_reply[0].SetString("hello, it's me");
|
|
sub_reply[1].SetInteger(422);
|
|
|
|
brpc::RedisReply r2(&arena);
|
|
r2.CopyFromDifferentArena(r);
|
|
ASSERT_TRUE(r2.is_array());
|
|
ASSERT_EQ((int)r2[0].size(), 2);
|
|
ASSERT_STREQ(r2[0][0].c_str(), sub_reply[0].c_str());
|
|
ASSERT_EQ(r2[0][1].integer(), sub_reply[1].integer());
|
|
}
|
|
// SetXXX can be called multiple times.
|
|
{
|
|
brpc::RedisReply r(&arena);
|
|
r.SetStatus("OK");
|
|
ASSERT_TRUE(r.is_string());
|
|
r.SetNullString();
|
|
ASSERT_TRUE(r.is_nil());
|
|
r.SetArray(2);
|
|
ASSERT_TRUE(r.is_array());
|
|
r.SetString("OK");
|
|
ASSERT_TRUE(r.is_string());
|
|
r.SetError("OK");
|
|
ASSERT_TRUE(r.is_error());
|
|
r.SetInteger(42);
|
|
ASSERT_TRUE(r.is_integer());
|
|
}
|
|
}
|
|
|
|
TEST_F(RedisTest, redis_reply_rejects_deep_nested_arrays) {
|
|
ScopedRedisMaxReplyDepth scoped_depth(4);
|
|
|
|
butil::IOBuf buf;
|
|
for (int i = 0; i <= brpc::FLAGS_redis_max_reply_depth; ++i) {
|
|
buf.append("*1\r\n");
|
|
}
|
|
buf.append(":0\r\n");
|
|
|
|
butil::Arena arena;
|
|
brpc::RedisReply reply(&arena);
|
|
EXPECT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, reply.ConsumePartialIOBuf(buf));
|
|
|
|
buf.clear();
|
|
for (int i = 0; i < brpc::FLAGS_redis_max_reply_depth; ++i) {
|
|
buf.append("*1\r\n");
|
|
}
|
|
buf.append(":0\r\n");
|
|
|
|
brpc::RedisReply valid_reply(&arena);
|
|
EXPECT_EQ(brpc::PARSE_OK, valid_reply.ConsumePartialIOBuf(buf));
|
|
EXPECT_TRUE(valid_reply.is_array());
|
|
}
|
|
|
|
butil::Mutex s_mutex;
|
|
std::unordered_map<std::string, std::string> m;
|
|
std::unordered_map<std::string, int64_t> int_map;
|
|
|
|
|
|
class RedisServiceImpl : public brpc::RedisService {
|
|
public:
|
|
RedisServiceImpl(std::string password)
|
|
: _batch_count(0)
|
|
, _password(std::move(password)) {}
|
|
|
|
brpc::RedisCommandHandlerResult OnBatched(const std::vector<butil::StringPiece>& args,
|
|
brpc::RedisReply* output, bool flush_batched) {
|
|
if (_batched_command.empty() && flush_batched) {
|
|
if (args[0] == "set") {
|
|
DoSet(args[1].as_string(), args[2].as_string(), output);
|
|
} else if (args[0] == "get") {
|
|
DoGet(args[1].as_string(), output);
|
|
}
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
std::vector<std::string> comm;
|
|
for (int i = 0; i < (int)args.size(); ++i) {
|
|
comm.push_back(args[i].as_string());
|
|
}
|
|
_batched_command.push_back(comm);
|
|
if (flush_batched) {
|
|
output->SetArray(_batched_command.size());
|
|
for (int i = 0; i < (int)_batched_command.size(); ++i) {
|
|
if (_batched_command[i][0] == "set") {
|
|
DoSet(_batched_command[i][1], _batched_command[i][2], &(*output)[i]);
|
|
} else if (_batched_command[i][0] == "get") {
|
|
DoGet(_batched_command[i][1], &(*output)[i]);
|
|
}
|
|
}
|
|
_batch_count++;
|
|
_batched_command.clear();
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
} else {
|
|
return brpc::REDIS_CMD_BATCHED;
|
|
}
|
|
}
|
|
|
|
void DoSet(const std::string& key, const std::string& value, brpc::RedisReply* output) {
|
|
m[key] = value;
|
|
output->SetStatus("OK");
|
|
}
|
|
|
|
void DoGet(const std::string& key, brpc::RedisReply* output) {
|
|
auto it = m.find(key);
|
|
if (it != m.end()) {
|
|
output->SetString(it->second);
|
|
} else {
|
|
output->SetNullString();
|
|
}
|
|
}
|
|
|
|
std::vector<std::vector<std::string> > _batched_command;
|
|
int _batch_count;
|
|
std::string _password;
|
|
};
|
|
|
|
|
|
class AuthSession : public brpc::Destroyable {
|
|
public:
|
|
explicit AuthSession(std::string password)
|
|
: _password(std::move(password)) {}
|
|
|
|
void Destroy() override {
|
|
delete this;
|
|
}
|
|
|
|
const std::string _password;
|
|
};
|
|
|
|
class AuthCommandHandler : public brpc::RedisCommandHandler {
|
|
public:
|
|
AuthCommandHandler(RedisServiceImpl* rs)
|
|
: _rs(rs) {}
|
|
|
|
brpc::RedisCommandHandlerResult Run(brpc::RedisConnContext* ctx,
|
|
const std::vector<butil::StringPiece>& args,
|
|
brpc::RedisReply* output,
|
|
bool flush_batched) {
|
|
if (args.size() < 1) {
|
|
output->SetError("ERR wrong number of arguments for 'AUTH' command");
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
const std::string password(args[1].data(), args[1].size());
|
|
if (_rs->_password != password) {
|
|
output->SetError("ERR invalid username/password");
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
auto auth_session = new AuthSession(password);
|
|
ctx->reset_session(auth_session);
|
|
output->SetStatus("OK");
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
|
|
private:
|
|
RedisServiceImpl* _rs;
|
|
};
|
|
|
|
class SetCommandHandler : public brpc::RedisCommandHandler {
|
|
public:
|
|
SetCommandHandler(RedisServiceImpl* rs, bool batch_process = false)
|
|
: _rs(rs)
|
|
, _batch_process(batch_process) {}
|
|
|
|
brpc::RedisCommandHandlerResult Run(brpc::RedisConnContext* ctx,
|
|
const std::vector<butil::StringPiece>& args,
|
|
brpc::RedisReply* output,
|
|
bool flush_batched) {
|
|
if (!ctx->session) {
|
|
output->SetError("ERR no auth");
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
AuthSession* session = static_cast<AuthSession*>(ctx->session);
|
|
if (!session || (session->_password != _rs->_password)) {
|
|
output->SetError("ERR no auth");
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
if (args.size() < 3) {
|
|
output->SetError("ERR wrong number of arguments for 'set' command");
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
if (_batch_process) {
|
|
return _rs->OnBatched(args, output, flush_batched);
|
|
} else {
|
|
DoSet(args[1].as_string(), args[2].as_string(), output);
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
}
|
|
|
|
void DoSet(const std::string& key, const std::string& value, brpc::RedisReply* output) {
|
|
m[key] = value;
|
|
output->SetStatus("OK");
|
|
}
|
|
|
|
private:
|
|
RedisServiceImpl* _rs;
|
|
bool _batch_process;
|
|
};
|
|
|
|
|
|
class GetCommandHandler : public brpc::RedisCommandHandler {
|
|
public:
|
|
GetCommandHandler(RedisServiceImpl* rs, bool batch_process = false)
|
|
: _rs(rs)
|
|
, _batch_process(batch_process) {}
|
|
|
|
brpc::RedisCommandHandlerResult Run(brpc::RedisConnContext* ctx,
|
|
const std::vector<butil::StringPiece>& args,
|
|
brpc::RedisReply* output,
|
|
bool flush_batched) {
|
|
if (!ctx->session) {
|
|
output->SetError("ERR no auth");
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
AuthSession* session = static_cast<AuthSession*>(ctx->session);
|
|
if (session->_password != _rs->_password) {
|
|
output->SetError("ERR no auth");
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
if (args.size() < 2) {
|
|
output->SetError("ERR wrong number of arguments for 'get' command");
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
if (_batch_process) {
|
|
return _rs->OnBatched(args, output, flush_batched);
|
|
} else {
|
|
DoGet(args[1].as_string(), output);
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
}
|
|
|
|
void DoGet(const std::string& key, brpc::RedisReply* output) {
|
|
auto it = m.find(key);
|
|
if (it != m.end()) {
|
|
output->SetString(it->second);
|
|
} else {
|
|
output->SetNullString();
|
|
}
|
|
}
|
|
|
|
private:
|
|
RedisServiceImpl* _rs;
|
|
bool _batch_process;
|
|
};
|
|
|
|
class IncrCommandHandler : public brpc::RedisCommandHandler {
|
|
public:
|
|
IncrCommandHandler(RedisServiceImpl* rs)
|
|
: _rs(rs) {}
|
|
|
|
brpc::RedisCommandHandlerResult Run(brpc::RedisConnContext* ctx,
|
|
const std::vector<butil::StringPiece>& args,
|
|
brpc::RedisReply* output,
|
|
bool flush_batched) {
|
|
if (!ctx->session) {
|
|
output->SetError("ERR no auth");
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
AuthSession* session = static_cast<AuthSession*>(ctx->session);
|
|
if (session->_password != _rs->_password) {
|
|
output->SetError("ERR no auth");
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
if (args.size() < 2) {
|
|
output->SetError("ERR wrong number of arguments for 'incr' command");
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
int64_t value;
|
|
s_mutex.lock();
|
|
value = ++int_map[args[1].as_string()];
|
|
s_mutex.unlock();
|
|
output->SetInteger(value);
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
|
|
private:
|
|
RedisServiceImpl* _rs;
|
|
};
|
|
|
|
TEST_F(RedisTest, server_sanity) {
|
|
std::string password = GeneratePassword();
|
|
std::unique_ptr<brpc::policy::RedisAuthenticator> redis_auth_holder(
|
|
new brpc::policy::RedisAuthenticator(password));
|
|
RedisServiceImpl* rsimpl = new RedisServiceImpl(password);
|
|
std::unique_ptr<GetCommandHandler> gh(new GetCommandHandler(rsimpl));
|
|
std::unique_ptr<SetCommandHandler> sh(new SetCommandHandler(rsimpl));
|
|
std::unique_ptr<AuthCommandHandler> ah(new AuthCommandHandler(rsimpl));
|
|
std::unique_ptr<IncrCommandHandler> ih(new IncrCommandHandler(rsimpl));
|
|
brpc::Server server;
|
|
brpc::ServerOptions server_options;
|
|
rsimpl->AddCommandHandler("get", gh.get());
|
|
rsimpl->AddCommandHandler("set", sh.get());
|
|
rsimpl->AddCommandHandler("incr", ih.get());
|
|
rsimpl->AddCommandHandler("auth", ah.get());
|
|
server_options.redis_service = rsimpl;
|
|
brpc::PortRange pr(8081, 8900);
|
|
ASSERT_EQ(0, server.Start("127.0.0.1", pr, &server_options));
|
|
|
|
brpc::ChannelOptions options;
|
|
options.protocol = brpc::PROTOCOL_REDIS;
|
|
options.auth = redis_auth_holder.get();
|
|
brpc::Channel channel;
|
|
ASSERT_EQ(0, channel.Init("127.0.0.1", server.listen_address().port, &options));
|
|
|
|
brpc::RedisRequest request;
|
|
brpc::RedisResponse response;
|
|
brpc::Controller cntl;
|
|
ASSERT_TRUE(request.AddCommand("get hello"));
|
|
ASSERT_TRUE(request.AddCommand("get hello2"));
|
|
ASSERT_TRUE(request.AddCommand("set key1 value1"));
|
|
ASSERT_TRUE(request.AddCommand("get key1"));
|
|
ASSERT_TRUE(request.AddCommand("set key2 value2"));
|
|
ASSERT_TRUE(request.AddCommand("get key2"));
|
|
ASSERT_TRUE(request.AddCommand("xxxcommand key2"));
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(7, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_NIL, response.reply(0).type());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_NIL, response.reply(1).type());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(2).type());
|
|
ASSERT_STREQ("OK", response.reply(2).c_str());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(3).type());
|
|
ASSERT_STREQ("value1", response.reply(3).c_str());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(4).type());
|
|
ASSERT_STREQ("OK", response.reply(4).c_str());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(5).type());
|
|
ASSERT_STREQ("value2", response.reply(5).c_str());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_ERROR, response.reply(6).type());
|
|
ASSERT_TRUE(butil::StringPiece(response.reply(6).error_message()).starts_with("ERR unknown command"));
|
|
|
|
cntl.Reset();
|
|
request.Clear();
|
|
response.Clear();
|
|
std::string value3("value3");
|
|
value3.append(1, '\0');
|
|
value3.append(1, 'a');
|
|
std::vector<butil::StringPiece> pieces;
|
|
pieces.push_back("set");
|
|
pieces.push_back("key3");
|
|
pieces.push_back(value3);
|
|
ASSERT_TRUE(request.AddCommandByComponents(&pieces[0], pieces.size()));
|
|
ASSERT_TRUE(request.AddCommand("set key4 \"\""));
|
|
ASSERT_TRUE(request.AddCommand("get key3"));
|
|
ASSERT_TRUE(request.AddCommand("get key4"));
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(4, response.reply_size());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(0).type());
|
|
ASSERT_STREQ("OK", response.reply(0).c_str());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(1).type());
|
|
ASSERT_STREQ("OK", response.reply(1).c_str());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(2).type());
|
|
ASSERT_STREQ("value3", response.reply(2).c_str());
|
|
ASSERT_NE("value3", response.reply(2).data());
|
|
ASSERT_EQ(value3, response.reply(2).data());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(3).type());
|
|
ASSERT_EQ("", response.reply(3).data());
|
|
}
|
|
|
|
void* incr_thread(void* arg) {
|
|
brpc::Channel* c = static_cast<brpc::Channel*>(arg);
|
|
for (int i = 0; i < 5000; ++i) {
|
|
brpc::RedisRequest request;
|
|
brpc::RedisResponse response;
|
|
brpc::Controller cntl;
|
|
EXPECT_TRUE(request.AddCommand("incr count"));
|
|
c->CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
EXPECT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
EXPECT_EQ(1, response.reply_size());
|
|
EXPECT_TRUE(response.reply(0).is_integer()) << response.reply(0);
|
|
}
|
|
return NULL;
|
|
}
|
|
|
|
TEST_F(RedisTest, server_concurrency) {
|
|
std::string password = GeneratePassword();
|
|
int N = 10;
|
|
std::unique_ptr<brpc::policy::RedisAuthenticator> redis_auth_holder(
|
|
new brpc::policy::RedisAuthenticator(password));
|
|
RedisServiceImpl* rsimpl = new RedisServiceImpl(password);
|
|
std::unique_ptr<AuthCommandHandler> ah(new AuthCommandHandler(rsimpl));
|
|
std::unique_ptr<IncrCommandHandler> ih(new IncrCommandHandler(rsimpl));
|
|
brpc::Server server;
|
|
brpc::ServerOptions server_options;
|
|
rsimpl->AddCommandHandler("incr", ih.get());
|
|
rsimpl->AddCommandHandler("auth", ah.get());
|
|
server_options.redis_service = rsimpl;
|
|
brpc::PortRange pr(8081, 8900);
|
|
ASSERT_EQ(0, server.Start("0.0.0.0", pr, &server_options));
|
|
|
|
brpc::ChannelOptions options;
|
|
options.protocol = brpc::PROTOCOL_REDIS;
|
|
options.connection_type = "pooled";
|
|
options.auth = redis_auth_holder.get();
|
|
std::vector<bthread_t> bths;
|
|
std::vector<brpc::Channel*> channels;
|
|
for (int i = 0; i < N; ++i) {
|
|
channels.push_back(new brpc::Channel);
|
|
ASSERT_EQ(0, channels.back()->Init("127.0.0.1", server.listen_address().port, &options));
|
|
bthread_t bth;
|
|
ASSERT_EQ(bthread_start_background(&bth, NULL, incr_thread, channels.back()), 0);
|
|
bths.push_back(bth);
|
|
}
|
|
|
|
for (int i = 0; i < N; ++i) {
|
|
bthread_join(bths[i], NULL);
|
|
delete channels[i];
|
|
}
|
|
ASSERT_EQ(int_map["count"], 10 * 5000LL);
|
|
}
|
|
|
|
class MultiCommandHandler : public brpc::RedisCommandHandler {
|
|
public:
|
|
MultiCommandHandler() {}
|
|
|
|
brpc::RedisCommandHandlerResult Run(const std::vector<butil::StringPiece>& args,
|
|
brpc::RedisReply* output,
|
|
bool flush_batched) override {
|
|
output->SetStatus("OK");
|
|
return brpc::REDIS_CMD_CONTINUE;
|
|
}
|
|
|
|
RedisCommandHandler* NewTransactionHandler() override {
|
|
return new MultiTransactionHandler;
|
|
}
|
|
|
|
class MultiTransactionHandler : public brpc::RedisCommandHandler {
|
|
public:
|
|
brpc::RedisCommandHandlerResult Run(const std::vector<butil::StringPiece>& args,
|
|
brpc::RedisReply* output,
|
|
bool flush_batched) override {
|
|
if (args[0] == "multi") {
|
|
output->SetError("ERR duplicate multi");
|
|
return brpc::REDIS_CMD_CONTINUE;
|
|
}
|
|
if (args[0] != "exec") {
|
|
std::vector<std::string> comm;
|
|
for (int i = 0; i < (int)args.size(); ++i) {
|
|
comm.push_back(args[i].as_string());
|
|
}
|
|
_commands.push_back(comm);
|
|
output->SetStatus("QUEUED");
|
|
return brpc::REDIS_CMD_CONTINUE;
|
|
}
|
|
output->SetArray(_commands.size());
|
|
s_mutex.lock();
|
|
for (size_t i = 0; i < _commands.size(); ++i) {
|
|
if (_commands[i][0] == "incr") {
|
|
int64_t value;
|
|
value = ++int_map[_commands[i][1]];
|
|
(*output)[i].SetInteger(value);
|
|
} else {
|
|
(*output)[i].SetStatus("unknown command");
|
|
}
|
|
}
|
|
s_mutex.unlock();
|
|
return brpc::REDIS_CMD_HANDLED;
|
|
}
|
|
private:
|
|
std::vector<std::vector<std::string> > _commands;
|
|
};
|
|
};
|
|
|
|
TEST_F(RedisTest, server_command_continue) {
|
|
std::string password = GeneratePassword();
|
|
std::unique_ptr<brpc::policy::RedisAuthenticator> redis_auth_holder(
|
|
new brpc::policy::RedisAuthenticator(password));
|
|
RedisServiceImpl* rsimpl = new RedisServiceImpl(password);
|
|
std::unique_ptr<AuthCommandHandler> ah(new AuthCommandHandler(rsimpl));
|
|
std::unique_ptr<GetCommandHandler> gh(new GetCommandHandler(rsimpl));
|
|
std::unique_ptr<SetCommandHandler> sh(new SetCommandHandler(rsimpl));
|
|
std::unique_ptr<IncrCommandHandler> ih(new IncrCommandHandler(rsimpl));
|
|
std::unique_ptr<MultiCommandHandler> mh(new MultiCommandHandler);
|
|
brpc::Server server;
|
|
brpc::ServerOptions server_options;
|
|
rsimpl->AddCommandHandler("auth", ah.get());
|
|
rsimpl->AddCommandHandler("get", gh.get());
|
|
rsimpl->AddCommandHandler("set", sh.get());
|
|
rsimpl->AddCommandHandler("incr", ih.get());
|
|
rsimpl->AddCommandHandler("multi", mh.get());
|
|
server_options.redis_service = rsimpl;
|
|
brpc::PortRange pr(8081, 8900);
|
|
ASSERT_EQ(0, server.Start("127.0.0.1", pr, &server_options));
|
|
|
|
brpc::ChannelOptions options;
|
|
options.protocol = brpc::PROTOCOL_REDIS;
|
|
options.auth = redis_auth_holder.get();
|
|
brpc::Channel channel;
|
|
ASSERT_EQ(0, channel.Init("127.0.0.1", server.listen_address().port, &options));
|
|
{
|
|
brpc::RedisRequest request;
|
|
brpc::RedisResponse response;
|
|
brpc::Controller cntl;
|
|
ASSERT_TRUE(request.AddCommand("set hello world"));
|
|
ASSERT_TRUE(request.AddCommand("get hello"));
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(2, response.reply_size());
|
|
ASSERT_STREQ("world", response.reply(1).c_str());
|
|
}
|
|
{
|
|
brpc::RedisRequest request;
|
|
brpc::RedisResponse response;
|
|
brpc::Controller cntl;
|
|
ASSERT_TRUE(request.AddCommand("multi"));
|
|
ASSERT_TRUE(request.AddCommand("mUltI"));
|
|
int count = 10;
|
|
for (int i = 0; i < count; ++i) {
|
|
ASSERT_TRUE(request.AddCommand("incr hello 1"));
|
|
}
|
|
ASSERT_TRUE(request.AddCommand("exec"));
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_EQ(13, response.reply_size());
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(0).type());
|
|
ASSERT_STREQ("OK", response.reply(0).c_str());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_ERROR, response.reply(1).type());
|
|
for (int i = 2; i < count + 2; ++i) {
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(i).type());
|
|
ASSERT_STREQ("QUEUED", response.reply(i).c_str());
|
|
}
|
|
const brpc::RedisReply& m = response.reply(count + 2);
|
|
ASSERT_EQ(count, (int)m.size());
|
|
for (int i = 0; i < count; ++i) {
|
|
ASSERT_EQ(i+1, m[i].integer());
|
|
}
|
|
}
|
|
// After 'multi', normal requests should be successful
|
|
{
|
|
brpc::RedisRequest request;
|
|
brpc::RedisResponse response;
|
|
brpc::Controller cntl;
|
|
ASSERT_TRUE(request.AddCommand("get hello"));
|
|
ASSERT_TRUE(request.AddCommand("get hello2"));
|
|
ASSERT_TRUE(request.AddCommand("set key1 value1"));
|
|
ASSERT_TRUE(request.AddCommand("get key1"));
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_STREQ("world", response.reply(0).c_str());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_NIL, response.reply(1).type());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(2).type());
|
|
ASSERT_STREQ("OK", response.reply(2).c_str());
|
|
ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(3).type());
|
|
ASSERT_STREQ("value1", response.reply(3).c_str());
|
|
}
|
|
}
|
|
|
|
TEST_F(RedisTest, server_handle_pipeline) {
|
|
std::string password = GeneratePassword();
|
|
std::unique_ptr<brpc::policy::RedisAuthenticator> redis_auth_holder(
|
|
new brpc::policy::RedisAuthenticator(password));
|
|
RedisServiceImpl* rsimpl = new RedisServiceImpl(password);
|
|
std::unique_ptr<GetCommandHandler> getch(new GetCommandHandler(rsimpl, true));
|
|
std::unique_ptr<SetCommandHandler> setch(new SetCommandHandler(rsimpl, true));
|
|
std::unique_ptr<AuthCommandHandler> authch(new AuthCommandHandler(rsimpl));
|
|
std::unique_ptr<MultiCommandHandler> multich(new MultiCommandHandler);
|
|
brpc::Server server;
|
|
brpc::ServerOptions server_options;
|
|
rsimpl->AddCommandHandler("auth", authch.get());
|
|
rsimpl->AddCommandHandler("get", getch.get());
|
|
rsimpl->AddCommandHandler("set", setch.get());
|
|
rsimpl->AddCommandHandler("multi", multich.get());
|
|
server_options.redis_service = rsimpl;
|
|
brpc::PortRange pr(8081, 8900);
|
|
ASSERT_EQ(0, server.Start("127.0.0.1", pr, &server_options));
|
|
|
|
brpc::ChannelOptions options;
|
|
options.protocol = brpc::PROTOCOL_REDIS;
|
|
options.auth = redis_auth_holder.get();
|
|
brpc::Channel channel;
|
|
ASSERT_EQ(0, channel.Init("127.0.0.1", server.listen_address().port, &options));
|
|
|
|
brpc::RedisRequest request;
|
|
brpc::RedisResponse response;
|
|
brpc::Controller cntl;
|
|
ASSERT_TRUE(request.AddCommand("set key1 v1"));
|
|
ASSERT_TRUE(request.AddCommand("set key2 v2"));
|
|
ASSERT_TRUE(request.AddCommand("set key3 v3"));
|
|
ASSERT_TRUE(request.AddCommand("get hello"));
|
|
ASSERT_TRUE(request.AddCommand("get hello"));
|
|
ASSERT_TRUE(request.AddCommand("set key1 world"));
|
|
ASSERT_TRUE(request.AddCommand("set key2 world"));
|
|
ASSERT_TRUE(request.AddCommand("get key2"));
|
|
channel.CallMethod(NULL, &cntl, &request, &response, NULL);
|
|
ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
|
|
ASSERT_EQ(8, response.reply_size());
|
|
ASSERT_EQ(1, rsimpl->_batch_count);
|
|
ASSERT_TRUE(response.reply(7).is_string());
|
|
ASSERT_STREQ(response.reply(7).c_str(), "world");
|
|
}
|
|
|
|
TEST_F(RedisTest, memory_allocation_limits) {
|
|
int32_t original_limit = brpc::FLAGS_redis_max_allocation_size;
|
|
brpc::FLAGS_redis_max_allocation_size = 1024;
|
|
|
|
butil::Arena arena;
|
|
|
|
// Test redis_reply.cpp limits
|
|
{
|
|
// Test bulk string exceeding limit
|
|
butil::IOBuf buf;
|
|
std::string large_string = "*1\r\n$2000\r\n";
|
|
large_string.append(2000, 'a');
|
|
large_string.append("\r\n");
|
|
buf.append(large_string);
|
|
|
|
brpc::RedisReply reply(&arena);
|
|
brpc::ParseError err = reply.ConsumePartialIOBuf(buf);
|
|
ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
|
|
}
|
|
|
|
{
|
|
// Test array allocation exceeding limit
|
|
butil::IOBuf buf;
|
|
int32_t large_count = brpc::FLAGS_redis_max_allocation_size / sizeof(brpc::RedisReply) + 1;
|
|
std::string large_array = "*" + std::to_string(large_count) + "\r\n";
|
|
buf.append(large_array);
|
|
|
|
brpc::RedisReply reply(&arena);
|
|
brpc::ParseError err = reply.ConsumePartialIOBuf(buf);
|
|
ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
|
|
}
|
|
|
|
{
|
|
// Test large count
|
|
butil::IOBuf buf;
|
|
int64_t large_count = 9223372036854775807;
|
|
std::string large_array = "*" + std::to_string(large_count) + "\r\n";
|
|
buf.append(large_array);
|
|
|
|
brpc::RedisReply reply(&arena);
|
|
brpc::ParseError err = reply.ConsumePartialIOBuf(buf);
|
|
ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
|
|
}
|
|
|
|
// Test redis_command.cpp limits
|
|
{
|
|
// Test command string exceeding limit
|
|
brpc::RedisCommandParser parser;
|
|
butil::IOBuf buf;
|
|
std::string large_cmd = "*2\r\n$3\r\nget\r\n$2000\r\n";
|
|
large_cmd.append(2000, 'b');
|
|
large_cmd.append("\r\n");
|
|
buf.append(large_cmd);
|
|
|
|
std::vector<butil::StringPiece> args;
|
|
brpc::ParseError err = parser.Consume(buf, &args, &arena);
|
|
ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
|
|
}
|
|
|
|
{
|
|
// Test command array size exceeding limit
|
|
brpc::RedisCommandParser parser;
|
|
butil::IOBuf buf;
|
|
int32_t large_array_size = brpc::FLAGS_redis_max_allocation_size / sizeof(butil::StringPiece) + 1;
|
|
std::string large_array_cmd = "*" + std::to_string(large_array_size) + "\r\n";
|
|
buf.append(large_array_cmd);
|
|
|
|
std::vector<butil::StringPiece> args;
|
|
brpc::ParseError err = parser.Consume(buf, &args, &arena);
|
|
ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
|
|
}
|
|
|
|
{
|
|
// Test inline command exceeding limit before CRLF arrives.
|
|
brpc::RedisCommandParser parser;
|
|
butil::IOBuf buf;
|
|
std::string large_inline_cmd = "get ";
|
|
large_inline_cmd.append(brpc::FLAGS_redis_max_allocation_size + 1, 'k');
|
|
buf.append(large_inline_cmd);
|
|
|
|
std::vector<butil::StringPiece> args;
|
|
brpc::ParseError err = parser.Consume(buf, &args, &arena);
|
|
ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
|
|
}
|
|
|
|
{
|
|
// A command line exactly at the limit may have CRLF split across reads.
|
|
brpc::RedisCommandParser parser;
|
|
butil::IOBuf buf;
|
|
std::string boundary_inline_cmd = "get ";
|
|
boundary_inline_cmd.append(brpc::FLAGS_redis_max_allocation_size -
|
|
boundary_inline_cmd.size(), 'k');
|
|
boundary_inline_cmd.push_back('\r');
|
|
buf.append(boundary_inline_cmd);
|
|
|
|
std::vector<butil::StringPiece> args;
|
|
brpc::ParseError err = parser.Consume(buf, &args, &arena);
|
|
ASSERT_EQ(brpc::PARSE_ERROR_NOT_ENOUGH_DATA, err);
|
|
|
|
buf.push_back('\n');
|
|
err = parser.Consume(buf, &args, &arena);
|
|
ASSERT_EQ(brpc::PARSE_OK, err);
|
|
ASSERT_EQ(2, (int)args.size());
|
|
ASSERT_EQ("get", args[0].as_string());
|
|
ASSERT_EQ(brpc::FLAGS_redis_max_allocation_size - 4, (int)args[1].size());
|
|
}
|
|
|
|
{
|
|
// Test that only the current inline command line is limited and copied.
|
|
brpc::RedisCommandParser parser;
|
|
butil::IOBuf buf;
|
|
std::string pipelined_inline_cmd = "PING\r\nget ";
|
|
pipelined_inline_cmd.append(brpc::FLAGS_redis_max_allocation_size + 1, 'k');
|
|
buf.append(pipelined_inline_cmd);
|
|
|
|
std::vector<butil::StringPiece> args;
|
|
brpc::ParseError err = parser.Consume(buf, &args, &arena);
|
|
ASSERT_EQ(brpc::PARSE_OK, err);
|
|
ASSERT_EQ(1, (int)args.size());
|
|
ASSERT_EQ("ping", args[0].as_string());
|
|
ASSERT_EQ(pipelined_inline_cmd.size() - 6, buf.size());
|
|
}
|
|
|
|
{
|
|
// Test large command array work
|
|
int32_t original_limit_tmp = brpc::FLAGS_redis_max_allocation_size;
|
|
brpc::FLAGS_redis_max_allocation_size = 1024 * 1024;
|
|
brpc::RedisCommandParser parser;
|
|
butil::IOBuf buf;
|
|
int32_t large_array_size = brpc::FLAGS_redis_max_allocation_size / sizeof(butil::StringPiece);
|
|
std::string large_array_cmd = "*" + std::to_string(large_array_size) + "\r\n";
|
|
for(int i = 0; i < large_array_size; i++){
|
|
large_array_cmd.append("$1\r\n1\r\n");
|
|
}
|
|
buf.append(large_array_cmd);
|
|
|
|
std::vector<butil::StringPiece> args;
|
|
brpc::ParseError err = parser.Consume(buf, &args, &arena);
|
|
ASSERT_EQ(brpc::PARSE_OK, err);
|
|
brpc::FLAGS_redis_max_allocation_size = original_limit_tmp;
|
|
}
|
|
|
|
// Test valid cases within limits
|
|
{
|
|
// Test small bulk string should work
|
|
butil::IOBuf buf;
|
|
std::string small_string = "*1\r\n$10\r\nhelloworld\r\n";
|
|
buf.append(small_string);
|
|
|
|
brpc::RedisReply reply(&arena);
|
|
brpc::ParseError err = reply.ConsumePartialIOBuf(buf);
|
|
ASSERT_EQ(brpc::PARSE_OK, err);
|
|
ASSERT_TRUE(reply.is_array());
|
|
ASSERT_EQ(1, (int)reply.size());
|
|
ASSERT_STREQ("helloworld", reply[0].c_str());
|
|
}
|
|
|
|
{
|
|
// Test small command should work
|
|
brpc::RedisCommandParser parser;
|
|
butil::IOBuf buf;
|
|
std::string small_cmd = "*2\r\n$3\r\nget\r\n$5\r\nmykey\r\n";
|
|
buf.append(small_cmd);
|
|
|
|
std::vector<butil::StringPiece> args;
|
|
brpc::ParseError err = parser.Consume(buf, &args, &arena);
|
|
ASSERT_EQ(brpc::PARSE_OK, err);
|
|
ASSERT_EQ(2, (int)args.size());
|
|
ASSERT_EQ("get", args[0].as_string());
|
|
ASSERT_EQ("mykey", args[1].as_string());
|
|
}
|
|
|
|
brpc::FLAGS_redis_max_allocation_size = original_limit;
|
|
}
|
|
|
|
} //namespace
|