simple dns server

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

src/factory/DnsTaskImpl.cc

@@ -25,6 +25,8 @@ using namespace protocol;
 
 #define DNS_KEEPALIVE_DEFAULT	(60 * 1000)
 
+/**********Client**********/
+
 class ComplexDnsTask : public WFComplexClientTask<DnsRequest, DnsResponse,
 							  std::function<void (WFDnsTask *)>>
 {
@@ -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<DnsRequest, DnsResponse>
+{
+public:
+	WFDnsServerTask(CommService *service,
+					 std::function<void (WFDnsTask *)>& 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<void (WFDnsTask *)>& process)
+{
+	return new WFDnsServerTask(service, process);
+}

src/factory/WFTaskFactory.inl

@@ -594,6 +594,8 @@ public:
 					std::function<void (WFHttpTask *)>& process);
 	static WFMySQLTask *create_mysql_task(CommService *service,
 					std::function<void (WFMySQLTask *)>& process);
+	static WFDnsTask *create_dns_task(CommService *service,
+					std::function<void (WFDnsTask *)>& process);
 };
 
 /**********Template Thread Task Factory**********/

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;

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;

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

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)

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<void (WFDnsTask *)>;
+using WFDnsServer = WFServer<protocol::DnsRequest,
+							 protocol::DnsResponse>;
+
+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