From ec660592ed0be8eb9a0d77ba5b1519f4d39bd5aa Mon Sep 17 00:00:00 2001 From: wzlsuccess Date: Sat, 19 Jun 2021 11:13:16 +0800 Subject: [PATCH] update sasl for scram --- src/client/WFKafkaClient.cc | 52 +++++++++++++++----- src/factory/KafkaTaskImpl.cc | 49 ++++++++++--------- src/protocol/KafkaDataTypes.h | 91 +++++++++++++++++++++++++++++++++-- src/protocol/KafkaMessage.cc | 9 ++-- 4 files changed, 159 insertions(+), 42 deletions(-) diff --git a/src/client/WFKafkaClient.cc b/src/client/WFKafkaClient.cc index 2497f2da..7dc5adda 100644 --- a/src/client/WFKafkaClient.cc +++ b/src/client/WFKafkaClient.cc @@ -124,7 +124,8 @@ public: this->meta_list = new KafkaMetaList; this->broker_list = new KafkaBrokerList; this->lock_status = new KafkaLockStatus; - this->broker_map = new KafkaBrokerMap; + this->broker_map_by_id = new KafkaBrokerMap; + this->broker_map_by_uri = new KafkaBrokerMap; this->meta_map = new std::map; } @@ -137,7 +138,8 @@ public: delete this->cgroup; delete this->meta_list; delete this->broker_list; - delete this->broker_map; + delete this->broker_map_by_id; + delete this->broker_map_by_uri; delete this->lock_status; delete this->meta_map; } @@ -147,7 +149,8 @@ public: KafkaCgroup *cgroup; KafkaMetaList *meta_list; KafkaBrokerList *broker_list; - KafkaBrokerMap *broker_map; + KafkaBrokerMap *broker_map_by_id; + KafkaBrokerMap *broker_map_by_uri; KafkaLockStatus *lock_status; std::map *meta_map; @@ -253,7 +256,8 @@ public: this->cgroup = *client->member->cgroup; this->client_meta_list = *client->member->meta_list; this->client_broker_list = *client->member->broker_list; - this->client_broker_map = *client->member->broker_map; + this->client_broker_map_by_id = *client->member->broker_map_by_id; + this->client_broker_map_by_uri = *client->member->broker_map_by_uri; this->query = query; if (!client->member->broker_hosts->empty()) @@ -323,7 +327,13 @@ private: inline KafkaBroker *get_broker(int node_id) { - return this->client_broker_map.find_item(node_id); + return this->client_broker_map_by_id.find_item(node_id); + } + + inline KafkaBroker *find_broker(const KafkaBrokerMap& map, + const KafkaBroker& broker) + { + return map.find_item(broker); } int get_node_id(const KafkaToppar *toppar); @@ -337,7 +347,8 @@ private: KafkaLockStatus lock_status; KafkaMetaList client_meta_list; KafkaBrokerList client_broker_list; - KafkaBrokerMap client_broker_map; + KafkaBrokerMap client_broker_map_by_id; + KafkaBrokerMap client_broker_map_by_uri; KafkaCgroup cgroup; std::map toppar_list_map; std::string url; @@ -594,7 +605,7 @@ void ComplexKafkaTask::kafka_process_broker_api(ComplexKafkaTask *t, __WFKafkaTa sizeof(kafka_api_version_t) * api_cnt); broker->get_raw_ptr()->status = KAFKA_BROKER_INITED; - t->client_broker_map.add_item(*broker); + t->client_broker_map_by_id.add_item(*broker, broker->get_node_id()); } else { @@ -628,9 +639,18 @@ void ComplexKafkaTask::kafka_process_broker_api(ComplexKafkaTask *t, __WFKafkaTa task->get_resp()->get_broker_list()->rewind(); while ((broker = task->get_resp()->get_broker_list()->get_next()) != NULL) { - exist = t->get_broker(broker->get_node_id()); + exist = t->find_broker(t->client_broker_map_by_uri, *broker); if (exist && exist->get_raw_ptr()->status == KAFKA_BROKER_INITED) + { + if (exist->get_node_id() != broker->get_node_id()) + { + broker->copy_from(*exist); + broker->get_raw_ptr()->status = KAFKA_BROKER_INITED; + t->client_broker_map_by_id.add_item(*broker, + broker->get_node_id()); + } continue; + } else if (exist && exist->get_raw_ptr()->status == KAFKA_BROKER_DOING) { char name[64]; @@ -641,7 +661,9 @@ void ComplexKafkaTask::kafka_process_broker_api(ComplexKafkaTask *t, __WFKafkaTa continue; } - t->client_broker_map.add_item(*broker); + broker->get_raw_ptr()->status = KAFKA_BROKER_DOING; + t->client_broker_map_by_uri.add_item(*broker); + t->client_broker_map_by_id.add_item(*broker, broker->get_node_id()); auto cb = std::bind(&ComplexKafkaTask::kafka_broker_api_callback, t, std::placeholders::_1); @@ -670,7 +692,6 @@ void ComplexKafkaTask::kafka_process_broker_api(ComplexKafkaTask *t, __WFKafkaTa ntask->get_req()->set_broker(*broker); ntask->get_req()->set_api(Kafka_ApiVersions); ntask->user_data = broker->get_raw_ptr(); - broker->get_raw_ptr()->status = KAFKA_BROKER_DOING; KafkaComplexTask *ctask = static_cast(ntask); *ctask->get_mutable_ctx() = cb; series = Workflow::create_series_work(ntask, nullptr); @@ -914,6 +935,15 @@ void ComplexKafkaTask::generate_info() this->userinfo += this->config.get_sasl_password(); this->userinfo += ":"; this->userinfo += this->config.get_sasl_mechanisms(); + + this->meta_list.rewind(); + KafkaMeta *meta; + while ((meta = this->meta_list.get_next()) != NULL) + { + this->userinfo += "%"; + this->userinfo += meta->get_topic(); + } + this->userinfo += "@"; this->url = "kafka://" + this->userinfo + this->url.substr(this->url.find("kafka://") + 8); @@ -1040,7 +1070,7 @@ void ComplexKafkaTask::dispatch() else if ((this->api_type == Kafka_Fetch || this->api_type == Kafka_OffsetCommit) && (*this->lock_status.get_status() & KAFKA_CGROUP_INIT)) { - KafkaBroker *broker = this->client_broker_map.get_first_entry(); + KafkaBroker *broker = this->client_broker_map_by_id.get_first_entry(); if (!broker) { this->state = WFT_STATE_TASK_ERROR; diff --git a/src/factory/KafkaTaskImpl.cc b/src/factory/KafkaTaskImpl.cc index 6d7c8c47..626619cb 100644 --- a/src/factory/KafkaTaskImpl.cc +++ b/src/factory/KafkaTaskImpl.cc @@ -48,17 +48,17 @@ protected: virtual bool finish_once(); private: - struct __KafkaConnectionInfo + struct KafkaConnectionInfo { kafka_sasl_t sasl; std::string mechanisms; - __KafkaConnectionInfo() + KafkaConnectionInfo() { kafka_sasl_init(&this->sasl); } - ~__KafkaConnectionInfo() + ~KafkaConnectionInfo() { kafka_sasl_deinit(&this->sasl); } @@ -109,16 +109,24 @@ CommMessageOut *__ComplexKafkaTask::message_out() long long seqid = this->get_seq(); KafkaBroker *broker = this->get_req()->get_broker(); - if (!broker->get_api()) + if (seqid == 0) { if (!this->get_req()->get_config()->get_broker_version()) { - KafkaRequest *req = new KafkaRequest; + if (!broker->get_api()) + { + KafkaRequest *req = new KafkaRequest; - req->duplicate(*this->get_req()); - req->set_api(Kafka_ApiVersions); - is_user_request_ = false; - return req; + req->duplicate(*this->get_req()); + req->set_api(Kafka_ApiVersions); + + if (this->get_req()->get_api() != Kafka_ApiVersions) + is_user_request_ = false; + + return req; + } + else + seqid++; } else { @@ -139,17 +147,16 @@ CommMessageOut *__ComplexKafkaTask::message_out() this->error = WFT_ERR_KAFKA_VERSION_DISALLOWED; return NULL; } + seqid++; } } - const char *sasl_mechanisms = this->get_req()->get_config()->get_sasl_mechanisms(); - if (seqid < 2 && sasl_mechanisms) + if (seqid == 1 && !this->get_connection()->get_context()) { - WFConnection *conn = this->get_connection(); - __KafkaConnectionInfo *conn_info = (__KafkaConnectionInfo *)conn->get_context(); - if (!conn_info) + const char *sasl_mechanisms = this->get_req()->get_config()->get_sasl_mechanisms(); + if (sasl_mechanisms) { - conn_info = new __KafkaConnectionInfo; + KafkaConnectionInfo *conn_info = new KafkaConnectionInfo; if (!conn_info->init(sasl_mechanisms)) { this->state = WFT_STATE_TASK_ERROR; @@ -158,16 +165,14 @@ CommMessageOut *__ComplexKafkaTask::message_out() } auto&& deleter = [] (void *ctx) - { - __KafkaConnectionInfo *conn_info = (__KafkaConnectionInfo *)ctx; - delete conn_info; - }; - conn->set_context(conn_info, std::move(deleter)); - + { + KafkaConnectionInfo *conn_info = (KafkaConnectionInfo *)ctx; + delete conn_info; + }; + this->get_connection()->set_context(conn_info, std::move(deleter)); this->get_req()->set_sasl(&conn_info->sasl); KafkaRequest *req = new KafkaRequest; - req->duplicate(*this->get_req()); if (broker->get_features() & KAFKA_FEATURE_SASL_HANDSHAKE) diff --git a/src/protocol/KafkaDataTypes.h b/src/protocol/KafkaDataTypes.h index 8ade7af6..dd95971e 100644 --- a/src/protocol/KafkaDataTypes.h +++ b/src/protocol/KafkaDataTypes.h @@ -234,7 +234,56 @@ public: return *this; } - T *find_item(int id) + T *find_item(const T& v) const + { + rb_node **p = &this->t_map->rb_node; + T *t; + + while (*p) + { + t = rb_entry(*p, T, rb); + + if (v < *t) + p = &(*p)->rb_left; + else if (v > *t) + p = &(*p)->rb_right; + else + break; + } + + return *p ? t : NULL; + } + + void add_item(T& obj) + { + rb_node **p = &this->t_map->rb_node; + rb_node *parent = NULL; + T *t; + + while (*p) + { + parent = *p; + t = rb_entry(*p, T, rb); + + if (obj < *t) + p = &(*p)->rb_left; + else if (obj > *t) + p = &(*p)->rb_right; + else + break; + } + + if (*p == NULL) + { + T *nt = new T; + + *nt = obj; + rb_link_node(nt->get_rb(), parent, p); + rb_insert_color(nt->get_rb(), this->t_map); + } + } + + T *find_item(int id) const { rb_node **p = &this->t_map->rb_node; T *t; @@ -254,12 +303,11 @@ public: return *p ? t : NULL; } - void add_item(T& obj) + void add_item(T& obj, int id) { rb_node **p = &this->t_map->rb_node; rb_node *parent = NULL; T *t; - int id = obj.get_id(); while (*p) { @@ -978,6 +1026,41 @@ public: return *this; } + bool operator< (const KafkaBroker& broker) const + { + return this->get_uri() < broker.get_uri(); + } + + bool operator> (const KafkaBroker& broker) const + { + return this->get_uri() > broker.get_uri(); + } + + bool copy_from(const KafkaBroker& broker) + { + const struct sockaddr *addr; + socklen_t socklen; + broker.get_broker_addr(&addr, &socklen); + memcpy(&this->ptr->addr, addr, socklen); + this->ptr->addrlen = socklen; + this->ptr->to_addr = 1; + + size_t size = (broker.get_raw_ptr()->api_elements) * + sizeof(kafka_api_version_t); + void *p = malloc(size); + + if (!p) + return false; + + memcpy(p, broker.get_raw_ptr()->api, size); + free(this->ptr->api); + this->ptr->api = (kafka_api_version_t *)p; + this->ptr->api_elements = broker.get_raw_ptr()->api_elements; + this->ptr->features = broker.get_raw_ptr()->features; + + return true; + } + kafka_broker_t *get_raw_ptr() const { return this->ptr; } struct list_head *get_list() { return &this->list; } @@ -1011,7 +1094,7 @@ public: return is_equal(broker.ptr->host, broker.ptr->port); } - void get_broker_addr(const struct sockaddr **addr, socklen_t *socklen) + void get_broker_addr(const struct sockaddr **addr, socklen_t *socklen) const { if (this->ptr->addrlen) { diff --git a/src/protocol/KafkaMessage.cc b/src/protocol/KafkaMessage.cc index b07e3a8d..f37d631a 100644 --- a/src/protocol/KafkaMessage.cc +++ b/src/protocol/KafkaMessage.cc @@ -3787,12 +3787,12 @@ int KafkaResponse::handle_sasl_continue() { int ret; std::vector iovecs(2); - if (this->encode(iovecs.data(), 2) == 2) + if (this->encode(iovecs.data(), 2) > 0) { for (auto vec : iovecs) { ret = this->feedback((const char *)vec.iov_base, vec.iov_len); - if (ret != vec.iov_len) + if (ret != (int)vec.iov_len) { if (ret > 0) errno = EAGAIN; @@ -3810,8 +3810,7 @@ int KafkaResponse::append(const void *buf, size_t *size) if (ret == 1) { - int resp_ret = this->parse_response(); - if (resp_ret == 0) + if (this->parse_response() == 0) { if (this->api_type == Kafka_SaslHandshake) { @@ -3830,7 +3829,7 @@ int KafkaResponse::append(const void *buf, size_t *size) } } else - ret = resp_ret; + return -1; } return ret;