From cde9fbbac9e82ef1abfa53deb2f890056085daad Mon Sep 17 00:00:00 2001 From: kedixa <1204837541@qq.com> Date: Mon, 7 Jun 2021 19:27:03 +0800 Subject: [PATCH] simple dns server --- CMakeLists_Headers.txt | 1 + src/factory/DnsTaskImpl.cc | 38 ++++++++++++++++++++++ src/factory/WFTaskFactory.inl | 2 ++ src/protocol/DnsMessage.cc | 3 +- src/protocol/DnsMessage.h | 5 +-- src/protocol/dns_parser.c | 1 + src/protocol/dns_parser.h | 4 +-- src/server/WFDnsServer.h | 60 +++++++++++++++++++++++++++++++++++ 8 files changed, 108 insertions(+), 6 deletions(-) create mode 100644 src/server/WFDnsServer.h diff --git a/CMakeLists_Headers.txt b/CMakeLists_Headers.txt index a4be8811..a7680d56 100644 --- a/CMakeLists_Headers.txt +++ b/CMakeLists_Headers.txt @@ -58,6 +58,7 @@ set(INCLUDE_HEADERS src/protocol/DnsMessage.h src/protocol/DnsUtil.h src/server/WFServer.h + src/server/WFDnsServer.h src/server/WFHttpServer.h src/server/WFRedisServer.h src/server/WFMySQLServer.h diff --git a/src/factory/DnsTaskImpl.cc b/src/factory/DnsTaskImpl.cc index e9464db9..c57d2abe 100644 --- a/src/factory/DnsTaskImpl.cc +++ b/src/factory/DnsTaskImpl.cc @@ -25,6 +25,8 @@ using namespace protocol; #define DNS_KEEPALIVE_DEFAULT (60 * 1000) +/**********Client**********/ + class ComplexDnsTask : public WFComplexClientTask> { @@ -164,6 +166,8 @@ bool ComplexDnsTask::need_redirect() return false; } +/**********Client Factory**********/ + WFDnsTask *WFTaskFactory::create_dns_task(const std::string& url, int retry_max, dns_callback_t callback) @@ -193,3 +197,37 @@ WFDnsTask *WFTaskFactory::create_dns_task(const ParsedURI& uri, task->set_keep_alive(DNS_KEEPALIVE_DEFAULT); return task; } + +/**********Server**********/ + +class WFDnsServerTask : public WFServerTask +{ +public: + WFDnsServerTask(CommService *service, + std::function& process): + WFServerTask(service, WFGlobal::get_scheduler(), process) + { } + +protected: + virtual CommMessageOut *message_out() + { + DnsResponse *resp = this->get_resp(); + resp->set_leading_length(true); + return this->WFServerTask::message_out(); + } + + virtual CommMessageIn *message_in() + { + DnsRequest *req = this->get_req(); + req->set_leading_length(true); + return this->WFServerTask::message_in(); + } +}; + +/**********Server Factory**********/ + +WFDnsTask *WFServerTaskFactory::create_dns_task(CommService *service, + std::function& process) +{ + return new WFDnsServerTask(service, process); +} diff --git a/src/factory/WFTaskFactory.inl b/src/factory/WFTaskFactory.inl index 0dee29b1..8f20f0a3 100644 --- a/src/factory/WFTaskFactory.inl +++ b/src/factory/WFTaskFactory.inl @@ -594,6 +594,8 @@ public: std::function& process); static WFMySQLTask *create_mysql_task(CommService *service, std::function& process); + static WFDnsTask *create_dns_task(CommService *service, + std::function& process); }; /**********Template Thread Task Factory**********/ diff --git a/src/protocol/DnsMessage.cc b/src/protocol/DnsMessage.cc index 187e3ae6..8b7d9ce9 100644 --- a/src/protocol/DnsMessage.cc +++ b/src/protocol/DnsMessage.cc @@ -22,6 +22,7 @@ #define DNS_LABELS_MAX 63 #define DNS_MESSAGE_MAX_UDP_SIZE 512 +#define DNS_HEADER_SIZE sizeof (struct dns_header) namespace protocol { @@ -90,7 +91,7 @@ int DnsMessage::encode_reply() h.arcount = htons(0); msgbuf.append((const char *)&h, sizeof (struct dns_header)); - p = parser->question.qname; + p = parser->question.qname ? parser->question.qname : "."; while (*p) { name = p; diff --git a/src/protocol/DnsMessage.h b/src/protocol/DnsMessage.h index a8d079b7..c4353225 100644 --- a/src/protocol/DnsMessage.h +++ b/src/protocol/DnsMessage.h @@ -230,8 +230,6 @@ public: return this->parser; } - int append(const void *buf, size_t *size); - void set_request_id(uint16_t id) { this->request_id = id; @@ -242,6 +240,9 @@ public: this->request_name = name; } +protected: + virtual int append(const void *buf, size_t *size); + private: uint16_t request_id; std::string request_name; diff --git a/src/protocol/dns_parser.c b/src/protocol/dns_parser.c index 080392ec..d7e5dd17 100644 --- a/src/protocol/dns_parser.c +++ b/src/protocol/dns_parser.c @@ -24,6 +24,7 @@ #define DNS_LABELS_MAX 63 #define DNS_NAMES_MAX 256 #define DNS_MSGBASE_INIT_SIZE 514 // 512 + 2(leading length) +#define DNS_HEADER_SIZE sizeof (struct dns_header) #define MAX(x, y) ((x) <= (y) ? (y) : (x)) struct __dns_record_entry diff --git a/src/protocol/dns_parser.h b/src/protocol/dns_parser.h index e66c0f0f..33d7b267 100644 --- a/src/protocol/dns_parser.h +++ b/src/protocol/dns_parser.h @@ -105,7 +105,7 @@ struct dns_header uint8_t z : 3; uint8_t rcode : 4; #else -#error "Your systems ENDIANNESS is broken, please fix!" +#error endian? #endif uint16_t qdcount; uint16_t ancount; @@ -155,8 +155,6 @@ struct dns_record void *rdata; }; -#define DNS_HEADER_SIZE sizeof(struct dns_header) - typedef struct __dns_parser { void *msgbuf; // Message with leading length (TCP) diff --git a/src/server/WFDnsServer.h b/src/server/WFDnsServer.h new file mode 100644 index 00000000..bf012c60 --- /dev/null +++ b/src/server/WFDnsServer.h @@ -0,0 +1,60 @@ +/* + 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. + + Authors: Liu Kai (liukaidx@sogou-inc.com) +*/ + +#ifndef _WFDNSSERVER_H_ +#define _WFDNSSERVER_H_ + +#include "DnsMessage.h" +#include "WFServer.h" +#include "WFTaskFactory.h" + +using dns_process_t = std::function; +using WFDnsServer = WFServer; + +static constexpr struct WFServerParams DNS_SERVER_PARAMS_DEFAULT = +{ + .max_connections = 2000, + .peer_response_timeout = 10 * 1000, + .receive_timeout = -1, + .keep_alive_timeout = 300 * 1000, + .request_size_limit = (size_t)-1, + .ssl_accept_timeout = 5000, +}; + +template<> +inline WFDnsServer::WFServer(dns_process_t proc) : + WFServerBase(&DNS_SERVER_PARAMS_DEFAULT), + process(std::move(proc)) +{ +} + +template<> +inline CommSession *WFDnsServer::new_session(long long seq, CommConnection *conn) +{ + WFDnsTask *task; + + task = WFServerTaskFactory::create_dns_task(this, this->process); + task->set_keep_alive(this->params.keep_alive_timeout); + task->set_receive_timeout(this->params.receive_timeout); + task->get_req()->set_size_limit(this->params.request_size_limit); + + return task; +} + +#endif