diff --git a/CMakeLists_Headers.txt b/CMakeLists_Headers.txt index e1233a5e..7997f65b 100644 --- a/CMakeLists_Headers.txt +++ b/CMakeLists_Headers.txt @@ -36,7 +36,7 @@ else () endif () set(INCLUDE_HEADERS - src/algorithm/DNSRoutine.h + src/algorithm/DnsRoutine.h src/algorithm/MapReduce.h src/algorithm/MapReduce.inl src/protocol/ProtocolMessage.h @@ -54,12 +54,15 @@ set(INCLUDE_HEADERS src/protocol/mysql_types.h src/protocol/mysql_byteorder.h src/protocol/SSLWrapper.h + src/protocol/dns_parser.h + src/protocol/DnsMessage.h src/server/WFServer.h src/server/WFHttpServer.h src/server/WFRedisServer.h src/server/WFMySQLServer.h src/client/WFMySQLConnection.h - src/manager/DNSCache.h + src/client/WFDnsClient.h + src/manager/DnsCache.h src/manager/WFGlobal.h src/manager/UpstreamManager.h src/manager/RouteManager.h @@ -72,6 +75,7 @@ set(INCLUDE_HEADERS src/util/StringUtil.h src/util/URIParser.h src/util/MD5Util.h + src/util/DnsUtil.h src/factory/WFConnection.h src/factory/WFTask.h src/factory/WFTask.inl @@ -84,7 +88,7 @@ set(INCLUDE_HEADERS src/factory/Workflow.h src/factory/WFOperator.h src/nameservice/WFNameService.h - src/nameservice/WFDNSResolver.h + src/nameservice/WFDnsResolver.h src/nameservice/WFServiceGovernance.h src/nameservice/UpstreamPolicies.h ) diff --git a/src/algorithm/CMakeLists.txt b/src/algorithm/CMakeLists.txt index 013c49c8..11602d8d 100644 --- a/src/algorithm/CMakeLists.txt +++ b/src/algorithm/CMakeLists.txt @@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.6) project(algorithm) set(SRC - DNSRoutine.cc + DnsRoutine.cc ) add_library(${PROJECT_NAME} OBJECT ${SRC}) diff --git a/src/algorithm/DNSRoutine.cc b/src/algorithm/DnsRoutine.cc similarity index 90% rename from src/algorithm/DNSRoutine.cc rename to src/algorithm/DnsRoutine.cc index c0938641..a8f4f8e1 100644 --- a/src/algorithm/DNSRoutine.cc +++ b/src/algorithm/DnsRoutine.cc @@ -25,11 +25,11 @@ #include #include #include -#include "DNSRoutine.h" +#include "DnsRoutine.h" #define PORT_STR_MAX 5 -DNSOutput::DNSOutput(DNSOutput&& move) +DnsOutput::DnsOutput(DnsOutput&& move) { error_ = move.error_; addrinfo_ = move.addrinfo_; @@ -38,7 +38,7 @@ DNSOutput::DNSOutput(DNSOutput&& move) move.addrinfo_ = NULL; } -DNSOutput& DNSOutput::operator= (DNSOutput&& move) +DnsOutput& DnsOutput::operator= (DnsOutput&& move) { if (this != &move) { @@ -55,7 +55,7 @@ DNSOutput& DNSOutput::operator= (DNSOutput&& move) return *this; } -void DNSRoutine::run_local_path(const std::string& path, DNSOutput *out) +void DnsRoutine::run_local_path(const std::string& path, DnsOutput *out) { struct sockaddr_un *sun = NULL; @@ -85,7 +85,7 @@ void DNSRoutine::run_local_path(const std::string& path, DNSOutput *out) out->error_ = EAI_SYSTEM; } -void DNSRoutine::run(const DNSInput *in, DNSOutput *out) +void DnsRoutine::run(const DnsInput *in, DnsOutput *out) { if (!in->host_.empty() && in->host_[0] == '/') { diff --git a/src/algorithm/DNSRoutine.h b/src/algorithm/DnsRoutine.h similarity index 69% rename from src/algorithm/DNSRoutine.h rename to src/algorithm/DnsRoutine.h index 0ec69aad..b697f100 100644 --- a/src/algorithm/DNSRoutine.h +++ b/src/algorithm/DnsRoutine.h @@ -23,17 +23,17 @@ #include #include -class DNSInput +class DnsInput { public: - DNSInput(): + DnsInput(): port_(0) {} //move constructor - DNSInput(DNSInput&& move) = default; + DnsInput(DnsInput&& move) = default; //move operator - DNSInput& operator= (DNSInput &&move) = default; + DnsInput& operator= (DnsInput &&move) = default; void reset(const std::string& host, unsigned short port) { @@ -48,32 +48,32 @@ protected: std::string host_; unsigned short port_; - friend class DNSRoutine; + friend class DnsRoutine; }; -class DNSOutput +class DnsOutput { public: - DNSOutput(): + DnsOutput(): error_(0), addrinfo_(NULL) {} - ~DNSOutput() + ~DnsOutput() { if (addrinfo_) freeaddrinfo(addrinfo_); } //move constructor - DNSOutput(DNSOutput&& move); + DnsOutput(DnsOutput&& move); //move operator - DNSOutput& operator= (DNSOutput&& move); + DnsOutput& operator= (DnsOutput&& move); int get_error() const { return error_; } const struct addrinfo *get_addrinfo() const { return addrinfo_; } - //if DONOT want DNSOutput release addrinfo, use move_addrinfo in callback + //if DONOT want DnsOutput release addrinfo, use move_addrinfo in callback struct addrinfo *move_addrinfo() { struct addrinfo *p = addrinfo_; @@ -85,19 +85,27 @@ protected: int error_; struct addrinfo *addrinfo_; - friend class DNSRoutine; + friend class DnsRoutine; }; -class DNSRoutine +class DnsRoutine { public: - static void run(const DNSInput *in, DNSOutput *out); + static void run(const DnsInput *in, DnsOutput *out); + static void create(DnsOutput *out, int error, struct addrinfo *ai) + { + if (out->addrinfo_) + freeaddrinfo(out->addrinfo_); + + out->error_ = error; + out->addrinfo_ = ai; + } private: - static void run_local_path(const std::string& path, DNSOutput *out); + static void run_local_path(const std::string& path, DnsOutput *out); }; -//new WFDNSTask(queue, executor, dns_routine, callback) +//new WFDnsTask(queue, executor, dns_routine, callback) //if donot want freeaddrinfo, please std::move output in callback #endif diff --git a/src/client/CMakeLists.txt b/src/client/CMakeLists.txt index 3785e976..876c3a38 100644 --- a/src/client/CMakeLists.txt +++ b/src/client/CMakeLists.txt @@ -2,6 +2,7 @@ cmake_minimum_required(VERSION 3.6) project(client) set(SRC + WFDnsClient.cc WFMySQLConnection.cc ) diff --git a/src/client/WFDnsClient.cc b/src/client/WFDnsClient.cc new file mode 100644 index 00000000..151ac910 --- /dev/null +++ b/src/client/WFDnsClient.cc @@ -0,0 +1,275 @@ +/* + 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 "URIParser.h" +#include "StringUtil.h" +#include "WFDnsClient.h" + +using namespace protocol; + +using DnsCtx = std::function; +using ComplexTask = WFComplexClientTask; + +class DnsParams +{ +public: + struct dns_params + { + std::vector uris; + std::vector search_list; + int ndots; + int attempts; + bool rotate; + }; + +public: + DnsParams() + { + this->ref = new std::atomic(1); + this->params = new dns_params(); + } + + DnsParams(const DnsParams& p) + { + this->ref = p.ref; + this->params = p.params; + this->incref(); + } + + DnsParams& operator=(const DnsParams& p) + { + if (this != &p) + { + this->decref(); + this->ref = p.ref; + this->params = p.params; + this->incref(); + } + return *this; + } + + ~DnsParams() { this->decref(); } + + const dns_params *get_params() const { return this->params; } + dns_params *get_params() { return this->params; } + +private: + void incref() { (*this->ref)++; } + void decref() + { + if (--*this->ref == 0) + { + delete this->params; + delete this->ref; + } + } + +private: + dns_params *params; + std::atomic *ref; +}; + +enum +{ + DNS_STATUS_TRY_ORIGIN_DONE = 0, + DNS_STATUS_TRY_ORIGIN_FIRST = 1, + DNS_STATUS_TRY_ORIGIN_LAST = 2 +}; + +struct DnsStatus +{ + std::string origin_name; + std::string current_name; + size_t next_server; // next server to try + size_t last_server; // last server to try + size_t next_domain; // next search domain to try + int attempt_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; +} + +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(); + DnsResponse *resp = task->get_resp(); + const auto *p = params.get_params(); + 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.attempt_left--; + if (s.attempt_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; + } +} + +int WFDnsClient::init(const std::string& url) +{ + return this->init(url, "", 1, 2, false); +} + +int WFDnsClient::init(const std::string& url, const std::string& search_list, + int ndots, int attempts, bool rotate) +{ + std::vector hosts; + std::vector uris; + std::string host; + ParsedURI uri; + + id = 0; + hosts = StringUtil::split_filter_empty(url, ','); + + for(size_t i = 0; i < hosts.size(); i++) + { + host = hosts[i]; + if (strncasecmp(host.c_str(), "dns://", 6) != 0 && + strncasecmp(host.c_str(), "dnss://", 7) != 0) + { + host = "dns://" + host; + } + + if (URIParser::parse(host, uri) != 0) + return -1; + + uris.emplace_back(std::move(uri)); + } + + if (uris.empty() || ndots < 0 || attempts < 1) + { + errno = EINVAL; + return -1; + } + + this->params = new DnsParams; + DnsParams::dns_params *q = ((DnsParams *)(this->params))->get_params(); + q->uris = std::move(uris); + q->search_list = StringUtil::split_filter_empty(search_list, ','); + q->ndots = ndots > 15 ? 15 : ndots; + q->attempts = attempts > 5 ? 5 : attempts; + q->rotate = rotate; + + return 0; +} + +void WFDnsClient::deinit() +{ + delete (DnsParams *)(this->params); + this->params = NULL; +} + +WFDnsTask *WFDnsClient::create_dns_task(const std::string& name, + dns_callback_t callback) +{ + DnsParams::dns_params *p = ((DnsParams *)(this->params))->get_params(); + struct DnsStatus status; + size_t next_server; + WFDnsTask *task; + DnsRequest *req; + + next_server = p->rotate ? this->id++ % p->uris.size() : 0; + + status.origin_name = name; + status.next_domain = 0; + status.attempt_left = p->attempts; + status.try_origin_state = DNS_STATUS_TRY_ORIGIN_FIRST; + + if (!name.empty() && name.back() == '.') + status.next_domain = p->uris.size(); + else if (__get_ndots(name) < p->ndots) + status.try_origin_state = DNS_STATUS_TRY_ORIGIN_LAST; + + __has_next_name(p, &status); + + 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(); + req->set_question(status.current_name.c_str(), DNS_TYPE_A, DNS_CLASS_IN); + req->set_rd(1); + + ComplexTask *ctask = static_cast(task); + *ctask->get_mutable_ctx() = std::bind(__callback_internal, + std::placeholders::_1, + *(DnsParams *)params, status); + + return task; +} diff --git a/src/client/WFDnsClient.h b/src/client/WFDnsClient.h new file mode 100644 index 00000000..7ae86e9b --- /dev/null +++ b/src/client/WFDnsClient.h @@ -0,0 +1,45 @@ +/* + 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) +*/ + +#ifndef _WFDNSCLIENT_H_ +#define _WFDNSCLIENT_H_ + +#include +#include "WFTaskFactory.h" +#include "DnsMessage.h" + +class WFDnsClient +{ +public: + WFDnsClient() : params(NULL) { } + virtual ~WFDnsClient() { } + + int init(const std::string& url); + int init(const std::string& url, const std::string& search_list, + int ndots, int attempts, bool rotate); + void deinit(); + + WFDnsTask *create_dns_task(const std::string& name, + dns_callback_t callback); + +private: + void *params; + std::atomic id; +}; + +#endif diff --git a/src/factory/CMakeLists.txt b/src/factory/CMakeLists.txt index 24689fa6..a11f6c5c 100644 --- a/src/factory/CMakeLists.txt +++ b/src/factory/CMakeLists.txt @@ -3,6 +3,7 @@ project(factory) set(SRC WFGraphTask.cc + DnsTaskImpl.cc HttpTaskImpl.cc RedisTaskImpl.cc MySQLTaskImpl.cc diff --git a/src/factory/DnsTaskImpl.cc b/src/factory/DnsTaskImpl.cc new file mode 100644 index 00000000..e9464db9 --- /dev/null +++ b/src/factory/DnsTaskImpl.cc @@ -0,0 +1,195 @@ +/* + 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 "WFTaskError.h" +#include "WFTaskFactory.h" +#include "DnsMessage.h" + +using namespace protocol; + +#define DNS_KEEPALIVE_DEFAULT (60 * 1000) + +class ComplexDnsTask : public WFComplexClientTask> +{ + static struct addrinfo hints; + +public: + ComplexDnsTask(int retry_max, dns_callback_t&& cb): + WFComplexClientTask(retry_max, std::move(cb)) + { + this->set_transport_type(TT_UDP); + } + +protected: + virtual CommMessageOut *message_out(); + virtual CommMessageIn *message_in(); + virtual bool init_success(); + virtual bool finish_once(); + +private: + bool need_redirect(); +}; + +struct addrinfo ComplexDnsTask::hints = +{ + .ai_flags = AI_NUMERICSERV | AI_NUMERICHOST, + .ai_family = AF_UNSPEC, + .ai_socktype = SOCK_STREAM, + .ai_protocol = 0, + .ai_addrlen = 0, + .ai_addr = NULL, + .ai_canonname = NULL, + .ai_next = NULL +}; + +CommMessageOut *ComplexDnsTask::message_out() +{ + DnsRequest *req = this->get_req(); + DnsResponse *resp = this->get_resp(); + TransportType type = this->get_transport_type(); + + /* Set these field every time, in case of reconstruct on redirect */ + resp->set_request_id(req->get_id()); + resp->set_request_name(req->get_question_name()); + req->set_leading_length(type != TT_UDP); + resp->set_leading_length(type != TT_UDP); + + return this->WFClientTask::message_out(); +} + +CommMessageIn *ComplexDnsTask::message_in() +{ + return this->WFClientTask::message_in(); +} + +bool ComplexDnsTask::init_success() +{ + static constexpr struct EndpointParams ep = + { + .max_connections = 65535, + .connect_timeout = 10 * 1000, + .response_timeout = 10 * 1000, + .ssl_connect_timeout = 10 * 1000, + .use_tls_sni = false + }; + + if (uri_.scheme && strcasecmp(uri_.scheme, "dnss") == 0) + this->WFComplexClientTask::set_transport_type(TT_TCP_SSL); + else if (uri_.scheme && strcasecmp(uri_.scheme, "dns") != 0) + { + this->state = WFT_STATE_TASK_ERROR; + this->error = WFT_ERR_URI_SCHEME_INVALID; + return false; + } + + if (!this->route_result_.request_object) + { + TransportType type = this->get_transport_type(); + struct addrinfo *addr; + int ret; + + ret = getaddrinfo(uri_.host, uri_.port, &hints, &addr); + if (ret != 0) + { + this->state = WFT_STATE_TASK_ERROR; + this->error = WFT_ERR_URI_PARSE_FAILED; + return false; + } + + ret = WFGlobal::get_route_manager()->get(type, addr, info_, &ep, + uri_.host, route_result_); + freeaddrinfo(addr); + if (ret < 0) + { + this->state = WFT_STATE_SYS_ERROR; + this->error = errno; + return false; + } + } + + return true; +} + +bool ComplexDnsTask::finish_once() +{ + if (this->state == WFT_STATE_SUCCESS) + { + if (need_redirect()) + this->set_redirect(uri_); + else if (this->state != WFT_STATE_SUCCESS) + this->disable_retry(); + } + + /* If retry times meet retry max and there is no redirect, + * we ask the client for a retry or redirect. + */ + if (retry_times_ == retry_max_ && !redirect_ && *this->get_mutable_ctx()) + { + /* Reset type to UDP before a client redirect. */ + this->set_transport_type(TT_UDP); + (*this->get_mutable_ctx())(this); + } + + return true; +} + +bool ComplexDnsTask::need_redirect() +{ + DnsResponse *client_resp = this->get_resp(); + TransportType type = this->get_transport_type(); + + if (type == TT_UDP && client_resp->get_tc() == 1) + { + this->set_transport_type(TT_TCP); + return true; + } + + return false; +} + +WFDnsTask *WFTaskFactory::create_dns_task(const std::string& url, + int retry_max, + dns_callback_t callback) +{ + ParsedURI uri; + + URIParser::parse(url, uri); + return WFTaskFactory::create_dns_task(uri, retry_max, std::move(callback)); +} + +WFDnsTask *WFTaskFactory::create_dns_task(const ParsedURI& uri, + int retry_max, + dns_callback_t callback) +{ + ComplexDnsTask *task = new ComplexDnsTask(retry_max, std::move(callback)); + const char *name; + + if (uri.path && uri.path[0] && uri.path[1]) + name = uri.path + 1; + else + name = "."; + + DnsRequest *req = task->get_req(); + req->set_question(name, DNS_TYPE_A, DNS_CLASS_IN); + + task->init(uri); + task->set_keep_alive(DNS_KEEPALIVE_DEFAULT); + return task; +} diff --git a/src/factory/WFTaskFactory.cc b/src/factory/WFTaskFactory.cc index 5ee626e4..ce42b3c0 100644 --- a/src/factory/WFTaskFactory.cc +++ b/src/factory/WFTaskFactory.cc @@ -293,20 +293,6 @@ void WFTaskFactory::count_by_name(const std::string& counter_name, unsigned int __CounterMap::get_instance()->count_n(counter_name, n); } -WFDNSTask *WFTaskFactory::create_dns_task(const std::string& host, - unsigned short port, - dns_callback_t callback) -{ - auto *task = WFThreadTaskFactory:: - create_thread_task(WFGlobal::get_dns_queue(), - WFGlobal::get_dns_executor(), - DNSRoutine::run, - std::move(callback)); - - task->get_input()->reset(host, port); - return task; -} - /********FileIOTask*************/ class WFFilepreadTask : public WFFileIOTask diff --git a/src/factory/WFTaskFactory.h b/src/factory/WFTaskFactory.h index 2057608a..430df40c 100644 --- a/src/factory/WFTaskFactory.h +++ b/src/factory/WFTaskFactory.h @@ -28,6 +28,7 @@ #include "RedisMessage.h" #include "HttpMessage.h" #include "MySQLMessage.h" +#include "DnsMessage.h" #include "DNSRoutine.h" #include "Workflow.h" #include "WFTask.h" @@ -49,6 +50,10 @@ using WFMySQLTask = WFNetworkTask; using mysql_callback_t = std::function; +using WFDnsTask = WFNetworkTask; +using dns_callback_t = std::function; + // File IO tasks struct FileIOArgs @@ -93,10 +98,6 @@ using WFEmptyTask = WFGenericTask; using WFDynamicTask = WFGenericTask; using dynamic_create_t = std::function; -// DNS task. For internal usage only. -using WFDNSTask = WFThreadTask; -using dns_callback_t = std::function; - class WFTaskFactory { public: @@ -138,6 +139,14 @@ public: int retry_max, mysql_callback_t callback); + static WFDnsTask *create_dns_task(const std::string& url, + int retry_max, + dns_callback_t callback); + + static WFDnsTask *create_dns_task(const ParsedURI& uri, + int retry_max, + dns_callback_t callback); + public: static WFFileIOTask *create_pread_task(int fd, void *buf, @@ -225,11 +234,6 @@ public: return new WFGraphTask(std::move(callback)); } -public: - static WFDNSTask *create_dns_task(const std::string& host, - unsigned short port, - dns_callback_t callback); - public: static WFEmptyTask *create_empty_task() { diff --git a/src/manager/CMakeLists.txt b/src/manager/CMakeLists.txt index 6a282556..91af981c 100644 --- a/src/manager/CMakeLists.txt +++ b/src/manager/CMakeLists.txt @@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.6) project(manager) set(SRC - DNSCache.cc + DnsCache.cc UpstreamManager.cc RouteManager.cc WFGlobal.cc diff --git a/src/manager/DNSCache.cc b/src/manager/DnsCache.cc similarity index 86% rename from src/manager/DNSCache.cc rename to src/manager/DnsCache.cc index fec5da4c..8b3ba049 100644 --- a/src/manager/DNSCache.cc +++ b/src/manager/DnsCache.cc @@ -18,16 +18,16 @@ #include #include -#include "DNSCache.h" +#include "DnsCache.h" #define GET_CURRENT_SECOND std::chrono::duration_cast(std::chrono::steady_clock::now().time_since_epoch()).count() #define CONFIDENT_INC 10 #define TTL_INC 10 -const DNSCache::DNSHandle *DNSCache::get_inner(const HostPort& host_port, int type) +const DnsCache::DnsHandle *DnsCache::get_inner(const HostPort& host_port, int type) { - const DNSHandle *handle = cache_pool_.get(host_port); + const DnsHandle *handle = cache_pool_.get(host_port); if (handle) { @@ -42,7 +42,7 @@ const DNSCache::DNSHandle *DNSCache::get_inner(const HostPort& host_port, int ty if (cur_time > handle->value.expire_time) { - const_cast(handle)->value.expire_time += TTL_INC; + const_cast(handle)->value.expire_time += TTL_INC; cache_pool_.release(handle); return NULL; } @@ -57,7 +57,7 @@ const DNSCache::DNSHandle *DNSCache::get_inner(const HostPort& host_port, int ty if (cur_time > handle->value.confident_time) { - const_cast(handle)->value.confident_time += CONFIDENT_INC; + const_cast(handle)->value.confident_time += CONFIDENT_INC; cache_pool_.release(handle); return NULL; } @@ -73,7 +73,7 @@ const DNSCache::DNSHandle *DNSCache::get_inner(const HostPort& host_port, int ty return handle; } -const DNSCache::DNSHandle *DNSCache::put(const HostPort& host_port, +const DnsCache::DnsHandle *DnsCache::put(const HostPort& host_port, struct addrinfo *addrinfo, unsigned int dns_ttl_default, unsigned int dns_ttl_min) diff --git a/src/manager/DNSCache.h b/src/manager/DnsCache.h similarity index 73% rename from src/manager/DNSCache.h rename to src/manager/DnsCache.h index 594c418e..fa719dae 100644 --- a/src/manager/DNSCache.h +++ b/src/manager/DnsCache.h @@ -31,7 +31,7 @@ #define GET_TYPE_TTL 0 #define GET_TYPE_CONFIDENT 1 -struct DNSCacheValue +struct DnsCacheValue { struct addrinfo *addrinfo; int64_t confident_time; @@ -41,7 +41,7 @@ struct DNSCacheValue class ValueDeleter { public: - void operator() (const DNSCacheValue& value) const + void operator() (const DnsCacheValue& value) const { freeaddrinfo(value.addrinfo); } @@ -50,20 +50,20 @@ public: // RAII: NO. Release handle by user // Thread safety: YES // MUST call release when handle no longer used -class DNSCache +class DnsCache { public: using HostPort = std::pair; - using DNSHandle = LRUHandle; + using DnsHandle = LRUHandle; public: // release handle by get/put - void release(DNSHandle *handle) + void release(DnsHandle *handle) { cache_pool_.release(handle); } - void release(const DNSHandle *handle) + void release(const DnsHandle *handle) { cache_pool_.release(handle); } @@ -71,57 +71,57 @@ public: // get handler // Need call release when handle no longer needed //Handle *get(const KEY &key); - const DNSHandle *get(const HostPort& host_port) + const DnsHandle *get(const HostPort& host_port) { return cache_pool_.get(host_port); } - const DNSHandle *get(const std::string& host, unsigned short port) + const DnsHandle *get(const std::string& host, unsigned short port) { return get(HostPort(host, port)); } - const DNSHandle *get(const char *host, unsigned short port) + const DnsHandle *get(const char *host, unsigned short port) { return get(std::string(host), port); } - const DNSHandle *get_ttl(const HostPort& host_port) + const DnsHandle *get_ttl(const HostPort& host_port) { return get_inner(host_port, GET_TYPE_TTL); } - const DNSHandle *get_ttl(const std::string& host, unsigned short port) + const DnsHandle *get_ttl(const std::string& host, unsigned short port) { return get_ttl(HostPort(host, port)); } - const DNSHandle *get_ttl(const char *host, unsigned short port) + const DnsHandle *get_ttl(const char *host, unsigned short port) { return get_ttl(std::string(host), port); } - const DNSHandle *get_confident(const HostPort& host_port) + const DnsHandle *get_confident(const HostPort& host_port) { return get_inner(host_port, GET_TYPE_CONFIDENT); } - const DNSHandle *get_confident(const std::string& host, unsigned short port) + const DnsHandle *get_confident(const std::string& host, unsigned short port) { return get_confident(HostPort(host, port)); } - const DNSHandle *get_confident(const char *host, unsigned short port) + const DnsHandle *get_confident(const char *host, unsigned short port) { return get_confident(std::string(host), port); } - const DNSHandle *put(const HostPort& host_port, + const DnsHandle *put(const HostPort& host_port, struct addrinfo *addrinfo, unsigned int dns_ttl_default, unsigned int dns_ttl_min); - const DNSHandle *put(const std::string& host, + const DnsHandle *put(const std::string& host, unsigned short port, struct addrinfo *addrinfo, unsigned int dns_ttl_default, @@ -130,7 +130,7 @@ public: return put(HostPort(host, port), addrinfo, dns_ttl_default, dns_ttl_min); } - const DNSHandle *put(const char *host, + const DnsHandle *put(const char *host, unsigned short port, struct addrinfo *addrinfo, unsigned int dns_ttl_default, @@ -156,10 +156,10 @@ public: } private: - const DNSHandle *get_inner(const HostPort& host_port, int type); + const DnsHandle *get_inner(const HostPort& host_port, int type); std::mutex mutex_; - LRUCache cache_pool_; + LRUCache cache_pool_; }; #endif diff --git a/src/manager/WFGlobal.cc b/src/manager/WFGlobal.cc index 79e58b04..ca3f7c19 100644 --- a/src/manager/WFGlobal.cc +++ b/src/manager/WFGlobal.cc @@ -34,13 +34,14 @@ #include "WFGlobal.h" #include "EndpointParams.h" #include "CommScheduler.h" -#include "DNSCache.h" +#include "DnsCache.h" #include "RouteManager.h" #include "Executor.h" #include "WFTask.h" #include "WFTaskError.h" #include "WFNameService.h" -#include "WFDNSResolver.h" +#include "WFDnsResolver.h" +#include "WFDnsClient.h" class __WFGlobal { @@ -122,6 +123,15 @@ private: __WFGlobal::__WFGlobal() : settings_(GLOBAL_SETTINGS_DEFAULT) { + static_scheme_port_["dns"] = "53"; + static_scheme_port_["Dns"] = "53"; + static_scheme_port_["DNS"] = "53"; + + static_scheme_port_["dnss"] = "853"; + static_scheme_port_["Dnss"] = "853"; + static_scheme_port_["DNSs"] = "853"; + static_scheme_port_["DNSS"] = "853"; + static_scheme_port_["http"] = "80"; static_scheme_port_["Http"] = "80"; static_scheme_port_["HTTP"] = "80"; @@ -278,13 +288,13 @@ private: bool flag_; }; -class __DNSManager +class __DnsManager { public: ExecQueue *get_dns_queue() { return &dns_queue_; } Executor *get_dns_executor() { return &dns_executor_; } - __DNSManager() + __DnsManager() { int ret; @@ -299,7 +309,7 @@ public: abort(); } - ~__DNSManager() + ~__DnsManager() { dns_executor_.deinit(); dns_queue_.deinit(); @@ -383,14 +393,14 @@ private: } } - __DNSManager *get_dns_manager_safe() + __DnsManager *get_dns_manager_safe() { if (!dns_flag_) { dns_mutex_.lock(); if (!dns_flag_) { - dns_manager_ = new __DNSManager(); + dns_manager_ = new __DnsManager(); dns_flag_ = true; } @@ -406,29 +416,29 @@ private: IOServer *io_server_; volatile bool io_flag_; std::mutex io_mutex_; - __DNSManager *dns_manager_; + __DnsManager *dns_manager_; volatile bool dns_flag_; std::mutex dns_mutex_; }; -class __DNSCache +class __DnsCache { public: - static __DNSCache *get_instance() + static __DnsCache *get_instance() { - static __DNSCache kInstance; + static __DnsCache kInstance; return &kInstance; } - DNSCache *get_dns_cache() { return &dns_cache_; } + DnsCache *get_dns_cache() { return &dns_cache_; } private: - __DNSCache() { } + __DnsCache() { } - ~__DNSCache() { } + ~__DnsCache() { } private: - DNSCache dns_cache_; + DnsCache dns_cache_; }; class __ExecManager @@ -529,23 +539,217 @@ public: WFNameService *get_name_service() { return &service_; } private: - static WFDNSResolver resolver_; + static WFDnsResolver resolver_; WFNameService service_; public: __NameServiceManager() : service_(&__NameServiceManager::resolver_) { } }; -WFDNSResolver __NameServiceManager::resolver_; +WFDnsResolver __NameServiceManager::resolver_; + +#define MAX(x, y) ((x) >= (y) ? (x) : (y)) +#define HOSTS_LINEBUF_INIT_SIZE 128 + +static int __read_line(FILE *fp, void **buf, size_t *bufsize) +{ + char *p = (char *)*buf; + size_t offset = 0; + size_t newsize; + void *newbase; + int nleft; + + while (1) + { + if (offset + 1 >= *bufsize) + { + newsize = MAX(HOSTS_LINEBUF_INIT_SIZE, *bufsize * 2); + newbase = realloc(*buf, newsize); + if (!newbase) + return -1; + + *buf = newbase; + *bufsize = newsize; + p = (char *)*buf; + } + + nleft = *bufsize - offset; + if (!fgets(&p[offset], nleft, fp)) + { + if (offset != 0 && !ferror(fp)) + break; + + return ferror(fp) ? -1 : 1; + } + + offset += strlen(&p[offset]); + if (p[offset - 1] == '\n') + { + p[offset - 1] = '\0'; + break; + } + } + + return 0; +} + +static void __split_merge_str(const char *p, std::string& result) +{ + const char *start; + + if (!isspace(*p)) + return; + + while (1) + { + while (isspace(*p)) + p++; + + start = p; + while (*p && *p != '#' && *p != ';' && !isspace(*p)) + p++; + + if (start == p) + break; + + if (!result.empty()) + result.push_back(','); + result.append(start, p); + } +} + +static inline const char *__try_options(const char *p, const char *q, + const char *r) +{ + size_t len = strlen(r); + if ((size_t)(q - p) >= len && strncmp(p, r, len) == 0) + return p + len; + return NULL; +} + +static void __set_options(const char *p, int& ndots, int& attempt, bool& rotate) +{ + const char *start; + const char *opt; + + if (!isspace(*p)) + return; + + while (1) + { + while (isspace(*p)) + p++; + + start = p; + while (*p && *p != '#' && *p != ';' && !isspace(*p)) + p++; + + if (start == p) + break; + + if ((opt = __try_options(start, p, "ndots:")) != NULL) + ndots = atoi(opt); + else if ((opt = __try_options(start, p, "attempts:")) != NULL) + attempt = atoi(opt); + else if ((opt = __try_options(start, p, "rotate")) != NULL) + rotate = true; + } +} + +static int __parse_resolv_conf(const char *filename, + std::string& url, std::string& search_list, + int& ndots, int& attempts, bool& rotate) +{ + size_t bufsize = 0; + void *line = NULL; + FILE *fp; + int ret; + + fp = fopen(filename, "r"); + if (!fp) + return -1; + + while ((ret = __read_line(fp, &line, &bufsize)) == 0) + { + const char *p = (const char *)line; + if (strncmp(p, "nameserver", 10) == 0) + __split_merge_str(p + 10, url); + else if (strncmp(p, "search", 6) == 0) + __split_merge_str(p + 6, search_list); + else if (strncmp(p, "options", 7) == 0) + __set_options(p + 7, ndots, attempts, rotate); + } + + free(line); + fclose(fp); + + return ret; +} + +class __DnsClientManager +{ +public: + static __DnsClientManager *get_instance() + { + static __DnsClientManager kInstance; + return &kInstance; + } + +public: + WFDnsClient *get_dns_client() { return client; } + +private: + __DnsClientManager() + { + const char *path = WFGlobal::get_global_settings()->resolv_conf_path; + + client = NULL; + if (path && path[0]) + { + int ndots = 1; + int attempts = 2; + bool rotate = false; + std::string url; + std::string search; + + client = new WFDnsClient; + if (__parse_resolv_conf(path, url, search, ndots, attempts, + rotate) >= 0) + { + if (client->init(url, search, ndots, attempts, rotate) >= 0) + return; + } + + delete client; + abort(); + } + } + + ~__DnsClientManager() + { + if (client) + { + client->deinit(); + delete client; + } + } + + WFDnsClient *client; +}; + +WFDnsClient *WFGlobal::get_dns_client() +{ + return __DnsClientManager::get_instance()->get_dns_client(); +} CommScheduler *WFGlobal::get_scheduler() { return __CommManager::get_instance()->get_scheduler(); } -DNSCache *WFGlobal::get_dns_cache() +DnsCache *WFGlobal::get_dns_cache() { - return __DNSCache::get_instance()->get_dns_cache(); + return __DnsCache::get_instance()->get_dns_cache(); } RouteManager *WFGlobal::get_route_manager() diff --git a/src/manager/WFGlobal.h b/src/manager/WFGlobal.h index 6ad04bbc..0e15621a 100644 --- a/src/manager/WFGlobal.h +++ b/src/manager/WFGlobal.h @@ -27,7 +27,7 @@ #include #include #include "CommScheduler.h" -#include "DNSCache.h" +#include "DnsCache.h" #include "RouteManager.h" #include "Executor.h" #include "EndpointParams.h" @@ -52,6 +52,8 @@ struct WFGlobalSettings int poller_threads; int handler_threads; int compute_threads; ///< auto-set by system CPU number if value<=0 + const char *resolv_conf_path; + const char *hosts_path; }; /** @@ -66,6 +68,8 @@ static constexpr struct WFGlobalSettings GLOBAL_SETTINGS_DEFAULT = .poller_threads = 4, .handler_threads = 20, .compute_threads = -1, + .resolv_conf_path = NULL, // use thread dns task for default + .hosts_path = NULL, }; /** @@ -108,7 +112,7 @@ public: public: // Internal usage only static CommScheduler *get_scheduler(); - static DNSCache *get_dns_cache(); + static DnsCache *get_dns_cache(); static RouteManager *get_route_manager(); static SSL_CTX *get_ssl_client_ctx(); static SSL_CTX *new_ssl_server_ctx(); @@ -118,6 +122,7 @@ public: static ExecQueue *get_dns_queue(); static Executor *get_dns_executor(); static WFNameService *get_name_service(); + static class WFDnsClient *get_dns_client(); static void sync_operation_begin(); static void sync_operation_end(); }; diff --git a/src/nameservice/CMakeLists.txt b/src/nameservice/CMakeLists.txt index 984a7325..504f6b27 100644 --- a/src/nameservice/CMakeLists.txt +++ b/src/nameservice/CMakeLists.txt @@ -3,7 +3,7 @@ project(nameservice) set(SRC WFNameService.cc - WFDNSResolver.cc + WFDnsResolver.cc WFServiceGovernance.cc UpstreamPolicies.cc ) diff --git a/src/nameservice/UpstreamPolicies.h b/src/nameservice/UpstreamPolicies.h index 89a23d6c..8bd5fb19 100644 --- a/src/nameservice/UpstreamPolicies.h +++ b/src/nameservice/UpstreamPolicies.h @@ -25,7 +25,7 @@ #include "URIParser.h" #include "EndpointParams.h" #include "WFNameService.h" -#include "WFDNSResolver.h" +#include "WFDnsResolver.h" #include "WFGlobal.h" #include "WFTaskError.h" #include "WFServiceGovernance.h" diff --git a/src/nameservice/WFDNSResolver.cc b/src/nameservice/WFDNSResolver.cc deleted file mode 100644 index 625122fd..00000000 --- a/src/nameservice/WFDNSResolver.cc +++ /dev/null @@ -1,282 +0,0 @@ -/* - Copyright (c) 2020 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. - - Authors: Wu Jiaxu (wujiaxu@sogou-inc.com) - Xie Han (xiehan@sogou-inc.com) -*/ - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include "DNSRoutine.h" -#include "EndpointParams.h" -#include "RouteManager.h" -#include "WFGlobal.h" -#include "WFTaskError.h" -#include "WFTaskFactory.h" -#include "WFNameService.h" -#include "WFDNSResolver.h" - -/* -DNS_CACHE_LEVEL_0 -> NO cache -DNS_CACHE_LEVEL_1 -> TTL MIN -DNS_CACHE_LEVEL_2 -> TTL [DEFAULT] -DNS_CACHE_LEVEL_3 -> Forever -*/ - -#define DNS_CACHE_LEVEL_0 0 -#define DNS_CACHE_LEVEL_1 1 -#define DNS_CACHE_LEVEL_2 2 -#define DNS_CACHE_LEVEL_3 3 - -class WFResolverTask : public WFRouterTask -{ -public: - WFResolverTask(const struct WFNSParams *params, int dns_cache_level, - unsigned int dns_ttl_default, unsigned int dns_ttl_min, - const struct EndpointParams *endpoint_params, - router_callback_t&& cb) : - WFRouterTask(std::move(cb)), - type_(params->type), - host_(params->uri.host ? params->uri.host : ""), - port_(params->uri.port ? atoi(params->uri.port) : 0), - info_(params->info), - dns_cache_level_(dns_cache_level), - dns_ttl_default_(dns_ttl_default), - dns_ttl_min_(dns_ttl_min), - endpoint_params_(*endpoint_params), - first_addr_only_(params->fixed_addr) - { - } - -private: - virtual void dispatch(); - virtual SubTask *done(); - void dns_callback(WFDNSTask *dns_task); - void dns_callback_internal(DNSOutput *dns_task, - unsigned int ttl_default, - unsigned int ttl_min); - -private: - TransportType type_; - std::string host_; - unsigned short port_; - std::string info_; - int dns_cache_level_; - unsigned int dns_ttl_default_; - unsigned int dns_ttl_min_; - struct EndpointParams endpoint_params_; - bool first_addr_only_; - bool insert_dns_; -}; - -void WFResolverTask::dispatch() -{ - insert_dns_ = true; - if (dns_cache_level_ != DNS_CACHE_LEVEL_0) - { - auto *dns_cache = WFGlobal::get_dns_cache(); - const DNSCache::DNSHandle *addr_handle = NULL; - - switch (dns_cache_level_) - { - case DNS_CACHE_LEVEL_1: - addr_handle = dns_cache->get_confident(host_, port_); - break; - - case DNS_CACHE_LEVEL_2: - addr_handle = dns_cache->get_ttl(host_, port_); - break; - - case DNS_CACHE_LEVEL_3: - addr_handle = dns_cache->get(host_, port_); - break; - - default: - break; - } - - if (addr_handle) - { - auto *route_manager = WFGlobal::get_route_manager(); - struct addrinfo *addrinfo = addr_handle->value.addrinfo; - struct addrinfo first; - - if (first_addr_only_ && addrinfo->ai_next) - { - first = *addrinfo; - first.ai_next = NULL; - addrinfo = &first; - } - - if (route_manager->get(type_, addrinfo, info_, &endpoint_params_, - host_, this->result) < 0) - { - this->state = WFT_STATE_SYS_ERROR; - this->error = errno; - } - else - this->state = WFT_STATE_SUCCESS; - - insert_dns_ = false; - dns_cache->release(addr_handle); - } - } - - if (insert_dns_ && !host_.empty()) - { - char front = host_.front(); - char back = host_.back(); - struct in6_addr addr; - int ret; - - if (host_.find(':') != std::string::npos) - ret = inet_pton(AF_INET6, host_.c_str(), &addr); - else if (isdigit(back) && isdigit(front)) - ret = inet_pton(AF_INET, host_.c_str(), &addr); - else if (front == '/') - ret = 1; - else - ret = 0; - - if (ret == 1) - { - DNSInput dns_in; - DNSOutput dns_out; - - dns_in.reset(host_, port_); - DNSRoutine::run(&dns_in, &dns_out); - dns_callback_internal(&dns_out, (unsigned int)-1, (unsigned int)-1); - insert_dns_ = false; - } - } - - if (insert_dns_) - { - auto&& cb = std::bind(&WFResolverTask::dns_callback, - this, - std::placeholders::_1); - WFDNSTask *dns_task = WFTaskFactory::create_dns_task(host_, port_, - std::move(cb)); - series_of(this)->push_front(dns_task); - } - - this->subtask_done(); -} - -SubTask *WFResolverTask::done() -{ - SeriesWork *series = series_of(this); - - if (!insert_dns_) - { - if (this->callback) - this->callback(this); - - delete this; - } - - return series->pop(); -} - -void WFResolverTask::dns_callback_internal(DNSOutput *dns_out, - unsigned int ttl_default, - unsigned int ttl_min) -{ - int dns_error = dns_out->get_error(); - - if (dns_error) - { - if (dns_error == EAI_SYSTEM) - { - this->state = WFT_STATE_SYS_ERROR; - this->error = errno; - } - else - { - this->state = WFT_STATE_DNS_ERROR; - this->error = dns_error; - } - } - else - { - auto *route_manager = WFGlobal::get_route_manager(); - auto *dns_cache = WFGlobal::get_dns_cache(); - struct addrinfo *addrinfo = dns_out->move_addrinfo(); - const DNSCache::DNSHandle *addr_handle; - - addr_handle = dns_cache->put(host_, port_, addrinfo, - (unsigned int)ttl_default, - (unsigned int)ttl_min); - if (route_manager->get(type_, addrinfo, info_, &endpoint_params_, - host_, this->result) < 0) - { - this->state = WFT_STATE_SYS_ERROR; - this->error = errno; - } - else - this->state = WFT_STATE_SUCCESS; - - dns_cache->release(addr_handle); - } -} - -void WFResolverTask::dns_callback(WFDNSTask *dns_task) -{ - if (dns_task->get_state() == WFT_STATE_SUCCESS) - dns_callback_internal(dns_task->get_output(), dns_ttl_default_, dns_ttl_min_); - else - { - this->state = dns_task->get_state(); - this->error = dns_task->get_error(); - } - - if (this->callback) - this->callback(this); - - delete this; -} - -WFRouterTask * -WFDNSResolver::create(const struct WFNSParams *params, int dns_cache_level, - unsigned int dns_ttl_default, unsigned int dns_ttl_min, - const struct EndpointParams *endpoint_params, - router_callback_t&& callback) -{ - return new WFResolverTask(params, dns_cache_level, - dns_ttl_default, dns_ttl_min, - endpoint_params, std::move(callback)); -} - -WFRouterTask *WFDNSResolver::create_router_task(const struct WFNSParams *params, - router_callback_t callback) -{ - const auto *settings = WFGlobal::get_global_settings(); - unsigned int dns_ttl_default = settings->dns_ttl_default; - unsigned int dns_ttl_min = settings->dns_ttl_min; - const struct EndpointParams *endpoint_params = &settings->endpoint_params; - int dns_cache_level = params->retry_times == 0 ? DNS_CACHE_LEVEL_2 : - DNS_CACHE_LEVEL_1; - return create(params, dns_cache_level, dns_ttl_default, dns_ttl_min, - endpoint_params, std::move(callback)); -} - diff --git a/src/nameservice/WFDnsResolver.cc b/src/nameservice/WFDnsResolver.cc new file mode 100644 index 00000000..6ec137f9 --- /dev/null +++ b/src/nameservice/WFDnsResolver.cc @@ -0,0 +1,637 @@ +/* + Copyright (c) 2020 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. + + Authors: Wu Jiaxu (wujiaxu@sogou-inc.com) + Xie Han (xiehan@sogou-inc.com) +*/ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "DnsRoutine.h" +#include "EndpointParams.h" +#include "RouteManager.h" +#include "WFGlobal.h" +#include "WFTaskError.h" +#include "WFTaskFactory.h" +#include "WFNameService.h" +#include "DnsUtil.h" +#include "WFDnsClient.h" +#include "WFDnsResolver.h" + +/* +DNS_CACHE_LEVEL_0 -> NO cache +DNS_CACHE_LEVEL_1 -> TTL MIN +DNS_CACHE_LEVEL_2 -> TTL [DEFAULT] +DNS_CACHE_LEVEL_3 -> Forever +*/ + +#define DNS_CACHE_LEVEL_0 0 +#define DNS_CACHE_LEVEL_1 1 +#define DNS_CACHE_LEVEL_2 2 +#define DNS_CACHE_LEVEL_3 3 + +#define MAX(x, y) ((x) >= (y) ? (x) : (y)) +#define HOSTS_LINEBUF_INIT_SIZE 128 +#define PORT_STR_MAX 5 + +// Dns Thread task. For internal usage only. +using ThreadDnsTask = WFThreadTask; +using thread_dns_callback_t = std::function; + +static constexpr struct addrinfo __ai_hints = +{ +#ifdef AI_ADDRCONFIG + .ai_flags = AI_ADDRCONFIG, +#else + .ai_flags = 0, +#endif + .ai_family = AF_UNSPEC, + .ai_socktype = SOCK_STREAM, + .ai_protocol = 0, + .ai_addrlen = 0, + .ai_addr = NULL, + .ai_canonname = NULL, + .ai_next = NULL +}; + +struct DnsContext +{ + int state; + int error; + int eai_error; + unsigned short port; + struct addrinfo *ai; +}; + +static int __default_family() +{ + struct addrinfo *res; + struct addrinfo *cur; + int family = AF_UNSPEC; + bool v4 = false; + bool v6 = false; + + if (getaddrinfo(NULL, "0", &__ai_hints, &res) == 0) + { + for(cur = res; cur; cur = cur->ai_next) + { + if(cur->ai_family == AF_INET) + v4 = true; + else if(cur->ai_family == AF_INET6) + v6 = true; + } + + freeaddrinfo(res); + if (v4 ^ v6) + family = v4 ? AF_INET : AF_INET6; + } + + return family; +} + +static int __read_line(FILE *fp, void **buf, size_t *bufsize) +{ + char *p = (char *)*buf; + size_t offset = 0; + size_t newsize; + void *newbase; + int nleft; + + while (1) + { + if (offset + 1 >= *bufsize) + { + newsize = MAX(HOSTS_LINEBUF_INIT_SIZE, *bufsize * 2); + newbase = realloc(*buf, newsize); + if (!newbase) + return -1; + + *buf = newbase; + *bufsize = newsize; + p = (char *)*buf; + } + + nleft = *bufsize - offset; + if (!fgets(&p[offset], nleft, fp)) + { + if (offset != 0 && !ferror(fp)) + break; + + return ferror(fp) ? -1 : 1; + } + + offset += strlen(&p[offset]); + if (p[offset - 1] == '\n') + { + p[offset - 1] = '\0'; + break; + } + } + + return 0; +} + +// hosts line format: IP canonical_name [aliases...] [# Comment] +static int __readaddrinfo_line(char *p, const char *name, const char *port, + const struct addrinfo *hints, + struct addrinfo **res) +{ + const char *ip = NULL; + char *start; + + start = p; + while (*start != '\0' && *start != '#') + start++; + *start = '\0'; + + while (1) + { + while (isspace(*p)) + p++; + + start = p; + while (*p != '\0' && !isspace(*p)) + p++; + + if (start == p) + break; + + if (*p != '\0') + *p++ = '\0'; + + if (ip == NULL) + { + ip = start; + continue; + } + + if (strcasecmp(name, start) == 0) + { + if (getaddrinfo(ip, port, hints, res) == 0) + return 0; + } + } + + return 1; +} + +static int __readaddrinfo(const char *filename, + const char *name, unsigned short port, + const struct addrinfo *hints, + struct addrinfo **res) +{ + char port_str[PORT_STR_MAX + 1]; + size_t bufsize = 0; + void *line = NULL; + int count = 0; + struct addrinfo h; + int errno_bak; + FILE *fp; + int ret; + + fp = fopen(filename, "r"); + if (!fp) + return EAI_SYSTEM; + + h = *hints; + h.ai_flags |= AI_NUMERICSERV | AI_NUMERICHOST, + snprintf(port_str, PORT_STR_MAX + 1, "%u", port); + + errno_bak = errno; + while ((ret = __read_line(fp, &line, &bufsize)) == 0) + { + if (__readaddrinfo_line((char*)line, name, port_str, &h, res) == 0) + { + count++; + res = &((*res)->ai_next); + } + } + + free(line); + fclose(fp); + if (count != 0) + { + errno = errno_bak; + return 0; + } + + return ret < 0 ? EAI_SYSTEM : EAI_NONAME; +} + +static ThreadDnsTask *__create_thread_dns_task(const std::string& host, + unsigned short port, + thread_dns_callback_t callback) +{ + auto *task = WFThreadTaskFactory:: + create_thread_task(WFGlobal::get_dns_queue(), + WFGlobal::get_dns_executor(), + DnsRoutine::run, + std::move(callback)); + + task->get_input()->reset(host, port); + return task; +} + +class WFResolverTask : public WFRouterTask +{ + static void dns_partial_callback(WFDnsTask *dns_task); + +public: + WFResolverTask(const struct WFNSParams *params, int dns_cache_level, + unsigned int dns_ttl_default, unsigned int dns_ttl_min, + const struct EndpointParams *endpoint_params, + router_callback_t&& cb) : + WFRouterTask(std::move(cb)), + type_(params->type), + host_(params->uri.host ? params->uri.host : ""), + port_(params->uri.port ? atoi(params->uri.port) : 0), + info_(params->info), + dns_cache_level_(dns_cache_level), + dns_ttl_default_(dns_ttl_default), + dns_ttl_min_(dns_ttl_min), + endpoint_params_(*endpoint_params), + first_addr_only_(params->fixed_addr) + { + } + +private: + virtual void dispatch(); + virtual SubTask *done(); + void dns_callback(ThreadDnsTask *dns_task); + void dns_single_callback(WFDnsTask *dns_task); + void dns_parallel_callback(const ParallelWork *pwork); + void dns_callback_internal(DnsOutput *dns_task, + unsigned int ttl_default, + unsigned int ttl_min); + +private: + TransportType type_; + std::string host_; + unsigned short port_; + std::string info_; + int dns_cache_level_; + unsigned int dns_ttl_default_; + unsigned int dns_ttl_min_; + struct EndpointParams endpoint_params_; + bool first_addr_only_; + bool insert_dns_; +}; + +void WFResolverTask::dispatch() +{ + insert_dns_ = true; + if (dns_cache_level_ != DNS_CACHE_LEVEL_0) + { + auto *dns_cache = WFGlobal::get_dns_cache(); + const DnsCache::DnsHandle *addr_handle = NULL; + + switch (dns_cache_level_) + { + case DNS_CACHE_LEVEL_1: + addr_handle = dns_cache->get_confident(host_, port_); + break; + + case DNS_CACHE_LEVEL_2: + addr_handle = dns_cache->get_ttl(host_, port_); + break; + + case DNS_CACHE_LEVEL_3: + addr_handle = dns_cache->get(host_, port_); + break; + + default: + break; + } + + if (addr_handle) + { + auto *route_manager = WFGlobal::get_route_manager(); + struct addrinfo *addrinfo = addr_handle->value.addrinfo; + struct addrinfo first; + + if (first_addr_only_ && addrinfo->ai_next) + { + first = *addrinfo; + first.ai_next = NULL; + addrinfo = &first; + } + + if (route_manager->get(type_, addrinfo, info_, &endpoint_params_, + host_, this->result) < 0) + { + this->state = WFT_STATE_SYS_ERROR; + this->error = errno; + } + else + this->state = WFT_STATE_SUCCESS; + + insert_dns_ = false; + dns_cache->release(addr_handle); + } + } + + if (insert_dns_ && !host_.empty()) + { + char front = host_.front(); + char back = host_.back(); + struct in6_addr addr; + int ret; + + if (host_.find(':') != std::string::npos) + ret = inet_pton(AF_INET6, host_.c_str(), &addr); + else if (isdigit(back) && isdigit(front)) + ret = inet_pton(AF_INET, host_.c_str(), &addr); + else if (front == '/') + ret = 1; + else + ret = 0; + + if (ret == 1) + { + DnsInput dns_in; + DnsOutput dns_out; + + dns_in.reset(host_, port_); + DnsRoutine::run(&dns_in, &dns_out); + dns_callback_internal(&dns_out, (unsigned int)-1, (unsigned int)-1); + insert_dns_ = false; + } + } + + static int family = __default_family(); + WFDnsClient *client = WFGlobal::get_dns_client(); + const char *hosts = WFGlobal::get_global_settings()->hosts_path; + + if (insert_dns_ && client && hosts) + { + struct addrinfo *ai; + int ret = __readaddrinfo(hosts, host_.c_str(), port_, &__ai_hints, &ai); + + if (ret == 0) + { + DnsOutput out; + DnsRoutine::create(&out, ret, ai); + dns_callback_internal(&out, dns_ttl_default_, dns_ttl_min_); + insert_dns_ = false; + } + } + + if (insert_dns_ && !client) + { + auto&& cb = std::bind(&WFResolverTask::dns_callback, + this, + std::placeholders::_1); + ThreadDnsTask *dns_task = __create_thread_dns_task(host_, port_, + std::move(cb)); + series_of(this)->push_front(dns_task); + } + else if (insert_dns_) + { + if (family == AF_INET || family == AF_INET6) + { + auto&& cb = std::bind(&WFResolverTask::dns_single_callback, + this, + std::placeholders::_1); + WFDnsTask *dns_task = client->create_dns_task(host_, std::move(cb)); + + if (family == AF_INET6) + dns_task->get_req()->set_question_type(DNS_TYPE_AAAA); + + series_of(this)->push_front(dns_task); + } + else + { + struct DnsContext *dctx = new struct DnsContext[2]; + WFDnsTask *task_v4; + WFDnsTask *task_v6; + ParallelWork *pwork; + + dctx[0].ai = NULL; + dctx[1].ai = NULL; + dctx[0].port = port_; + dctx[1].port = port_; + + task_v4 = client->create_dns_task(host_, dns_partial_callback); + task_v4->user_data = dctx; + + task_v6 = client->create_dns_task(host_, dns_partial_callback); + task_v6->get_req()->set_question_type(DNS_TYPE_AAAA); + task_v6->user_data = dctx + 1; + + auto&& cb = std::bind(&WFResolverTask::dns_parallel_callback, + this, + std::placeholders::_1); + + pwork = Workflow::create_parallel_work(std::move(cb)); + pwork->set_context(dctx); + pwork->add_series(Workflow::create_series_work(task_v4, nullptr)); + pwork->add_series(Workflow::create_series_work(task_v6, nullptr)); + + series_of(this)->push_front(pwork); + } + } + + this->subtask_done(); +} + +SubTask *WFResolverTask::done() +{ + SeriesWork *series = series_of(this); + + if (!insert_dns_) + { + if (this->callback) + this->callback(this); + + delete this; + } + + return series->pop(); +} + +void WFResolverTask::dns_callback_internal(DnsOutput *dns_out, + unsigned int ttl_default, + unsigned int ttl_min) +{ + int dns_error = dns_out->get_error(); + + if (dns_error) + { + if (dns_error == EAI_SYSTEM) + { + this->state = WFT_STATE_SYS_ERROR; + this->error = errno; + } + else + { + this->state = WFT_STATE_DNS_ERROR; + this->error = dns_error; + } + } + else + { + auto *route_manager = WFGlobal::get_route_manager(); + auto *dns_cache = WFGlobal::get_dns_cache(); + struct addrinfo *addrinfo = dns_out->move_addrinfo(); + const DnsCache::DnsHandle *addr_handle; + + addr_handle = dns_cache->put(host_, port_, addrinfo, + (unsigned int)ttl_default, + (unsigned int)ttl_min); + if (route_manager->get(type_, addrinfo, info_, &endpoint_params_, + host_, this->result) < 0) + { + this->state = WFT_STATE_SYS_ERROR; + this->error = errno; + } + else + this->state = WFT_STATE_SUCCESS; + + dns_cache->release(addr_handle); + } +} + +void WFResolverTask::dns_single_callback(WFDnsTask *dns_task) +{ + if (dns_task->get_state() == WFT_STATE_SUCCESS) + { + struct addrinfo *ai = NULL; + int ret = DnsUtil::getaddrinfo(dns_task->get_resp(), port_, &ai); + DnsOutput out; + DnsRoutine::create(&out, ret, ai); + dns_callback_internal(&out, dns_ttl_default_, dns_ttl_min_); + } + else + { + this->state = dns_task->get_state(); + this->error = dns_task->get_error(); + } + + if (this->callback) + this->callback(this); + + delete this; +} + +void WFResolverTask::dns_partial_callback(WFDnsTask *dns_task) +{ + struct DnsContext *ctx = (struct DnsContext *)dns_task->user_data; + + ctx->state = dns_task->get_state(); + ctx->error = dns_task->get_error(); + if (ctx->state == WFT_STATE_SUCCESS) + { + auto *resp = dns_task->get_resp(); + ctx->ai = NULL; + ctx->eai_error = DnsUtil::getaddrinfo(resp, ctx->port, &ctx->ai); + } + else + { + ctx->eai_error = EAI_NONAME; + ctx->ai = NULL; + } +} + +void WFResolverTask::dns_parallel_callback(const ParallelWork *pwork) +{ + struct DnsContext *c4 = (struct DnsContext *)(pwork->get_context()); + struct DnsContext *c6 = c4 + 1; + DnsOutput out; + + if (c4->state != WFT_STATE_SUCCESS && c6->state != WFT_STATE_SUCCESS) + { + this->state = c4->state; + this->error = c4->error; + } + else if (c4->eai_error != 0 && c6->eai_error != 0) + { + DnsRoutine::create(&out, c4->eai_error, NULL); + dns_callback_internal(&out, dns_ttl_default_, dns_ttl_min_); + } + else + { + struct addrinfo *ai = NULL; + struct addrinfo **pai = &ai; + + if (c4->ai != NULL) + { + *pai = c4->ai; + while (*pai) + pai = &(*pai)->ai_next; + } + + if (c6->ai != NULL) + *pai = c6->ai; + + DnsRoutine::create(&out, 0, ai); + dns_callback_internal(&out, dns_ttl_default_, dns_ttl_min_); + } + + delete[] c4; + + if (this->callback) + this->callback(this); + + delete this; +} + +void WFResolverTask::dns_callback(ThreadDnsTask *dns_task) +{ + if (dns_task->get_state() == WFT_STATE_SUCCESS) + dns_callback_internal(dns_task->get_output(), dns_ttl_default_, dns_ttl_min_); + else + { + this->state = dns_task->get_state(); + this->error = dns_task->get_error(); + } + + if (this->callback) + this->callback(this); + + delete this; +} + +WFRouterTask * +WFDnsResolver::create(const struct WFNSParams *params, int dns_cache_level, + unsigned int dns_ttl_default, unsigned int dns_ttl_min, + const struct EndpointParams *endpoint_params, + router_callback_t&& callback) +{ + return new WFResolverTask(params, dns_cache_level, + dns_ttl_default, dns_ttl_min, + endpoint_params, std::move(callback)); +} + +WFRouterTask *WFDnsResolver::create_router_task(const struct WFNSParams *params, + router_callback_t callback) +{ + const auto *settings = WFGlobal::get_global_settings(); + unsigned int dns_ttl_default = settings->dns_ttl_default; + unsigned int dns_ttl_min = settings->dns_ttl_min; + const struct EndpointParams *endpoint_params = &settings->endpoint_params; + int dns_cache_level = params->retry_times == 0 ? DNS_CACHE_LEVEL_2 : + DNS_CACHE_LEVEL_1; + return create(params, dns_cache_level, dns_ttl_default, dns_ttl_min, + endpoint_params, std::move(callback)); +} + diff --git a/src/nameservice/WFDNSResolver.h b/src/nameservice/WFDnsResolver.h similarity index 96% rename from src/nameservice/WFDNSResolver.h rename to src/nameservice/WFDnsResolver.h index 2419abb4..e8574c84 100644 --- a/src/nameservice/WFDNSResolver.h +++ b/src/nameservice/WFDnsResolver.h @@ -22,7 +22,7 @@ #include "EndpointParams.h" #include "WFNameService.h" -class WFDNSResolver : public WFNSPolicy +class WFDnsResolver : public WFNSPolicy { public: virtual WFRouterTask *create_router_task(const struct WFNSParams *params, diff --git a/src/nameservice/WFServiceGovernance.cc b/src/nameservice/WFServiceGovernance.cc index 8129f525..76de7e4e 100644 --- a/src/nameservice/WFServiceGovernance.cc +++ b/src/nameservice/WFServiceGovernance.cc @@ -20,7 +20,7 @@ #include "URIParser.h" #include "StringUtil.h" #include "WFNameService.h" -#include "WFDNSResolver.h" +#include "WFDnsResolver.h" #include "WFServiceGovernance.h" #include "UpstreamManager.h" @@ -230,7 +230,7 @@ void WFServiceGovernance::success(RouteManager::RouteResult *result, this->recover_server_from_breaker(server); pthread_rwlock_unlock(&this->rwlock); - WFDNSResolver::success(result, tracing, target); + WFDnsResolver::success(result, tracing, target); } void WFServiceGovernance::failed(RouteManager::RouteResult *result, @@ -253,7 +253,7 @@ void WFServiceGovernance::failed(RouteManager::RouteResult *result, pthread_rwlock_unlock(&this->rwlock); - WFDNSResolver::failed(result, tracing, target); + WFDnsResolver::failed(result, tracing, target); } void WFServiceGovernance::check_breaker() diff --git a/src/nameservice/WFServiceGovernance.h b/src/nameservice/WFServiceGovernance.h index fde24b4a..08448fbe 100644 --- a/src/nameservice/WFServiceGovernance.h +++ b/src/nameservice/WFServiceGovernance.h @@ -26,7 +26,7 @@ #include "URIParser.h" #include "EndpointParams.h" #include "WFNameService.h" -#include "WFDNSResolver.h" +#include "WFDnsResolver.h" #include "WFGlobal.h" #include "WFTaskError.h" @@ -102,7 +102,7 @@ public: virtual ~EndpointAddress() { delete this->params; } }; -class WFServiceGovernance : public WFDNSResolver +class WFServiceGovernance : public WFDnsResolver { public: virtual WFRouterTask *create_router_task(const struct WFNSParams *params, diff --git a/src/protocol/CMakeLists.txt b/src/protocol/CMakeLists.txt index a476c510..734d7c2b 100644 --- a/src/protocol/CMakeLists.txt +++ b/src/protocol/CMakeLists.txt @@ -2,11 +2,13 @@ cmake_minimum_required(VERSION 3.6) project(protocol) set(SRC + dns_parser.c http_parser.c redis_parser.c mysql_stream.c mysql_parser.c mysql_byteorder.c + DnsMessage.cc MySQLMessage.cc MySQLResult.cc HttpMessage.cc diff --git a/src/protocol/DnsMessage.cc b/src/protocol/DnsMessage.cc new file mode 100644 index 00000000..187e3ae6 --- /dev/null +++ b/src/protocol/DnsMessage.cc @@ -0,0 +1,203 @@ +/* + 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 "DnsMessage.h" + +#define DNS_LABELS_MAX 63 +#define DNS_MESSAGE_MAX_UDP_SIZE 512 + +namespace protocol +{ + +static inline void __append_uint8(std::string& s, uint8_t tmp) +{ + s.append((const char *)&tmp, sizeof (uint8_t)); +} + +static inline void __append_uint16(std::string& s, uint16_t tmp) +{ + tmp = htons(tmp); + s.append((const char *)&tmp, sizeof (uint16_t)); +} + +DnsMessage::DnsMessage(DnsMessage&& msg) : + ProtocolMessage(std::move(msg)) +{ + this->parser = msg.parser; + msg.parser = NULL; + + this->cur_size = msg.cur_size; + msg.cur_size = 0; +} + +DnsMessage& DnsMessage::operator = (DnsMessage&& msg) +{ + if (&msg != this) + { + *(ProtocolMessage *)this = std::move(msg); + + if (this->parser) + { + dns_parser_deinit(this->parser); + delete this->parser; + } + + this->parser = msg.parser; + msg.parser = NULL; + + this->cur_size = msg.cur_size; + msg.cur_size = 0; + } + return *this; +} + +int DnsMessage::encode_reply() +{ + struct dns_header h; + const char *name; + const char *p; + size_t len; + + msgbuf.clear(); + msgbuf.reserve(DNS_HEADER_SIZE); + msgsize = 0; + + // TODO encode other field + // pointers can only be used for occurances of a domain name where + // the format is not class specific + h = this->parser->header; + h.id = htons(h.id); + h.qdcount = htons(1); + h.ancount = htons(0); + h.nscount = htons(0); + h.arcount = htons(0); + + msgbuf.append((const char *)&h, sizeof (struct dns_header)); + p = parser->question.qname; + while (*p) + { + name = p; + while (*p && *p != '.') + p++; + + len = p - name; + if (len > DNS_LABELS_MAX || (len == 0 && *p && *(p + 1))) + { + errno = EINVAL; + return -1; + } + + if (len > 0) + { + __append_uint8(msgbuf, len); + msgbuf.append(name, len); + } + + if (*p == '.') + p++; + } + + len = 0; + __append_uint8(msgbuf, len); + __append_uint16(msgbuf, parser->question.qtype); + __append_uint16(msgbuf, parser->question.qclass); + + if (msgbuf.size() >= (1 << 16)) + { + errno = EOVERFLOW; + return -1; + } + + msgsize = htons(msgbuf.size()); + + return 0; +} + +int DnsMessage::encode(struct iovec vectors[], int) +{ + struct iovec *p = vectors; + + if (this->encode_reply() < 0) + return -1; + + // TODO + // if this is a request, it won't exceed the 512 bytes UDP limit + // if this is a response and exceed 512 bytes, we need a TrunCation reply + + if (this->has_leading_length()) + { + p->iov_base = &this->msgsize; + p->iov_len = sizeof (uint16_t); + p++; + } + + p->iov_base = (void *)this->msgbuf.data(); + p->iov_len = msgbuf.size(); + return p - vectors + 1; +} + +int DnsMessage::append(const void *buf, size_t *size) +{ + int ret = dns_parser_append_message(buf, size, this->parser); + + if (ret >= 0) + { + this->cur_size += *size; + if (this->cur_size > this->size_limit) + { + errno = EMSGSIZE; + ret = -1; + } + } + else if (ret == -2) + { + errno = EBADMSG; + ret = -1; + } + + return ret; +} + +std::atomic DnsRequest::req_id_; + +int DnsResponse::append(const void *buf, size_t *size) +{ + int ret = this->DnsMessage::append(buf, size); + const char *qname = this->parser->question.qname; + if (ret >= 1 && (this->request_id != this->get_id() || + strcasecmp(this->request_name.c_str(), qname) != 0)) + { + if (this->has_leading_length()) + { + errno = EBADMSG; + ret = -1; + } + else + { + dns_parser_deinit(this->parser); + dns_parser_init(this->parser); + ret = 0; + } + } + + return ret; +} + +} diff --git a/src/protocol/DnsMessage.h b/src/protocol/DnsMessage.h new file mode 100644 index 00000000..a8d079b7 --- /dev/null +++ b/src/protocol/DnsMessage.h @@ -0,0 +1,252 @@ +/* + 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) +*/ + +#ifndef _DNSMESSAGE_H_ +#define _DNSMESSAGE_H_ + +/** + * @file DnsMessage.h + * @brief Dns Protocol Interface + */ + +#include +#include +#include +#include "ProtocolMessage.h" +#include "dns_parser.h" + +namespace protocol +{ + +class DnsMessage : public ProtocolMessage +{ +protected: + virtual int encode(struct iovec vectors[], int max); + virtual int append(const void *buf, size_t *size); + +public: + int get_id() const + { + return parser->header.id; + } + int get_qr() const + { + return parser->header.qr; + } + int get_opcode() const + { + return parser->header.opcode; + } + int get_aa() const + { + return parser->header.aa; + } + int get_tc() const + { + return parser->header.tc; + } + int get_rd() const + { + return parser->header.rd; + } + int get_ra() const + { + return parser->header.ra; + } + int get_rcode() const + { + return parser->header.rcode; + } + int get_qdcount() const + { + return parser->header.qdcount; + } + int get_ancount() const + { + return parser->header.ancount; + } + int get_nscount() const + { + return parser->header.nscount; + } + int get_arcount() const + { + return parser->header.arcount; + } + + void set_id(int id) + { + parser->header.id = id; + } + void set_qr(int qr) + { + parser->header.qr = qr; + } + void set_opcode(int opcode) + { + parser->header.opcode = opcode; + } + void set_aa(int aa) + { + parser->header.aa = aa; + } + void set_tc(int tc) + { + parser->header.tc = tc; + } + void set_rd(int rd) + { + parser->header.rd = rd; + } + void set_ra(int ra) + { + parser->header.ra = ra; + } + void set_rcode(int rcode) + { + parser->header.rcode = rcode; + } + + int get_question_type() const + { + return parser->question.qtype; + } + int get_question_class() const + { + return parser->question.qclass; + } + std::string get_question_name() const + { + const char *name = parser->question.qname; + if (name == NULL) + return ""; + return name; + } + void set_question_type(int qtype) + { + parser->question.qtype = qtype; + } + void set_question_class(int qclass) + { + parser->question.qclass = qclass; + } + void set_question_name(const std::string& name) + { + char *pname = parser->question.qname; + if (pname != NULL) + free(pname); + parser->question.qname = strdup(name.c_str()); + } + + // Inner use only + bool has_leading_length() const + { + return parser->leading_length; + } + void set_leading_length(bool leading_length) + { + parser->leading_length = leading_length; + } + +public: + DnsMessage() : + parser(new dns_parser_t) + { + dns_parser_init(parser); + this->cur_size = 0; + } + + virtual ~DnsMessage() + { + if (this->parser) + { + dns_parser_deinit(parser); + delete this->parser; + } + } + + DnsMessage(DnsMessage&& msg); + DnsMessage& operator = (DnsMessage&& msg); + +protected: + dns_parser_t *parser; + std::string msgbuf; + size_t cur_size; + +private: + int encode_reply(); + int encode_truncation_reply(); + + // size of msgbuf, but in network byte order + uint16_t msgsize; +}; + +class DnsRequest : public DnsMessage +{ + static std::atomic req_id_; +public: + DnsRequest() + { + dns_parser_set_id(req_id_++, this->parser); + } + + DnsRequest(DnsRequest&& req) = default; + DnsRequest& operator = (DnsRequest&& req) = default; + + void set_question(const char *host, uint16_t qtype, uint16_t qclass) + { + dns_parser_set_question(host, qtype, qclass, this->parser); + } +}; + +class DnsResponse : public DnsMessage +{ +public: + DnsResponse() + { + this->request_id = 0; + } + + DnsResponse(DnsResponse&& req) = default; + DnsResponse& operator = (DnsResponse&& req) = default; + + const dns_parser_t *get_parser() const + { + return this->parser; + } + + int append(const void *buf, size_t *size); + + void set_request_id(uint16_t id) + { + this->request_id = id; + } + + void set_request_name(const std::string& name) + { + this->request_name = name; + } + +private: + uint16_t request_id; + std::string request_name; +}; + +} + +#endif diff --git a/src/protocol/dns_parser.c b/src/protocol/dns_parser.c new file mode 100644 index 00000000..8261d3c8 --- /dev/null +++ b/src/protocol/dns_parser.c @@ -0,0 +1,1018 @@ +/* + 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 "dns_parser.h" + +#define DNS_LABELS_MAX 63 +#define DNS_NAMES_MAX 256 +#define DNS_MSGBASE_INIT_SIZE 514 // 512 + 2(leading length) +#define MAX(x, y) ((x) <= (y) ? (y) : (x)) + +struct __dns_record_entry +{ + struct list_head entry_list; + struct dns_record record; +}; + + +static uint8_t __dns_parser_uint8(const char *ptr) +{ + return (unsigned char)ptr[0]; +} + +static uint16_t __dns_parser_uint16(const char *ptr) +{ + const unsigned char *p = (const unsigned char *)ptr; + return ((uint16_t)p[0] << 8) + + ((uint16_t)p[1]); +} + +static uint32_t __dns_parser_uint32(const char *ptr) +{ + const unsigned char *p = (const unsigned char *)ptr; + return ((uint32_t)p[0] << 24) + + ((uint32_t)p[1] << 16) + + ((uint32_t)p[2] << 8) + + ((uint32_t)p[3]); +} + +/* + * Parse a single . + * is a domain name represented as a series of labels, and + * terminated by a label with zero length. + * + * phost must point to an char array with at least DNS_NAMES_MAX+1 size + */ +static int __dns_parser_parse_host(char *phost, dns_parser_t *parser) +{ + uint8_t len; + uint16_t pointer; + size_t hcur; + const char *msgend; + const char **cur; + const char *curbackup; // backup cur when host label is pointer + + msgend = (const char *)parser->msgbuf + parser->msgsize; + cur = &(parser->cur); + curbackup = NULL; + hcur = 0; + + if (*cur >= msgend) + return -2; + + while (*cur < msgend) + { + len = __dns_parser_uint8(*cur); + + if ((len & 0xC0) == 0) + { + (*cur)++; + if (len == 0) + break; + if (len > DNS_LABELS_MAX || *cur + len > msgend || + hcur + len + 1 > DNS_NAMES_MAX) + return -2; + + memcpy(phost + hcur, *cur, len); + *cur += len; + hcur += len; + phost[hcur++] = '.'; + } + // RFC 1035, 4.1.4 Message compression + else if ((len & 0xC0) == 0xC0) + { + pointer = __dns_parser_uint16(*cur) & 0x3FFF; + *cur += 2; + + if (pointer >= parser->msgsize) + return -2; + + // backup cur only when the first pointer occurs + if (curbackup == NULL) + curbackup = *cur; + *cur = (const char *)parser->msgbase + pointer; + } + else + return -2; + } + if (curbackup != NULL) + *cur = curbackup; + + if (hcur > 1 && phost[hcur - 1] == '.') + hcur--; + + if (hcur == 0) + phost[hcur++] = '.'; + phost[hcur++] = '\0'; + + return 0; +} + +static void __dns_parser_free_record(struct __dns_record_entry *r) +{ + switch (r->record.type) + { + case DNS_TYPE_SOA: + { + struct dns_record_soa *soa; + soa = (struct dns_record_soa *)(r->record.rdata); + free(soa->mname); + free(soa->rname); + break; + } + case DNS_TYPE_SRV: + { + struct dns_record_srv *srv; + srv = (struct dns_record_srv *)(r->record.rdata); + free(srv->target); + break; + } + case DNS_TYPE_MX: + { + struct dns_record_mx *mx; + mx = (struct dns_record_mx *)(r->record.rdata); + free(mx->exchange); + break; + } + } + free(r->record.name); + free(r); +} + +static void __dns_parser_free_record_list(struct list_head *head) +{ + struct list_head *pos, *tmp; + struct __dns_record_entry *entry; + + list_for_each_safe(pos, tmp, head) + { + entry = list_entry(pos, struct __dns_record_entry, entry_list); + list_del(pos); + __dns_parser_free_record(entry); + } +} + +/* + * A RDATA format, from RFC 1035: + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | ADDRESS | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * + * ADDRESS: A 32 bit Internet address. + * Hosts that have multiple Internet addresses will have multiple A records. + */ +static int __dns_parser_parse_a(struct __dns_record_entry **r, + uint16_t rdlength, + dns_parser_t *parser) +{ + const char **cur; + struct __dns_record_entry *entry; + size_t entry_size; + + if (sizeof (struct in_addr) != rdlength) + return -2; + + cur = &(parser->cur); + entry_size = sizeof (struct __dns_record_entry) + sizeof (struct in_addr); + entry = (struct __dns_record_entry *)malloc(entry_size); + if (!entry) + return -1; + + entry->record.rdata = (void *)(entry + 1); + memcpy(entry->record.rdata, *cur, rdlength); + *cur += rdlength; + *r = entry; + + return 0; +} + +/* + * AAAA RDATA format, from RFC 3596: + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | ADDRESS | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * + * ADDRESS: A 128 bit Internet address. + * Hosts that have multiple addresses will have multiple AAAA records. + */ +static int __dns_parser_parse_aaaa(struct __dns_record_entry **r, + uint16_t rdlength, + dns_parser_t *parser) +{ + const char **cur; + struct __dns_record_entry *entry; + size_t entry_size; + + if (sizeof (struct in6_addr) != rdlength) + return -2; + + cur = &(parser->cur); + entry_size = sizeof (struct __dns_record_entry) + sizeof (struct in6_addr); + entry = (struct __dns_record_entry *)malloc(entry_size); + if (!entry) + return -1; + + entry->record.rdata = (void *)(entry + 1); + memcpy(entry->record.rdata, *cur, rdlength); + *cur += rdlength; + *r = entry; + + return 0; +} + +/* + * Parse any record. + */ +static int __dns_parser_parse_names(struct __dns_record_entry **r, + uint16_t rdlength, + dns_parser_t *parser) +{ + const char *rcdend; + const char **cur; + struct __dns_record_entry *entry; + size_t entry_size; + size_t name_len; + char name[DNS_NAMES_MAX + 2]; + int ret; + + cur = &(parser->cur); + rcdend = *cur + rdlength; + ret = __dns_parser_parse_host(name, parser); + if (ret < 0) + return ret; + + if (*cur != rcdend) + return -2; + + name_len = strlen(name); + entry_size = sizeof (struct __dns_record_entry) + name_len + 1; + entry = (struct __dns_record_entry *)malloc(entry_size); + if (!entry) + return -1; + + entry->record.rdata = (void *)(entry + 1); + memcpy(entry->record.rdata, name, name_len + 1); + *r = entry; + + return 0; +} + +/* + * SOA RDATA format, from RFC 1035: + * 1 1 1 1 1 1 + * 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * / MNAME / + * / / + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * / RNAME / + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | SERIAL | + * | | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | REFRESH | + * | | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | RETRY | + * | | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | EXPIRE | + * | | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | MINIMUM | + * | | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * + * MNAME: + * RNAME: + * SERIAL: The unsigned 32 bit version number. + * REFRESH: A 32 bit time interval. + * RETRY: A 32 bit time interval. + * EXPIRE: A 32 bit time value. + * MINIMUM: The unsigned 32 bit integer. + */ +static int __dns_parser_parse_soa(struct __dns_record_entry **r, + uint16_t rdlength, + dns_parser_t *parser) +{ + const char *rcdend; + const char **cur; + struct __dns_record_entry *entry; + struct dns_record_soa *soa; + size_t entry_size; + char mname[DNS_NAMES_MAX + 2]; + char rname[DNS_NAMES_MAX + 2]; + int ret; + + cur = &(parser->cur); + rcdend = *cur + rdlength; + ret = __dns_parser_parse_host(mname, parser); + if (ret < 0) + return ret; + ret = __dns_parser_parse_host(rname, parser); + if (ret < 0) + return ret; + + if (*cur + 20 != rcdend) + return -2; + + entry_size = sizeof (struct __dns_record_entry) + + sizeof (struct dns_record_soa); + entry = (struct __dns_record_entry *)malloc(entry_size); + if (!entry) + return -1; + + entry->record.rdata = (void *)(entry + 1); + soa = (struct dns_record_soa *)(entry->record.rdata); + + soa->mname = strdup(mname); + soa->rname = strdup(rname); + soa->serial = __dns_parser_uint32(*cur); + soa->refresh = __dns_parser_uint32(*cur + 4); + soa->retry = __dns_parser_uint32(*cur + 8); + soa->expire = __dns_parser_uint32(*cur + 12); + soa->minimum = __dns_parser_uint32(*cur + 16); + + if (!soa->mname || !soa->rname) + { + free(soa->mname); + free(soa->rname); + free(entry); + return -1; + } + + *cur += 20; + *r = entry; + + return 0; +} + +/* + * SRV RDATA format, from RFC 2782: + * 1 1 1 1 1 1 + * 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | PRIORITY | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | WEIGHT | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | PORT | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * / TARGET / + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * + * PRIORITY: A 16 bit unsigned integer in network byte order. + * WEIGHT: A 16 bit unsigned integer in network byte order. + * PORT: A 16 bit unsigned integer in network byte order. + * TARGET: + */ +static int __dns_parser_parse_srv(struct __dns_record_entry **r, + uint16_t rdlength, + dns_parser_t *parser) +{ + const char *rcdend; + const char **cur; + struct __dns_record_entry *entry; + struct dns_record_srv *srv; + size_t entry_size; + char target[DNS_NAMES_MAX + 2]; + uint16_t priority; + uint16_t weight; + uint16_t port; + int ret; + + cur = &(parser->cur); + rcdend = *cur + rdlength; + + if (*cur + 6 > rcdend) + return -2; + + priority = __dns_parser_uint16(*cur); + weight = __dns_parser_uint16(*cur + 2); + port = __dns_parser_uint16(*cur + 4); + *cur += 6; + + ret = __dns_parser_parse_host(target, parser); + if (ret < 0) + return ret; + if (*cur != rcdend) + return -2; + + entry_size = sizeof (struct __dns_record_entry) + + sizeof (struct dns_record_srv); + entry = (struct __dns_record_entry *)malloc(entry_size); + if (!entry) + return -1; + + entry->record.rdata = (void *)(entry + 1); + srv = (struct dns_record_srv *)(entry->record.rdata); + + srv->priority = priority; + srv->weight = weight; + srv->port = port; + srv->target = strdup(target); + + if (!srv->target) + { + free(entry); + return -1; + } + + *r = entry; + + return 0; +} + +static int __dns_parser_parse_mx(struct __dns_record_entry **r, + uint16_t rdlength, + dns_parser_t *parser) +{ + const char *rcdend; + const char **cur; + struct __dns_record_entry *entry; + struct dns_record_mx *mx; + size_t entry_size; + char exchange[DNS_NAMES_MAX + 2]; + int16_t preference; + int ret; + + cur = &(parser->cur); + rcdend = *cur + rdlength; + + if (*cur + 2 > rcdend) + return -2; + preference = __dns_parser_uint16(*cur); + *cur += 2; + + ret = __dns_parser_parse_host(exchange, parser); + if (ret < 0) + return ret; + if (*cur != rcdend) + return -2; + + entry_size = sizeof (struct __dns_record_entry) + + sizeof (struct dns_record_mx); + entry = (struct __dns_record_entry *)malloc(entry_size); + if (!entry) + return -1; + + entry->record.rdata = (void *)(entry + 1); + mx = (struct dns_record_mx *)(entry->record.rdata); + mx->exchange = strdup(exchange); + mx->preference = preference; + + if (!mx->exchange) + { + free(entry); + return -1; + } + + *r = entry; + + return 0; +} + +static int __dns_parser_parse_others(struct __dns_record_entry **r, + uint16_t rdlength, + dns_parser_t *parser) +{ + const char **cur; + struct __dns_record_entry *entry; + size_t entry_size; + + cur = &(parser->cur); + entry_size = sizeof (struct __dns_record_entry) + rdlength; + entry = (struct __dns_record_entry *)malloc(entry_size); + if (!entry) + return -1; + + entry->record.rdata = (void *)(entry + 1); + memcpy(entry->record.rdata, *cur, rdlength); + *cur += rdlength; + *r = entry; + + return 0; +} + +/* + * RR format, from RFC 1035: + * 1 1 1 1 1 1 + * 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | | + * / NAME / + * / / + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | TYPE | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | CLASS | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | TTL | + * | | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | RDLENGTH | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * / RDATA / + * / / + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * + */ +int __dns_parser_parse_record(int idx, dns_parser_t *parser) +{ + uint16_t i; + uint16_t type; + uint16_t rclass; + uint32_t ttl; + uint16_t rdlength; + uint16_t count; + const char *msgend; + const char **cur; + int ret; + struct __dns_record_entry *entry; + char host[DNS_NAMES_MAX + 2]; + struct list_head *list; + + switch (idx) + { + case 0: + count = parser->header.ancount; + list = &parser->answer_list; + break; + case 1: + count = parser->header.nscount; + list = &parser->authority_list; + break; + case 2: + count = parser->header.arcount; + list = &parser->additional_list; + break; + default: + return -2; + } + + msgend = (const char *)parser->msgbuf + parser->msgsize; + cur = &(parser->cur); + + for (i = 0; i < count; i++) + { + ret = __dns_parser_parse_host(host, parser); + if (ret < 0) + return ret; + + // TYPE(2) + CLASS(2) + TTL(4) + RDLENGTH(2) = 10 + if (*cur + 10 > msgend) + return -2; + type = __dns_parser_uint16(*cur); + rclass = __dns_parser_uint16(*cur + 2); + ttl = __dns_parser_uint32(*cur + 4); + rdlength = __dns_parser_uint16(*cur + 8); + *cur += 10; + if (*cur + rdlength > msgend) + return -2; + + entry = NULL; + switch (type) + { + case DNS_TYPE_A: + ret = __dns_parser_parse_a(&entry, rdlength, parser); + break; + case DNS_TYPE_AAAA: + ret = __dns_parser_parse_aaaa(&entry, rdlength, parser); + break; + case DNS_TYPE_NS: + case DNS_TYPE_CNAME: + case DNS_TYPE_PTR: + ret = __dns_parser_parse_names(&entry, rdlength, parser); + break; + case DNS_TYPE_SOA: + ret = __dns_parser_parse_soa(&entry, rdlength, parser); + break; + case DNS_TYPE_SRV: + ret = __dns_parser_parse_srv(&entry, rdlength, parser); + break; + case DNS_TYPE_MX: + ret = __dns_parser_parse_mx(&entry, rdlength, parser); + break; + default: + ret = __dns_parser_parse_others(&entry, rdlength, parser); + } + + if (ret < 0) + return ret; + + entry->record.name = strdup(host); + if (!entry->record.name) + { + __dns_parser_free_record(entry); + return -1; + } + + entry->record.type = type; + entry->record.rclass = rclass; + entry->record.ttl = ttl; + entry->record.rdlength = rdlength; + list_add_tail(&entry->entry_list, list); + } + + return 0; +} + +/* + * Question format, from RFC 1035: + * 1 1 1 1 1 1 + * 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | | + * / QNAME / + * / / + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | QTYPE | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * | QCLASS | + * +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ + * + * The query name is encoded as a series of labels, each represented + * as a one-byte length (maximum 63) followed by the text of the + * label. The list is terminated by a label of length zero (which can + * be thought of as the root domain). + */ +int __dns_parser_parse_question(dns_parser_t *parser) +{ + uint16_t qtype; + uint16_t qclass; + const char *msgend; + const char **cur; + int ret; + char host[DNS_NAMES_MAX + 2]; + + msgend = (const char *)parser->msgbuf + parser->msgsize; + cur = &(parser->cur); + + // question count != 1 is an error + if (parser->header.qdcount != 1) + return -2; + + // parse qname + ret = __dns_parser_parse_host(host, parser); + if (ret < 0) + return ret; + + // parse qtype and qclass + if (*cur + 4 > msgend) + return -2; + + qtype = __dns_parser_uint16(*cur); + qclass = __dns_parser_uint16(*cur + 2); + *cur += 4; + + if (parser->question.qname) + free(parser->question.qname); + + parser->question.qname = strdup(host); + if (parser->question.qname == NULL) + return -1; + + parser->question.qtype = qtype; + parser->question.qclass = qclass; + + return 0; +} + +void dns_parser_init(dns_parser_t *parser) +{ + parser->msgbuf = NULL; + parser->msgbase = NULL; + parser->cur = NULL; + parser->msgsize = 0; + parser->bufsize = 0; + parser->complete = 0; + parser->leading_length = 0; + memset(&parser->header, 0, DNS_HEADER_SIZE); + memset(&parser->question, 0, sizeof (struct dns_question)); + INIT_LIST_HEAD(&parser->answer_list); + INIT_LIST_HEAD(&parser->authority_list); + INIT_LIST_HEAD(&parser->additional_list); +} + +int dns_parser_set_question(const char *name, + uint16_t qtype, + uint16_t qclass, + dns_parser_t *parser) +{ + int ret; + + ret = dns_parser_set_question_name(name, parser); + if (ret < 0) + return ret; + + parser->question.qtype = qtype; + parser->question.qclass = qclass; + parser->header.qdcount = 1; + + return 0; +} + +int dns_parser_set_question_name(const char *name, dns_parser_t *parser) +{ + char *newname; + size_t len; + + len = strlen(name); + newname = (char *)malloc(len + 1); + + if (!newname) + return -1; + + memcpy(newname, name, len + 1); + // Remove trailing dot, except name is "." + if (len > 1 && newname[len - 1] == '.') + newname[len - 1] = '\0'; + + if (parser->question.qname) + free(parser->question.qname); + parser->question.qname = newname; + + return 0; +} + +void dns_parser_set_id(uint16_t id, dns_parser_t *parser) +{ + parser->header.id = id; +} + +int dns_parser_parse_all(dns_parser_t *parser) +{ + struct dns_header *h; + int ret; + int i; + + parser->complete = 1; + parser->cur = (const char *)parser->msgbase; + h = &parser->header; + + if (parser->msgsize < DNS_HEADER_SIZE) + return -2; + + memcpy(h, parser->msgbase, DNS_HEADER_SIZE); + h->id = ntohs(h->id); + h->qdcount = ntohs(h->qdcount); + h->ancount = ntohs(h->ancount); + h->nscount = ntohs(h->nscount); + h->arcount = ntohs(h->arcount); + parser->cur += DNS_HEADER_SIZE; + + ret = __dns_parser_parse_question(parser); + if (ret < 0) + return ret; + + for (i = 0; i < 3; i++) + { + ret = __dns_parser_parse_record(i, parser); + if (ret < 0) + return ret; + } + + return 0; +} + +int dns_parser_append_message(const void *buf, + size_t *n, + dns_parser_t *parser) +{ + int ret; + size_t total; + size_t new_size; + size_t msgsize_bak; + void *new_buf; + + if (parser->complete) + { + *n = 0; + return 1; + } + + if (parser->leading_length) + { + msgsize_bak = parser->msgsize; + if (parser->msgsize + *n > parser->bufsize) + { + new_size = MAX(DNS_MSGBASE_INIT_SIZE, 2 * parser->bufsize); + + while (new_size < parser->msgsize + *n) + new_size *= 2; + + new_buf = realloc(parser->msgbuf, new_size); + if (!new_buf) + return -1; + + parser->msgbuf = new_buf; + parser->bufsize = new_size; + } + + memcpy((char*)parser->msgbuf + parser->msgsize, buf, *n); + parser->msgsize += *n; + + if (parser->msgsize < 2) + return 0; + + total = __dns_parser_uint16((char*)parser->msgbuf); + if (parser->msgsize < total + 2) + return 0; + + *n = total + 2 - msgsize_bak; + parser->msgsize = total + 2; + parser->msgbase = (char*)parser->msgbuf + 2; + } + else + { + parser->msgbuf = malloc(*n); + memcpy(parser->msgbuf, buf, *n); + parser->msgbase = parser->msgbuf; + parser->msgsize = *n; + parser->bufsize = *n; + } + + ret = dns_parser_parse_all(parser); + if (ret < 0) + return ret; + + return 1; +} + +void dns_parser_deinit(dns_parser_t *parser) +{ + free(parser->msgbuf); + free(parser->question.qname); + + __dns_parser_free_record_list(&parser->answer_list); + __dns_parser_free_record_list(&parser->authority_list); + __dns_parser_free_record_list(&parser->additional_list); +} + +int dns_record_cursor_next(struct dns_record **record, + dns_record_cursor_t *cursor) +{ + struct __dns_record_entry *e; + + if (cursor->next->next != cursor->head) + { + cursor->next = cursor->next->next; + e = list_entry(cursor->next, struct __dns_record_entry, entry_list); + *record = &e->record; + return 0; + } + + return 1; +} + +int dns_record_cursor_find_cname(const char *name, + const char **cname, + dns_record_cursor_t *cursor) +{ + struct __dns_record_entry *e; + + if (!name || !cname) + return 1; + + cursor->next = cursor->head; + while (cursor->next->next != cursor->head) + { + cursor->next = cursor->next->next; + e = list_entry(cursor->next, struct __dns_record_entry, entry_list); + + if (e->record.type == DNS_TYPE_CNAME && + strcasecmp(name, e->record.name) == 0) + { + *cname = (const char *)e->record.rdata; + return 0; + } + } + + return 1; +} + +const char *dns_type2str(int type) +{ + switch (type) + { + case DNS_TYPE_A: + return "A"; + case DNS_TYPE_NS: + return "NS"; + case DNS_TYPE_MD: + return "MD"; + case DNS_TYPE_MF: + return "MF"; + case DNS_TYPE_CNAME: + return "CNAME"; + case DNS_TYPE_SOA: + return "SOA"; + case DNS_TYPE_MB: + return "MB"; + case DNS_TYPE_MG: + return "MG"; + case DNS_TYPE_MR: + return "MR"; + case DNS_TYPE_NULL: + return "NULL"; + case DNS_TYPE_WKS: + return "WKS"; + case DNS_TYPE_PTR: + return "PTR"; + case DNS_TYPE_HINFO: + return "HINFO"; + case DNS_TYPE_MINFO: + return "MINFO"; + case DNS_TYPE_MX: + return "MX"; + case DNS_TYPE_AAAA: + return "AAAA"; + case DNS_TYPE_SRV: + return "SRV"; + case DNS_TYPE_TXT: + return "TXT"; + case DNS_TYPE_AXFR: + return "AXFR"; + case DNS_TYPE_MAILB: + return "MAILB"; + case DNS_TYPE_MAILA: + return "MAILA"; + case DNS_TYPE_ALL: + return "ALL"; + default: + return "Unknown"; + } +} + +const char *dns_class2str(int dnsclass) +{ + switch (dnsclass) + { + case DNS_CLASS_IN: + return "IN"; + case DNS_CLASS_CS: + return "CS"; + case DNS_CLASS_CH: + return "CH"; + case DNS_CLASS_HS: + return "HS"; + case DNS_CLASS_ALL: + return "ALL"; + default: + return "Unknown"; + } +} + +const char *dns_opcode2str(int opcode) +{ + switch (opcode) + { + case DNS_OPCODE_QUERY: + return "QUERY"; + case DNS_OPCODE_IQUERY: + return "IQUERY"; + case DNS_OPCODE_STATUS: + return "STATUS"; + default: + return "Unknown"; + } +} + +const char *dns_rcode2str(int rcode) +{ + switch (rcode) + { + case DNS_RCODE_NO_ERROR: + return "NO_ERROR"; + case DNS_RCODE_FORMAT_ERROR: + return "FORMAT_ERROR"; + case DNS_RCODE_SERVER_FAILURE: + return "SERVER_FAILURE"; + case DNS_RCODE_NAME_ERROR: + return "NAME_ERROR"; + case DNS_RCODE_NOT_IMPLEMENTED: + return "NOT_IMPLEMENTED"; + case DNS_RCODE_REFUSED: + return "REFUSED"; + default: + return "Unknown"; + } +} diff --git a/src/protocol/dns_parser.h b/src/protocol/dns_parser.h new file mode 100644 index 00000000..e66c0f0f --- /dev/null +++ b/src/protocol/dns_parser.h @@ -0,0 +1,244 @@ +/* + 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) +*/ + +#ifndef _DNS_PARSER_H_ +#define _DNS_PARSER_H_ + +#include +#include +#include "list.h" + +enum +{ + DNS_TYPE_A = 1, + DNS_TYPE_NS, + DNS_TYPE_MD, + DNS_TYPE_MF, + DNS_TYPE_CNAME, + DNS_TYPE_SOA = 6, + DNS_TYPE_MB, + DNS_TYPE_MG, + DNS_TYPE_MR, + DNS_TYPE_NULL, + DNS_TYPE_WKS = 11, + DNS_TYPE_PTR, + DNS_TYPE_HINFO, + DNS_TYPE_MINFO, + DNS_TYPE_MX, + DNS_TYPE_TXT = 16, + + DNS_TYPE_AAAA = 28, + DNS_TYPE_SRV = 33, + + DNS_TYPE_AXFR = 252, + DNS_TYPE_MAILB = 253, + DNS_TYPE_MAILA = 254, + DNS_TYPE_ALL = 255 +}; + +enum +{ + DNS_CLASS_IN = 1, + DNS_CLASS_CS, + DNS_CLASS_CH, + DNS_CLASS_HS, + + DNS_CLASS_ALL = 255 +}; + +enum +{ + DNS_OPCODE_QUERY = 0, + DNS_OPCODE_IQUERY, + DNS_OPCODE_STATUS, +}; + +enum +{ + DNS_RCODE_NO_ERROR = 0, + DNS_RCODE_FORMAT_ERROR, + DNS_RCODE_SERVER_FAILURE, + DNS_RCODE_NAME_ERROR, + DNS_RCODE_NOT_IMPLEMENTED, + DNS_RCODE_REFUSED +}; + +/** + * dns_header_t is a struct to describe the header of a dns + * request or response packet, but the byte order is not + * transformed. + */ +struct dns_header +{ + uint16_t id; +#if __BYTE_ORDER == __LITTLE_ENDIAN + uint8_t rd : 1; + uint8_t tc : 1; + uint8_t aa : 1; + uint8_t opcode : 4; + uint8_t qr : 1; + uint8_t rcode : 4; + uint8_t z : 3; + uint8_t ra : 1; +#elif __BYTE_ORDER == __BIG_ENDIAN + uint8_t qr : 1; + uint8_t opcode : 4; + uint8_t aa : 1; + uint8_t tc : 1; + uint8_t rd : 1; + uint8_t ra : 1; + uint8_t z : 3; + uint8_t rcode : 4; +#else +#error "Your systems ENDIANNESS is broken, please fix!" +#endif + uint16_t qdcount; + uint16_t ancount; + uint16_t nscount; + uint16_t arcount; +}; + +struct dns_question +{ + char *qname; + uint16_t qtype; + uint16_t qclass; +}; + +struct dns_record_soa +{ + char *mname; + char *rname; + uint32_t serial; + int32_t refresh; + int32_t retry; + int32_t expire; + uint32_t minimum; +}; + +struct dns_record_srv +{ + uint16_t priority; + uint16_t weight; + uint16_t port; + char *target; +}; + +struct dns_record_mx +{ + int16_t preference; + char *exchange; +}; + +struct dns_record +{ + char *name; + uint16_t type; + uint16_t rclass; + uint32_t ttl; + uint16_t rdlength; + void *rdata; +}; + +#define DNS_HEADER_SIZE sizeof(struct dns_header) + +typedef struct __dns_parser +{ + void *msgbuf; // Message with leading length (TCP) + void *msgbase; // Message without leading length + const char *cur; // Current parser position + size_t msgsize; + size_t bufsize; + char complete; // Whether parse completed + char leading_length; // Response has leading length When TCP + struct dns_header header; + struct dns_question question; + struct list_head answer_list; + struct list_head authority_list; + struct list_head additional_list; +} dns_parser_t; + +typedef struct __dns_record_cursor +{ + const struct list_head *head; + const struct list_head *next; +} dns_record_cursor_t; + +#ifdef __cplusplus +extern "C" +{ +#endif + +void dns_parser_init(dns_parser_t *parser); + +void dns_parser_set_id(uint16_t id, dns_parser_t *parser); +int dns_parser_set_question(const char *name, + uint16_t qtype, + uint16_t qclass, + dns_parser_t *parser); +int dns_parser_set_question_name(const char *name, + dns_parser_t *parser); + +int dns_parser_parse_all(dns_parser_t *parser); +int dns_parser_append_message(const void *buf, size_t *n, + dns_parser_t *parser); + +void dns_parser_deinit(dns_parser_t *parser); + +int dns_record_cursor_next(struct dns_record **record, + dns_record_cursor_t *cursor); + +int dns_record_cursor_find_cname(const char *name, + const char **cname, + dns_record_cursor_t *cursor); + +const char *dns_type2str(int type); +const char *dns_class2str(int dnsclass); +const char *dns_opcode2str(int opcode); +const char *dns_rcode2str(int rcode); + +#ifdef __cplusplus +} +#endif + +static inline void dns_answer_cursor_init(dns_record_cursor_t *cursor, + const dns_parser_t *parser) +{ + cursor->head = &parser->answer_list; + cursor->next = cursor->head; +} + +static inline void dns_authority_cursor_init(dns_record_cursor_t *cursor, + const dns_parser_t *parser) +{ + cursor->head = &parser->authority_list; + cursor->next = cursor->head; +} + +static inline void dns_additional_cursor_init(dns_record_cursor_t *cursor, + const dns_parser_t *parser) +{ + cursor->head = &parser->additional_list; + cursor->next = cursor->head; +} + +static inline void dns_record_cursor_deinit(dns_record_cursor_t *cursor) +{ +} + +#endif diff --git a/src/util/CMakeLists.txt b/src/util/CMakeLists.txt index cb5def89..875825d0 100644 --- a/src/util/CMakeLists.txt +++ b/src/util/CMakeLists.txt @@ -2,6 +2,7 @@ cmake_minimum_required(VERSION 3.6) project(util) set(SRC + DnsUtil.cc EncodeStream.cc StringUtil.cc URIParser.cc diff --git a/src/util/DnsUtil.cc b/src/util/DnsUtil.cc new file mode 100644 index 00000000..f34ac51b --- /dev/null +++ b/src/util/DnsUtil.cc @@ -0,0 +1,126 @@ +/* + 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 "DnsUtil.h" + +int DnsUtil::getaddrinfo(const protocol::DnsResponse *resp, + unsigned short port, + struct addrinfo **addrinfo) +{ + int ancount = resp->get_ancount(); + int rcode = resp->get_rcode(); + int status = 0; + struct addrinfo *res = NULL; + struct addrinfo **pres = &res; + struct dns_record *record; + struct addrinfo *ai; + std::string qname; + const char *cname; + int family; + int addrlen; + + switch (rcode) + { + case DNS_RCODE_NAME_ERROR: + status = EAI_NONAME; + break; + case DNS_RCODE_SERVER_FAILURE: + status = EAI_AGAIN; + break; + case DNS_RCODE_FORMAT_ERROR: + case DNS_RCODE_NOT_IMPLEMENTED: + case DNS_RCODE_REFUSED: + status = EAI_FAIL; + break; + } + + qname = resp->get_question_name(); + cname = qname.c_str(); + + DnsResultCursor cursor(resp); + cursor.reset_answer_cursor(); + /* Forbid loop in cname chain */ + while (cursor.find_cname(cname, &cname) && ancount-- > 0) { } + + if (rcode == DNS_RCODE_NO_ERROR && ancount <= 0) + status = EAI_NODATA; + if (status != 0) + return status; + + cursor.reset_answer_cursor(); + while (cursor.next(&record)) + { + if (!(record->rclass == DNS_CLASS_IN && + (record->type == DNS_TYPE_A || record->type == DNS_TYPE_AAAA) && + strcasecmp(record->name, cname) == 0)) + continue; + + if (record->type == DNS_TYPE_A) + { + family = AF_INET; + addrlen = sizeof (struct sockaddr_in); + } + else + { + family = AF_INET6; + addrlen = sizeof (struct sockaddr_in6); + } + + ai = (struct addrinfo *)calloc(sizeof (struct addrinfo) + addrlen, 1); + if (ai == NULL) + { + if (res) + freeaddrinfo(res); + return EAI_MEMORY; + } + + ai->ai_family = family; + ai->ai_addrlen = addrlen; + ai->ai_addr = (struct sockaddr *)(ai + 1); + ai->ai_addr->sa_family = family; + + if (family == AF_INET) + { + struct sockaddr_in *in = (struct sockaddr_in *)(ai->ai_addr); + in->sin_port = htons(port); + memcpy(&in->sin_addr, record->rdata, sizeof (struct in_addr)); + } + else + { + struct sockaddr_in6 *in = (struct sockaddr_in6 *)(ai->ai_addr); + in->sin6_port = htons(port); + memcpy(&in->sin6_addr, record->rdata, sizeof (struct in6_addr)); + } + + *pres = ai; + pres = &ai->ai_next; + } + + if (res == NULL) + return EAI_NODATA; + + if (cname) + res->ai_canonname = strdup(cname); + + *addrinfo = res; + + return 0; +} diff --git a/src/util/DnsUtil.h b/src/util/DnsUtil.h new file mode 100644 index 00000000..4fa696a4 --- /dev/null +++ b/src/util/DnsUtil.h @@ -0,0 +1,89 @@ +/* + 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) +*/ + +#ifndef _DNSUTIL_H_ +#define _DNSUTIL_H_ + +#include "DnsMessage.h" + +/** + * @file DnsUtil.h + * @brief Dns toolbox + */ + +class DnsUtil +{ +public: + static int getaddrinfo(const protocol::DnsResponse *resp, + unsigned short port, + struct addrinfo **res); +}; + +class DnsResultCursor +{ +public: + DnsResultCursor(const protocol::DnsResponse *resp) : + parser(resp->get_parser()) + { + dns_answer_cursor_init(&cursor, parser); + record = NULL; + } + + DnsResultCursor(DnsResultCursor&& move) = delete; + DnsResultCursor& operator=(DnsResultCursor&& move) = delete; + + virtual ~DnsResultCursor() { } + + void reset_answer_cursor() + { + dns_answer_cursor_init(&cursor, parser); + } + + void reset_authority_cursor() + { + dns_authority_cursor_init(&cursor, parser); + } + + void reset_additional_cursor() + { + dns_additional_cursor_init(&cursor, parser); + } + + bool next(dns_record **next_record) + { + int ret = dns_record_cursor_next(&record, &cursor); + if (ret != 0) + record = NULL; + else + *next_record = record; + + return ret == 0; + } + + bool find_cname(const char *name, const char **cname) + { + return dns_record_cursor_find_cname(name, cname, &cursor) == 0; + } + +private: + const dns_parser_t *parser; + dns_record_cursor_t cursor; + dns_record *record; +}; + +#endif diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 68c3bd57..c673b168 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -38,6 +38,7 @@ set(TEST_LIST graph_unittest memory_unittest upstream_unittest + dns_unittest ) if (APPLE) diff --git a/test/dns_unittest.cc b/test/dns_unittest.cc new file mode 100644 index 00000000..047de099 --- /dev/null +++ b/test/dns_unittest.cc @@ -0,0 +1,127 @@ +/* + 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 "workflow/WFTaskFactory.h" +#include "workflow/WFDnsClient.h" + +#define RETRY_MAX 3 + +TEST(dns_unittest, WFDnsTaskCreate1) +{ + std::string url = "dns://119.29.29.29/www.sogou.com"; + auto *task = WFTaskFactory::create_dns_task(url, 0, NULL); + task->dismiss(); +} + +TEST(dns_unittest, WFDnsTaskCreate2) +{ + std::string url = "http://119.29.29.29:dns/"; + std::promise done; + auto *task = WFTaskFactory::create_dns_task(url, 0, + [&done] (WFDnsTask *task) + { + done.set_value(); + }); + task->start(); + done.get_future().get(); +} + +TEST(dns_unittest, WFDnsTask) +{ + std::string url = "dns://119.29.29.29/www.sogou.com"; + unsigned short req_id = 0x1234; + std::promise done; + + auto *task = WFTaskFactory::create_dns_task(url, RETRY_MAX, + [&done, req_id] (WFDnsTask *task) + { + int state = task->get_state(); + + if (state == WFT_STATE_SUCCESS) + { + unsigned short resp_id = task->get_resp()->get_id(); + EXPECT_TRUE(req_id == resp_id); + } + + done.set_value(); + }); + + auto *req = task->get_req(); + req->set_id(req_id); + req->set_rd(1); + req->set_question_type(DNS_TYPE_A); + task->start(); + + auto fut = done.get_future(); + fut.get(); +} + +TEST(dns_unittest, WFDnsClientInit1) +{ + WFDnsClient client; + if (client.init("bad") >= 0) + client.deinit(); +} + +TEST(dns_unittest, WFDnsClientInit2) +{ + WFDnsClient client; + int ret = client.init("0.0.0.0,0.0.0.1:1,dns://0.0.0.2,dnss://0.0.0.3"); + EXPECT_TRUE(ret >= 0); + client.deinit(); +} + +TEST(dns_unittest, WFDnsClient) +{ + unsigned short req_id = 0x4321; + std::promise done; + WFDnsClient client; + + client.init("dns://119.29.29.29/"); + + auto *task = client.create_dns_task("www.sogou.com", + [&done, req_id] (WFDnsTask *task) + { + int state = task->get_state(); + + if (state == WFT_STATE_SUCCESS) + { + unsigned short resp_id = task->get_resp()->get_id(); + EXPECT_TRUE(req_id == resp_id); + } + + done.set_value(); + }); + + client.deinit(); + + auto *req = task->get_req(); + req->set_id(req_id); + task->start(); + + auto fut = done.get_future(); + fut.get(); +} + +int main(int argc, char *argv[]) +{ + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +}