Merge pull request #109 from wchengcheng/master

support PLAIN authentication of couchbase in memcache client

example/memcache_c++/client.cpp

@@ -20,11 +20,15 @@
 #include <butil/string_printf.h>
 #include <brpc/channel.h>
 #include <brpc/memcache.h>
+#include <brpc/policy/couchbase_authenticator.h>
 
 DEFINE_int32(thread_num, 10, "Number of threads to send requests");
 DEFINE_bool(use_bthread, false, "Use bthread to send requests");
+DEFINE_bool(use_couchbase, false, "Use couchbase.");
 DEFINE_string(connection_type, "", "Connection type. Available values: single, pooled, short");
 DEFINE_string(server, "0.0.0.0:11211", "IP Address of server");
+DEFINE_string(bucket_name, "", "Couchbase bucktet name");
+DEFINE_string(bucket_password, "", "Couchbase bucket password");
 DEFINE_string(load_balancer, "", "The algorithm for load balancing");
 DEFINE_int32(timeout_ms, 100, "RPC timeout in milliseconds");
 DEFINE_int32(max_retry, 3, "Max retries(not including the first RPC)"); 
@@ -109,6 +113,13 @@ int main(int argc, char* argv[]) {
     options.connection_type = FLAGS_connection_type;
     options.timeout_ms = FLAGS_timeout_ms/*milliseconds*/;
     options.max_retry = FLAGS_max_retry;
+    if (FLAGS_use_couchbase && !FLAGS_bucket_name.empty()) {
+        brpc::policy::CouchbaseAuthenticator* auth =
+            new brpc::policy::CouchbaseAuthenticator(FLAGS_bucket_name,
+                                                     FLAGS_bucket_password);
+        options.auth = auth;
+    }
+
     if (channel.Init(FLAGS_server.c_str(), FLAGS_load_balancer.c_str(), &options) != 0) {
         LOG(ERROR) << "Fail to initialize channel";
         return -1;
@@ -180,6 +191,9 @@ int main(int argc, char* argv[]) {
             bthread_join(tids[i], NULL);
         }
     }
+    if (options.auth) {
+        delete options.auth;
+    }
 
     return 0;
 }

src/brpc/policy/couchbase_authenticator.cpp

+// Copyright (c) 2017 Baidu, Inc.
+//
+// Licensed 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.
+//
+// Author(s): Chengcheng Wu <wuchengcheng@qiyi.com>
+
+#include "brpc/policy/couchbase_authenticator.h"
+
+#include "butil/base64.h"
+#include "butil/iobuf.h"
+#include "butil/string_printf.h"
+#include "butil/sys_byteorder.h"
+#include "brpc/policy/memcache_binary_header.h"
+
+namespace brpc {
+namespace policy {
+
+namespace {
+
+constexpr char kPlainAuthCommand[] = "PLAIN";
+constexpr char kPadding[1] = {'\0'};
+
+}  // namespace
+
+// To get the couchbase authentication protocol, see
+// https://developer.couchbase.com/documentation/server/3.x/developer/dev-guide-3.0/sasl.html
+int CouchbaseAuthenticator::GenerateCredential(std::string* auth_str) const {
+    const brpc::policy::MemcacheRequestHeader header = {
+        brpc::policy::MC_MAGIC_REQUEST, brpc::policy::MC_BINARY_SASL_AUTH,
+        butil::HostToNet16(sizeof(kPlainAuthCommand) - 1), 0, 0, 0,
+        butil::HostToNet32(sizeof(kPlainAuthCommand) + 1 +
+                           bucket_name_.length() * 2 + bucket_password_.length()),
+        0, 0};
+    auth_str->clear();
+    auth_str->append(reinterpret_cast<const char*>(&header), sizeof(header));
+    auth_str->append(kPlainAuthCommand, sizeof(kPlainAuthCommand) - 1);
+    auth_str->append(bucket_name_);
+    auth_str->append(kPadding, sizeof(kPadding));
+    auth_str->append(bucket_name_);
+    auth_str->append(kPadding, sizeof(kPadding));
+    auth_str->append(bucket_password_);
+    return 0;
+}
+
+}  // namespace policy
+}  // namespace brpc

src/brpc/policy/couchbase_authenticator.h

+// Copyright (c) 2017 Baidu, Inc.
+//
+// Licensed 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.
+//
+// Author(s): Chengcheng Wu <wuchengcheng@qiyi.com>
+
+#ifndef BRPC_POLICY_COUCHBASE_AUTHENTICATOR_H
+#define BRPC_POLICY_COUCHBASE_AUTHENTICATOR_H
+
+#include "brpc/authenticator.h"
+
+namespace brpc {
+namespace policy {
+
+// Request to couchbase for authentication.
+// Notice that authentication for couchbase in special SASLAuthProtocol.
+// Couchbase Server 2.2 provide CRAM-MD5 support for SASL authentication,
+// but Couchbase Server prior to 2.2 using PLAIN SASL authentication.
+class CouchbaseAuthenticator : public Authenticator {
+ public:
+  CouchbaseAuthenticator(const std::string& bucket_name,
+                         const std::string& bucket_password)
+      : bucket_name_(bucket_name), bucket_password_(bucket_password) {}
+
+  int GenerateCredential(std::string* auth_str) const;
+
+  int VerifyCredential(const std::string&, const butil::EndPoint&,
+                       brpc::AuthContext*) const {
+    return 0;
+  }
+
+ private:
+  const std::string bucket_name_;
+  const std::string bucket_password_;
+};
+
+}  // namespace policy
+}  // namespace brpc
+
+#endif  // BRPC_POLICY_COUCHBASE_AUTHENTICATOR_H

src/brpc/policy/memcache_binary_protocol.cpp

@@ -63,6 +63,7 @@ static void InitSupportedCommandMap() {
     butil::bit_array_set(supported_cmd_map, MC_BINARY_PREPEND);
     butil::bit_array_set(supported_cmd_map, MC_BINARY_STAT);
     butil::bit_array_set(supported_cmd_map, MC_BINARY_TOUCH);
+    butil::bit_array_set(supported_cmd_map, MC_BINARY_SASL_AUTH);
 }
 
 inline bool IsSupportedCommand(uint8_t command) {
@@ -125,14 +126,26 @@ ParseResult ParseMemcacheMessage(butil::IOBuf* source,
         msg->meta.append(&local_header, sizeof(local_header));
         source->pop_front(sizeof(*header));
         source->cutn(&msg->meta, total_body_length);
-        if (++msg->pi.count >= pi.count) {
-            CHECK_EQ(msg->pi.count, pi.count);
-            msg = static_cast<MostCommonMessage*>(socket->release_parsing_context());
-            msg->pi = pi;
-            return MakeMessage(msg);
-        } else {
+        if (header->command == MC_BINARY_SASL_AUTH) {
+            if (header->status != 0) {
+                LOG(ERROR) << "Failed to authenticate the couchbase bucket."
+                           << "All the following commands will result in auth failure.";
+                return MakeParseError(PARSE_ERROR_NO_RESOURCE, 
+                                      "Fail to authenticate with the couchbase bucket");
+            }
+            DestroyingPtr<MostCommonMessage> auth_msg(
+                 static_cast<MostCommonMessage*>(socket->release_parsing_context()));
             socket->GivebackPipelinedInfo(pi);
-        }
+        } else {
+            if (++msg->pi.count >= pi.count) {
+                CHECK_EQ(msg->pi.count, pi.count);
+                msg = static_cast<MostCommonMessage*>(socket->release_parsing_context());
+                msg->pi = pi;
+                return MakeMessage(msg);
+            } else {
+                socket->GivebackPipelinedInfo(pi);
+            }
+        }    
     }
 }
 
@@ -196,9 +209,16 @@ void PackMemcacheRequest(butil::IOBuf* buf,
                          SocketMessage**,
                          uint64_t /*correlation_id*/,
                          const google::protobuf::MethodDescriptor*,
-                         Controller*,
+                         Controller* cntl,
                          const butil::IOBuf& request,
-                         const Authenticator* /*auth*/) {
+                         const Authenticator* auth) {
+    if (auth) {
+        std::string auth_str;
+        if (auth->GenerateCredential(&auth_str) != 0) {
+            return cntl->SetFailed(EREQUEST, "Fail to generate credential");
+        }
+        buf->append(auth_str);
+    }
     buf->append(request);
 }