p9prynkfo
/
workflow
1
0
Fork 0
Code Issues Pull Requests 1 Packages Projects Releases Wiki Activity

ADD file via upload

Browse Source
liangtingyu
pm5lov8b3 3 months ago
parent e4b163747e
commit 5b5972a172
1 changed files with 308 additions and 0 deletions
Show Stats Download Patch File Download Diff File

308
WFDnsClient.cc
Unescape View File

@ -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 <string>
#include <vector>
#include <atomic>
#include "URIParser.h"
#include "StringUtil.h"
#include "WFDnsClient.h"
using namespace protocol;
using DnsCtx = std::function<void (WFDnsTask *)>;
using ComplexTask = WFComplexClientTask<DnsRequest, DnsResponse, DnsCtx>;
// 定义一个DnsParams类，用于存储DNS查询参数
class DnsParams
{
public:
// 定义一个内部结构体dns_params，用于存储具体的DNS参数
struct dns_params
{
std::vector<ParsedURI> uris; // 存储解析后的URI
std::vector<std::string> search_list; // 搜索列表
int ndots; // 域名中允许的最小点数
int attempts; // 最大尝试次数
bool rotate; // 是否轮询
};
public:
// DnsParams类的构造函数
DnsParams()
{
this->ref = new std::atomic<size_t>(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<size_t> *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<ComplexTask *>(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<std::string> hosts; // 用于存储分割后的主机名
std::vector<ParsedURI> 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<ComplexTask *>(task); // 将任务转换为ComplexTask类型
// 设置任务的上下文回调，当DNS任务完成时，将调用__callback_internal函数
*ctask->get_mutable_ctx() = std::bind(__callback_internal,
std::placeholders::_1,
*(DnsParams *)params, status);
return task; // 返回新创建的DNS任务
}
Write Preview
Loading…
Cancel
Save