update sasl for scram

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<std::string, MetaStatus>;
 	}
 
@@ -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<std::string, MetaStatus> *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<int, KafkaTopparList> 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<KafkaComplexTask *>(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;

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)

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)
 		{

src/protocol/KafkaMessage.cc

@@ -3787,12 +3787,12 @@ int KafkaResponse::handle_sasl_continue()
 {
 	int ret;
 	std::vector<struct iovec> 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;