diff --git a/WFDnsClient.cc b/WFDnsClient.cc new file mode 100644 index 0000000..b74afca --- /dev/null +++ b/WFDnsClient.cc @@ -0,0 +1,308 @@ +/* + Copyright (c) 2021 Sogou, 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: Liu Kai (liukaidx@sogou-inc.com) +*/ + +#include +#include +#include +#include "URIParser.h" +#include "StringUtil.h" +#include "WFDnsClient.h" + +using namespace protocol; + +using DnsCtx = std::function; +using ComplexTask = WFComplexClientTask; + +// 定义一个DnsParams类,用于存储DNS查询参数 +class DnsParams +{ +public: + // 定义一个内部结构体dns_params,用于存储具体的DNS参数 + struct dns_params + { + std::vector uris; // 存储解析后的URI + std::vector search_list; // 搜索列表 + int ndots; // 域名中允许的最小点数 + int attempts; // 最大尝试次数 + bool rotate; // 是否轮询 + }; + +public: + // DnsParams类的构造函数 + DnsParams() + { + this->ref = new std::atomic(1); // 初始化引用计数为1 + this->params = new dns_params(); // 初始化参数结构体 + } + + // DnsParams类的拷贝构造函数 + DnsParams(const DnsParams& p) + { + this->ref = p.ref; // 拷贝引用计数指针 + this->params = p.params; // 拷贝参数结构体指针 + this->incref(); // 增加引用计数 + } + + // DnsParams类的赋值运算符 + DnsParams& operator=(const DnsParams& p) + { + if (this != &p) // 如果不是自赋值 + { + this->decref(); // 减少当前对象的引用计数 + this->ref = p.ref; // 拷贝引用计数指针 + this->params = p.params; // 拷贝参数结构体指针 + this->incref(); // 增加引用计数 + } + return *this; // 返回当前对象的引用 + } + + // DnsParams类的析构函数 + ~DnsParams() { this->decref(); } // 减少引用计数 + + // 获取const类型的参数指针 + const dns_params *get_params() const { return this->params; } + // 获取非const类型的参数指针 + dns_params *get_params() { return this->params; } + +private: + // 增加引用计数 + void incref() { (*this->ref)++; } + // 减少引用计数,并在计数为0时释放资源 + void decref() + { + if (--*this->ref == 0) + { + delete this->params; // 删除参数结构体 + delete this->ref; // 删除引用计数 + } + } + +private: + dns_params *params; // 指向参数结构体的指针 + std::atomic *ref; // 指向引用计数的指针 +}; + +// 定义DNS状态枚举 +enum +{ + DNS_STATUS_TRY_ORIGIN_DONE = 0, + DNS_STATUS_TRY_ORIGIN_FIRST = 1, + DNS_STATUS_TRY_ORIGIN_LAST = 2 +}; + +// 定义DnsStatus结构体,用于存储DNS查询状态 +struct DnsStatus +{ + std::string origin_name; // 原始域名 + std::string current_name; // 当前域名 + size_t next_server; // 下一个要尝试的服务器 + size_t last_server; // 上一个尝试的服务器 + size_t next_domain; // 下一个要尝试的搜索域 + int attempts_left; // 剩余尝试次数 + int try_origin_state; // 尝试原始域名的状态 +}; + +// 计算字符串中点的数量 +static int __get_ndots(const std::string& s) +{ + int ndots = 0; + for (size_t i = 0; i < s.size(); i++) + ndots += s[i] == '.'; // 统计点的数量 + return ndots; +} + +// 检查是否有下一个域名要尝试 +static bool __has_next_name(const DnsParams::dns_params *p, + struct DnsStatus *s) +{ + if (s->try_origin_state == DNS_STATUS_TRY_ORIGIN_FIRST) + { + s->current_name = s->origin_name; // 设置当前域名为原始域名 + s->try_origin_state = DNS_STATUS_TRY_ORIGIN_DONE; // 更新状态 + return true; + } + + if (s->next_domain < p->search_list.size()) // 如果还有搜索域要尝试 + { + s->current_name = s->origin_name; // 设置当前域名为原始域名 + s->current_name.push_back('.'); // 添加点 + s->current_name.append(p->search_list[s->next_domain]); // 添加搜索域 + + s->next_domain++; // 移动到下一个搜索域 + return true; + } + + if (s->try_origin_state == DNS_STATUS_TRY_ORIGIN_LAST) + { + s->current_name = s->origin_name; // 设置当前域名为原始域名 + s->try_origin_state = DNS_STATUS_TRY_ORIGIN_DONE; // 更新状态 + return true; + } + + return false; // 没有下一个域名要尝试 +} + +// DNS查询回调内部函数 +static void __callback_internal(WFDnsTask *task, const DnsParams& params, + struct DnsStatus& s) +{ + ComplexTask *ctask = static_cast(task); // 转换任务类型 + int state = task->get_state(); // 获取任务状态 + DnsRequest *req = task->get_req(); // 获取DNS请求 + DnsResponse *resp = task->get_resp(); // 获取DNS响应 + const auto *p = params.get_params(); // 获取DNS参数 + int rcode = resp->get_rcode(); // 获取响应码 + + bool try_next_server = state != WFT_STATE_SUCCESS || // 如果状态不是成功 + rcode == DNS_RCODE_SERVER_FAILURE || // 或者响应码是服务器失败 + rcode == DNS_RCODE_NOT_IMPLEMENTED || // 或者响应码是未实现 + rcode == DNS_RCODE_REFUSED; // 或者响应码是拒绝 + bool try_next_name = rcode == DNS_RCODE_FORMAT_ERROR || // 如果响应码是格式错误 + rcode == DNS_RCODE_NAME_ERROR || // 或者响应码是名字错误 + resp->get_ancount() == 0; // 或者响应中没有答案 + + if (try_next_server) + { + if (s.last_server == s.next_server) // 如果已经是最后一个服务器 + s.attempts_left--; // 减少尝试次数 + if (s.attempts_left <= 0) // 如果尝试次数用完 + return; // 返回 + + s.next_server = (s.next_server + 1) % p->uris.size(); // 计算下一个服务器 + ctask->set_redirect(p->uris[s.next_server]); // 设置重定向 + return; // 返回 + } + + if (try_next_name && __has_next_name(p, &s)) // 如果需要尝试下一个名字 + { + req->set_question_name(s.current_name.c_str()); // 设置查询名字 + ctask->set_redirect(p->uris[s.next_server]); // 设置重定向 + return; // 返回 + } +} + +// WFDnsClient类的初始化函数,带一个URL参数 +int WFDnsClient::init(const std::string& url) +{ + return this->init(url, "", 1, 2, false); // 调用另一个初始化函数 +} + +// WFDnsClient类的初始化函数,用于配置DNS客户端 +int WFDnsClient::init(const std::string& url, const std::string& search_list, + int ndots, int attempts, bool rotate) +{ + std::vector hosts; // 用于存储分割后的主机名 + std::vector uris; // 用于存储解析后的URI + std::string host; // 单个主机名字符串 + ParsedURI uri; // 用于存储解析后的单个URI + + this->id = 0; // 初始化客户端ID为0 + hosts = StringUtil::split_filter_empty(url, ','); // 根据逗号分割URL字符串,获取主机名列表 + + // 遍历主机名列表,对每个主机名进行处理 + for (size_t i = 0; i < hosts.size(); i++) + { + host = hosts[i]; // 获取当前主机名 + // 检查主机名是否以"dns://"或"dnss://"开头,如果不是,则添加"dns://"前缀 + if (strncasecmp(host.c_str(), "dns://", 6) != 0 && + strncasecmp(host.c_str(), "dnss://", 7) != 0) + { + host = "dns://" + host; + } + + // 使用URIParser解析当前主机名,如果解析失败则返回错误码-1 + if (URIParser::parse(host, uri) != 0) + return -1; + + // 将解析后的URI添加到uris列表中 + uris.emplace_back(std::move(uri)); + } + + // 如果uris列表为空,或者ndots小于0,或者attempts小于1,则设置errno为EINVAL并返回错误码-1 + if (uris.empty() || ndots < 0 || attempts < 1) + { + errno = EINVAL; + return -1; + } + + // 创建一个新的DnsParams对象来存储DNS参数 + this->params = new DnsParams; + DnsParams::dns_params *q = ((DnsParams *)this->params)->get_params(); // 获取DNS参数的指针 + q->uris = std::move(uris); // 将解析后的URI列表移动到DNS参数中 + q->search_list = StringUtil::split_filter_empty(search_list, ','); // 根据逗号分割search_list字符串,获取搜索域列表 + // 设置ndots的值,如果大于15则设置为15 + q->ndots = ndots > 15 ? 15 : ndots; + // 设置attempts的值,如果大于5则设置为5 + q->attempts = attempts > 5 ? 5 : attempts; + q->rotate = rotate; // 设置是否轮询 + + return 0; // 初始化成功,返回0 +} + +// WFDnsClient类的析构函数,用于释放资源 +void WFDnsClient::deinit() +{ + delete (DnsParams *)this->params; // 删除分配的DnsParams对象 + this->params = NULL; // 将params指针设置为NULL +} + +// WFDnsClient类的方法,用于创建一个新的DNS任务 +WFDnsTask *WFDnsClient::create_dns_task(const std::string& name, + dns_callback_t callback) +{ + DnsParams::dns_params *p = ((DnsParams *)this->params)->get_params(); // 获取DNS参数 + struct DnsStatus status; // 创建DNS状态结构体 + size_t next_server; // 下一个要尝试的服务器索引 + WFDnsTask *task; // 指向新创建的DNS任务的指针 + DnsRequest *req; // 指向DNS请求的指针 + + // 如果启用轮询,则计算下一个服务器索引,否则使用第一个服务器 + next_server = p->rotate ? this->id++ % p->uris.size() : 0; + + status.origin_name = name; // 设置原始域名 + status.next_domain = 0; // 设置下一个要尝试的搜索域索引 + status.attempts_left = p->attempts; // 设置剩余尝试次数 + status.try_origin_state = DNS_STATUS_TRY_ORIGIN_FIRST; // 设置尝试原始域名的状态 + + // 如果域名以点结尾,则跳过搜索域 + if (!name.empty() && name.back() == '.') + status.next_domain = p->search_list.size(); + // 如果域名中的点数少于ndots,则设置尝试原始域名的状态为最后尝试 + else if (__get_ndots(name) < p->ndots) + status.try_origin_state = DNS_STATUS_TRY_ORIGIN_LAST; + + // 检查是否有下一个域名要尝试,并更新状态 + __has_next_name(p, &status); + + // 创建一个新的DNS任务,使用下一个服务器的URI和提供的回调函数 + task = WFTaskFactory::create_dns_task(p->uris[next_server], 0, std::move(callback)); + status.next_server = next_server; // 设置当前服务器索引 + status.last_server = (next_server + p->uris.size() - 1) % p->uris.size(); // 设置最后一个服务器索引 + + req = task->get_req(); // 获取DNS请求对象 + req->set_question(status.current_name.c_str(), DNS_TYPE_A, DNS_CLASS_IN); // 设置DNS请求的问题部分 + req->set_rd(1); // 设置DNS请求的递归标志 + + ComplexTask *ctask = static_cast(task); // 将任务转换为ComplexTask类型 + // 设置任务的上下文回调,当DNS任务完成时,将调用__callback_internal函数 + *ctask->get_mutable_ctx() = std::bind(__callback_internal, + std::placeholders::_1, + *(DnsParams *)params, status); + + return task; // 返回新创建的DNS任务 +}