This is an automated email from the ASF dual-hosted git repository.

wwbmmm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/brpc.git


The following commit(s) were added to refs/heads/master by this push:
     new cffbd243 keep session info in RedisConnContext (#2902)
cffbd243 is described below

commit cffbd243e5100d4b5d69bfd1b12473e7abef3993
Author: Tanghui Lin <xmutang...@gmail.com>
AuthorDate: Thu Mar 6 19:25:52 2025 +0800

    keep session info in RedisConnContext (#2902)
    
    * keep session info in RedisConnContext
    
    * fix comment
    
    * add uint test for redis ctx
    
    * fix uinttest
---
 example/BUILD.bazel                |  11 +++
 example/redis_c++/redis_server.cpp | 114 ++++++++++++++++++++++++++----
 src/brpc/policy/redis_protocol.cpp |  34 +--------
 src/brpc/redis.cpp                 |  17 +++++
 src/brpc/redis.h                   |  46 ++++++++++++-
 test/BUILD.bazel                   |   1 +
 test/brpc_redis_unittest.cpp       | 137 ++++++++++++++++++++++++++++++++++---
 7 files changed, 305 insertions(+), 55 deletions(-)

diff --git a/example/BUILD.bazel b/example/BUILD.bazel
index b38cd5a1..df2722a4 100644
--- a/example/BUILD.bazel
+++ b/example/BUILD.bazel
@@ -123,3 +123,14 @@ cc_binary(
         "//:brpc",
     ],
 )
+
+cc_binary(
+    name = "redis_c++_server",
+    srcs = [
+        "redis_c++/redis_server.cpp",
+    ],
+    copts = COPTS,
+    deps = [
+        "//:brpc",
+    ],
+)
\ No newline at end of file
diff --git a/example/redis_c++/redis_server.cpp 
b/example/redis_c++/redis_server.cpp
index 6ebc3853..a53ee26e 100644
--- a/example/redis_c++/redis_server.cpp
+++ b/example/redis_c++/redis_server.cpp
@@ -31,21 +31,54 @@
 
 DEFINE_int32(port, 6379, "TCP Port of this server");
 
+class AuthSession : public brpc::Destroyable {
+public:
+    explicit AuthSession(const std::string& user_name, const std::string& 
password)
+        : _user_name(user_name), _password(password) {}    
+
+    void Destroy() override {
+        delete this;
+    }    
+
+    const std::string _user_name;
+    const std::string _password;
+};
+
 class RedisServiceImpl : public brpc::RedisService {
 public:
-    bool Set(const std::string& key, const std::string& value) {
+    RedisServiceImpl() {
+        _user_password["db1"] = "123456";
+        _user_password["db2"] = "123456";
+        _db_map["db1"].resize(kHashSlotNum);
+        _db_map["db2"].resize(kHashSlotNum);
+    }
+
+    bool Set(const std::string& db_name, const std::string& key, const 
std::string& value) {
         int slot = butil::crc32c::Value(key.c_str(), key.size()) % 
kHashSlotNum;
         _mutex[slot].lock();
-        _map[slot][key] = value;
+        auto& kv = _db_map[db_name];
+        kv[slot][key] = value;
         _mutex[slot].unlock();
         return true;
     }
 
-    bool Get(const std::string& key, std::string* value) {
+    bool Auth(const std::string& db_name, const std::string& password) {
+        if (_user_password.find(db_name) == _user_password.end()) {
+            return false;
+        } else {
+            if (_user_password[db_name] != password) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    bool Get(const std::string& db_name, const std::string& key, std::string* 
value) {
         int slot = butil::crc32c::Value(key.c_str(), key.size()) % 
kHashSlotNum;
         _mutex[slot].lock();
-        auto it = _map[slot].find(key);
-        if (it == _map[slot].end()) {
+        auto& kv = _db_map[db_name];
+        auto it = kv[slot].find(key);
+        if (it == kv[slot].end()) {
             _mutex[slot].unlock();
             return false;
         }
@@ -56,7 +89,9 @@ public:
 
 private:
     const static int kHashSlotNum = 32;
-    std::unordered_map<std::string, std::string> _map[kHashSlotNum];
+    typedef std::unordered_map<std::string, std::string> KVStore;
+    std::unordered_map<std::string, std::vector<KVStore>> _db_map;
+    std::unordered_map<std::string, std::string> _user_password;
     butil::Mutex _mutex[kHashSlotNum];
 };
 
@@ -65,16 +100,27 @@ public:
     explicit GetCommandHandler(RedisServiceImpl* rsimpl)
         : _rsimpl(rsimpl) {}
 
-    brpc::RedisCommandHandlerResult Run(const std::vector<butil::StringPiece>& 
args,
+    brpc::RedisCommandHandlerResult Run(brpc::RedisConnContext* ctx, 
+                                        const std::vector<butil::StringPiece>& 
args,
                                         brpc::RedisReply* output,
                                         bool /*flush_batched*/) override {
+
+        AuthSession* session = static_cast<AuthSession*>(ctx->get_session());
+        if (session == nullptr) {
+            output->FormatError("No auth session");
+            return brpc::REDIS_CMD_HANDLED;
+        }
+        if (session->_user_name.empty()) {
+            output->FormatError("No user name");
+            return brpc::REDIS_CMD_HANDLED;
+        }
         if (args.size() != 2ul) {
             output->FormatError("Expect 1 arg for 'get', actually %lu", 
args.size()-1);
             return brpc::REDIS_CMD_HANDLED;
         }
         const std::string key(args[1].data(), args[1].size());
         std::string value;
-        if (_rsimpl->Get(key, &value)) {
+        if (_rsimpl->Get(session->_user_name, key, &value)) {
             output->SetString(value);
         } else {
             output->SetNullString();
@@ -91,22 +137,64 @@ public:
     explicit SetCommandHandler(RedisServiceImpl* rsimpl)
         : _rsimpl(rsimpl) {}
 
-    brpc::RedisCommandHandlerResult Run(const std::vector<butil::StringPiece>& 
args,
+    brpc::RedisCommandHandlerResult Run(brpc::RedisConnContext* ctx, 
+                                        const std::vector<butil::StringPiece>& 
args,
                                         brpc::RedisReply* output,
                                         bool /*flush_batched*/) override {
+        AuthSession* session = static_cast<AuthSession*>(ctx->get_session());
+        if (session == nullptr) {
+            output->FormatError("No auth session");
+            return brpc::REDIS_CMD_HANDLED;
+        }
+        if (session->_user_name.empty()) {
+            output->FormatError("No user name");
+            return brpc::REDIS_CMD_HANDLED;
+        }                                            
         if (args.size() != 3ul) {
             output->FormatError("Expect 2 args for 'set', actually %lu", 
args.size()-1);
             return brpc::REDIS_CMD_HANDLED;
         }
         const std::string key(args[1].data(), args[1].size());
         const std::string value(args[2].data(), args[2].size());
-        _rsimpl->Set(key, value);
+        _rsimpl->Set(session->_user_name, key, value);
         output->SetStatus("OK");
         return brpc::REDIS_CMD_HANDLED;
        }
 
 private:
-       RedisServiceImpl* _rsimpl;
+    RedisServiceImpl* _rsimpl;
+};
+
+
+
+class AuthCommandHandler : public brpc::RedisCommandHandler {
+public:
+    explicit AuthCommandHandler(RedisServiceImpl* rsimpl)
+        : _rsimpl(rsimpl) {}
+    brpc::RedisCommandHandlerResult Run(brpc::RedisConnContext* ctx, 
+                                        const std::vector<butil::StringPiece>& 
args,
+                                        brpc::RedisReply* output,
+                                        bool /*flush_batched*/) override {
+        if (args.size() != 3ul) {
+            output->FormatError("Expect 2 args for 'auth', actually %lu", 
args.size()-1);
+            return brpc::REDIS_CMD_HANDLED;
+        }
+        
+        const std::string db_name(args[1].data(), args[1].size());
+        const std::string password(args[2].data(), args[2].size());
+        
+        if (_rsimpl->Auth(db_name, password)) {
+            output->SetStatus("OK");
+            auto auth_session = new AuthSession(db_name, password);
+            ctx->reset_session(auth_session);
+        } else {
+            output->FormatError("Invalid password for database '%s'", 
db_name.c_str());
+        }
+        return brpc::REDIS_CMD_HANDLED;
+    }
+
+private:
+    RedisServiceImpl* _rsimpl;
 };
 
 int main(int argc, char* argv[]) {
@@ -114,9 +202,11 @@ int main(int argc, char* argv[]) {
     RedisServiceImpl *rsimpl = new RedisServiceImpl;
     auto get_handler =std::unique_ptr<GetCommandHandler>(new 
GetCommandHandler(rsimpl));
     auto set_handler =std::unique_ptr<SetCommandHandler>( new 
SetCommandHandler(rsimpl));
+    auto auth_handler = std::unique_ptr<AuthCommandHandler>(new 
AuthCommandHandler(rsimpl));
     rsimpl->AddCommandHandler("get", get_handler.get());
     rsimpl->AddCommandHandler("set", set_handler.get());
-
+    rsimpl->AddCommandHandler("auth", auth_handler.get());
+    
     brpc::Server server;
     brpc::ServerOptions server_options;
     server_options.redis_service = rsimpl;
diff --git a/src/brpc/policy/redis_protocol.cpp 
b/src/brpc/policy/redis_protocol.cpp
index 94524e8b..f8acf49d 100644
--- a/src/brpc/policy/redis_protocol.cpp
+++ b/src/brpc/policy/redis_protocol.cpp
@@ -54,27 +54,6 @@ struct InputResponse : public InputMessageBase {
     }
 };
 
-// This class is as parsing_context in socket.
-class RedisConnContext : public Destroyable  {
-public:
-    explicit RedisConnContext(const RedisService* rs)
-        : redis_service(rs)
-        , batched_size(0) {}
-
-    ~RedisConnContext();
-    // @Destroyable
-    void Destroy() override;
-
-    const RedisService* redis_service;
-    // If user starts a transaction, transaction_handler indicates the
-    // handler pointer that runs the transaction command.
-    std::unique_ptr<RedisCommandHandler> transaction_handler;
-    // >0 if command handler is run in batched mode.
-    int batched_size;
-
-    RedisCommandParser parser;
-    butil::Arena arena;
-};
 
 int ConsumeCommand(RedisConnContext* ctx,
                    const std::vector<butil::StringPiece>& args,
@@ -83,7 +62,7 @@ int ConsumeCommand(RedisConnContext* ctx,
     RedisReply output(&ctx->arena);
     RedisCommandHandlerResult result = REDIS_CMD_HANDLED;
     if (ctx->transaction_handler) {
-        result = ctx->transaction_handler->Run(args, &output, flush_batched);
+        result = ctx->transaction_handler->Run(ctx, args, &output, 
flush_batched);
         if (result == REDIS_CMD_HANDLED) {
             ctx->transaction_handler.reset(NULL);
         } else if (result == REDIS_CMD_BATCHED) {
@@ -97,7 +76,7 @@ int ConsumeCommand(RedisConnContext* ctx,
             snprintf(buf, sizeof(buf), "ERR unknown command `%s`", 
args[0].as_string().c_str());
             output.SetError(buf);
         } else {
-            result = ch->Run(args, &output, flush_batched);
+            result = ch->Run(ctx, args, &output, flush_batched);
             if (result == REDIS_CMD_CONTINUE) {
                 if (ctx->batched_size != 0) {
                     LOG(ERROR) << "CONTINUE should not be returned in a 
batched process.";
@@ -134,15 +113,6 @@ int ConsumeCommand(RedisConnContext* ctx,
     return 0;
 }
 
-// ========== impl of RedisConnContext ==========
-
-RedisConnContext::~RedisConnContext() { }
-
-void RedisConnContext::Destroy() {
-    delete this;
-}
-
-// ========== impl of RedisConnContext ==========
 
 ParseResult ParseRedisMessage(butil::IOBuf* source, Socket* socket,
                               bool read_eof, const void* arg) {
diff --git a/src/brpc/redis.cpp b/src/brpc/redis.cpp
index 777e1999..f8870ae5 100644
--- a/src/brpc/redis.cpp
+++ b/src/brpc/redis.cpp
@@ -370,4 +370,21 @@ RedisCommandHandler* 
RedisCommandHandler::NewTransactionHandler() {
     return NULL;
 }
 
+// ========== impl of RedisConnContext ==========
+RedisConnContext::~RedisConnContext() { }
+
+void RedisConnContext::Destroy() {
+    if (session) {
+        session->Destroy();
+    }
+    delete this;
+}
+
+void RedisConnContext::reset_session(Destroyable* s){
+    if (session) {
+        session->Destroy();
+    }
+    session = s;
+}
+
 } // namespace brpc
diff --git a/src/brpc/redis.h b/src/brpc/redis.h
index c6b0ea21..50064519 100644
--- a/src/brpc/redis.h
+++ b/src/brpc/redis.h
@@ -21,8 +21,10 @@
 
 #include <unordered_map>
 
+#include "brpc/destroyable.h"
 #include "brpc/nonreflectable_message.h"
 #include "brpc/parse_result.h"
+#include "brpc/redis_command.h"
 #include "brpc/pb_compat.h"
 #include "brpc/redis_reply.h"
 #include "butil/arena.h"
@@ -210,6 +212,39 @@ enum RedisCommandHandlerResult {
     REDIS_CMD_BATCHED = 2,
 };
 
+class RedisCommandParser;
+
+// This class is as parsing_context in socket.
+class RedisConnContext : public Destroyable  {
+public:
+    explicit RedisConnContext(const RedisService* rs)
+        : redis_service(rs)
+        , batched_size(0)
+        , session(nullptr) {}
+
+    ~RedisConnContext();
+    // @Destroyable
+    void Destroy() override;
+    void reset_session(Destroyable* s);
+
+    Destroyable* get_session() { return session; }
+
+    const RedisService* redis_service;
+    // If user starts a transaction, transaction_handler indicates the
+    // handler pointer that runs the transaction command.
+    std::unique_ptr<RedisCommandHandler> transaction_handler;
+    // >0 if command handler is run in batched mode.
+    int batched_size;
+
+    RedisCommandParser parser;
+    butil::Arena arena;
+
+private:
+    // If user is authenticated, session is set.
+    // Keep auth session info in RedisConnContext to distinguish diffrent 
users( or diffrent db).
+    Destroyable* session;
+};
+
 // The Command handler for a redis request. User should impletement Run().
 class RedisCommandHandler {
 public:
@@ -235,8 +270,15 @@ public:
     // it returns REDIS_CMD_HANDLED. Read the comment below.
     virtual RedisCommandHandlerResult Run(const 
std::vector<butil::StringPiece>& args,
                                           brpc::RedisReply* output,
-                                          bool flush_batched) = 0;
-
+                                          bool flush_batched) {
+        return REDIS_CMD_HANDLED;                                    
+    };
+    virtual RedisCommandHandlerResult Run(RedisConnContext* ctx, 
+                                          const 
std::vector<butil::StringPiece>& args,
+                                          brpc::RedisReply* output,
+                                          bool flush_batched) {
+        return Run(args, output, flush_batched);
+    }
     // The Run() returns CONTINUE for "multi", which makes brpc call this 
method to
     // create a transaction_handler to process following commands until 
transaction_handler
     // returns OK. For example, for command "multi; set k1 v1; set k2 v2; set 
k3 v3;
diff --git a/test/BUILD.bazel b/test/BUILD.bazel
index 9817b45f..1e0ef966 100644
--- a/test/BUILD.bazel
+++ b/test/BUILD.bazel
@@ -192,6 +192,7 @@ cc_test(
     ],
 )
 
+
 cc_test(
     name = "bvar_test",
     srcs = glob(
diff --git a/test/brpc_redis_unittest.cpp b/test/brpc_redis_unittest.cpp
index 573ab2ed..017d5c7e 100644
--- a/test/brpc_redis_unittest.cpp
+++ b/test/brpc_redis_unittest.cpp
@@ -811,10 +811,13 @@ 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()
-        : _batch_count(0) {}
+        : _batch_count(0)
+        , _user("user1")
+        , _password("password1") {}
 
     brpc::RedisCommandHandlerResult OnBatched(const 
std::vector<butil::StringPiece>& args,
                    brpc::RedisReply* output, bool flush_batched) {
@@ -864,8 +867,52 @@ public:
 
     std::vector<std::vector<std::string> > _batched_command;
     int _batch_count;
+    std::string _user;
+    std::string _password;
+};
+
+
+class AuthSession : public brpc::Destroyable {
+public:
+    explicit AuthSession(const std::string& user_name, const std::string& 
password)
+        : _user_name(user_name), _password(password) {}  
+
+    void Destroy() override {
+        delete this;
+    }   
+
+    const std::string _user_name;
+    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() < 2) {
+            output->SetError("ERR wrong number of arguments for 'AUTH' 
command");
+            return brpc::REDIS_CMD_HANDLED;
+        }
+        const std::string user(args[1].data(), args[1].size());
+        const std::string password(args[2].data(), args[2].size());
+        if (_rs->_user != user || _rs->_password != password) {
+            output->SetError("ERR invalid username/password");
+            return brpc::REDIS_CMD_HANDLED;
+        }
+        auto auth_session = new AuthSession(user, password);
+        ctx->reset_session(auth_session);
+        output->SetStatus("OK");
+        return brpc::REDIS_CMD_HANDLED;
+    }
+
+private:
+    RedisServiceImpl* _rs;
+};
 
 class SetCommandHandler : public brpc::RedisCommandHandler {
 public:
@@ -873,9 +920,19 @@ public:
         : _rs(rs)
         , _batch_process(batch_process) {}
 
-    brpc::RedisCommandHandlerResult Run(const std::vector<butil::StringPiece>& 
args,
+    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) || 
(session->_user_name != _rs->_user)) {
+            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;
@@ -898,15 +955,26 @@ private:
     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(const std::vector<butil::StringPiece>& 
args,
+    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) || 
(session->_user_name != _rs->_user)) {
+               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;
@@ -935,11 +1003,22 @@ private:
 
 class IncrCommandHandler : public brpc::RedisCommandHandler {
 public:
-    IncrCommandHandler() {}
+    IncrCommandHandler(RedisServiceImpl* rs)
+        : _rs(rs) {}
 
-    brpc::RedisCommandHandlerResult Run(const std::vector<butil::StringPiece>& 
args,
+    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) || 
(session->_user_name != _rs->_user)) {
+            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;
@@ -951,6 +1030,9 @@ public:
         output->SetInteger(value);
         return brpc::REDIS_CMD_HANDLED;
     }
+
+private:
+    RedisServiceImpl* _rs;
 };
 
 TEST_F(RedisTest, server_sanity) {
@@ -959,10 +1041,12 @@ TEST_F(RedisTest, server_sanity) {
     RedisServiceImpl* rsimpl = new RedisServiceImpl;
     GetCommandHandler *gh = new GetCommandHandler(rsimpl);
     SetCommandHandler *sh = new SetCommandHandler(rsimpl);
-    IncrCommandHandler *ih = new IncrCommandHandler;
+    AuthCommandHandler *ah = new AuthCommandHandler(rsimpl);
+    IncrCommandHandler *ih = new IncrCommandHandler(rsimpl);
     rsimpl->AddCommandHandler("get", gh);
     rsimpl->AddCommandHandler("set", sh);
     rsimpl->AddCommandHandler("incr", ih);
+    rsimpl->AddCommandHandler("auth", ah);
     server_options.redis_service = rsimpl;
     brpc::PortRange pr(8081, 8900);
     ASSERT_EQ(0, server.Start("127.0.0.1", pr, &server_options));
@@ -975,6 +1059,15 @@ TEST_F(RedisTest, server_sanity) {
     brpc::RedisRequest request;
     brpc::RedisResponse response;
     brpc::Controller cntl;
+    ASSERT_TRUE(request.AddCommand("auth user1 password1"));
+    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_STREQ("OK", response.reply(0).c_str());
+    request.Clear();
+    response.Clear();
+    cntl.Reset();
     ASSERT_TRUE(request.AddCommand("get hello"));
     ASSERT_TRUE(request.AddCommand("get hello2"));
     ASSERT_TRUE(request.AddCommand("set key1 value1"));
@@ -1029,7 +1122,13 @@ TEST_F(RedisTest, server_sanity) {
 
 void* incr_thread(void* arg) {
     brpc::Channel* c = static_cast<brpc::Channel*>(arg);
-
+    // do auth
+    brpc::RedisRequest auth_req;
+    brpc::RedisResponse auth_resp;
+    brpc::Controller auth_cntl;
+    EXPECT_TRUE(auth_req.AddCommand("auth user1 password1"));
+    c->CallMethod(NULL, &auth_cntl, &auth_req, &auth_resp, NULL);
+    EXPECT_FALSE(auth_cntl.Failed()) << auth_cntl.ErrorText();
     for (int i = 0; i < 5000; ++i) {
         brpc::RedisRequest request;
         brpc::RedisResponse response;
@@ -1048,8 +1147,10 @@ TEST_F(RedisTest, server_concurrency) {
     brpc::Server server;
     brpc::ServerOptions server_options;
     RedisServiceImpl* rsimpl = new RedisServiceImpl;
-    IncrCommandHandler *ih = new IncrCommandHandler;
+    AuthCommandHandler *ah = new AuthCommandHandler(rsimpl);
+    IncrCommandHandler *ih = new IncrCommandHandler(rsimpl);
     rsimpl->AddCommandHandler("incr", ih);
+    rsimpl->AddCommandHandler("auth", ah);
     server_options.redis_service = rsimpl;
     brpc::PortRange pr(8081, 8900);
     ASSERT_EQ(0, server.Start("0.0.0.0", pr, &server_options));
@@ -1130,9 +1231,10 @@ TEST_F(RedisTest, server_command_continue) {
     brpc::Server server;
     brpc::ServerOptions server_options;
     RedisServiceImpl* rsimpl = new RedisServiceImpl;
+    rsimpl->AddCommandHandler("auth", new AuthCommandHandler(rsimpl));
     rsimpl->AddCommandHandler("get", new GetCommandHandler(rsimpl));
     rsimpl->AddCommandHandler("set", new SetCommandHandler(rsimpl));
-    rsimpl->AddCommandHandler("incr", new IncrCommandHandler);
+    rsimpl->AddCommandHandler("incr", new IncrCommandHandler(rsimpl));
     rsimpl->AddCommandHandler("multi", new MultiCommandHandler);
     server_options.redis_service = rsimpl;
     brpc::PortRange pr(8081, 8900);
@@ -1142,6 +1244,13 @@ TEST_F(RedisTest, server_command_continue) {
     options.protocol = brpc::PROTOCOL_REDIS;
     brpc::Channel channel;
     ASSERT_EQ(0, channel.Init("127.0.0.1", server.listen_address().port, 
&options));
+    // do auth
+    brpc::RedisRequest auth_req;
+    brpc::RedisResponse auth_resp;
+    brpc::Controller auth_cntl;
+    ASSERT_TRUE(auth_req.AddCommand("auth user1 password1"));
+    channel.CallMethod(NULL, &auth_cntl, &auth_req, &auth_resp, NULL);
+    ASSERT_FALSE(auth_cntl.Failed()) << auth_cntl.ErrorText();
 
     {
         brpc::RedisRequest request;
@@ -1207,6 +1316,8 @@ TEST_F(RedisTest, server_handle_pipeline) {
     RedisServiceImpl* rsimpl = new RedisServiceImpl;
     GetCommandHandler* getch = new GetCommandHandler(rsimpl, true);
     SetCommandHandler* setch = new SetCommandHandler(rsimpl, true);
+    AuthCommandHandler* authch = new AuthCommandHandler(rsimpl);
+    rsimpl->AddCommandHandler("auth", authch);
     rsimpl->AddCommandHandler("get", getch);
     rsimpl->AddCommandHandler("set", setch);
     rsimpl->AddCommandHandler("multi", new MultiCommandHandler);
@@ -1222,6 +1333,14 @@ TEST_F(RedisTest, server_handle_pipeline) {
     brpc::RedisRequest request;
     brpc::RedisResponse response;
     brpc::Controller cntl;
+    ASSERT_TRUE(request.AddCommand("auth user1 password1"));
+    channel.CallMethod(NULL, &cntl, &request, &response, NULL);
+    ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
+    ASSERT_EQ(1, response.reply_size());
+    ASSERT_STREQ("OK", response.reply(0).c_str());
+    request.Clear();
+    response.Clear();
+    cntl.Reset();
     ASSERT_TRUE(request.AddCommand("set key1 v1"));
     ASSERT_TRUE(request.AddCommand("set key2 v2"));
     ASSERT_TRUE(request.AddCommand("set key3 v3"));


---------------------------------------------------------------------
To unsubscribe, e-mail: dev-unsubscr...@brpc.apache.org
For additional commands, e-mail: dev-h...@brpc.apache.org

Reply via email to