[Feature] Add support of fake DNS records

4 files changed, 165 insertions, 68 deletions

diff --git a/contrib/librdns/dns_private.h b/contrib/librdns/dns_private.h
index 37d0240b6..44bb3dd84 100644
--- a/contrib/librdns/dns_private.h
+++ b/contrib/librdns/dns_private.h
@@ -56,6 +56,15 @@ struct rdns_server {
 	upstream_entry_t up;
 };
 
+enum rdns_request_state {
+	RDNS_REQUEST_NEW = 0,
+	RDNS_REQUEST_REGISTERED = 1,
+	RDNS_REQUEST_WAIT_SEND,
+	RDNS_REQUEST_WAIT_REPLY,
+	RDNS_REQUEST_REPLIED,
+	RDNS_REQUEST_FAKE,
+};
+
 struct rdns_request {
 	struct rdns_resolver *resolver;
 	struct rdns_async_context *async;
@@ -69,14 +78,7 @@ struct rdns_request {
 	int id;
 	struct rdns_request_name *requested_names;
 	unsigned int qcount;
-
-	enum {
-		RDNS_REQUEST_NEW = 0,
-		RDNS_REQUEST_REGISTERED = 1,
-		RDNS_REQUEST_WAIT_SEND,
-		RDNS_REQUEST_WAIT_REPLY,
-		RDNS_REQUEST_REPLIED
-	} state;
+	enum rdns_request_state state;
 
 	uint8_t *packet;
 	off_t pos;
@@ -109,6 +111,13 @@ struct rdns_io_channel {
 	ref_entry_t ref;
 };
 
+struct rdns_fake_reply {
+	char *request;
+	enum dns_rcode rcode;
+	struct rdns_reply_entry *result;
+	UT_hash_handle hh;
+};
+
 
 struct rdns_resolver {
 	struct rdns_server *servers;
@@ -117,6 +126,7 @@ struct rdns_resolver {
 	void *periodic; /** periodic event for resolver */
 	struct rdns_upstream_context *ups;
 	struct rdns_plugin *curve_plugin;
+	struct rdns_fake_reply *fake_elts;
 
 	rdns_log_function logger;
 	void *log_data;
diff --git a/contrib/librdns/rdns.h b/contrib/librdns/rdns.h
index 5c44900f1..545d9421f 100644
--- a/contrib/librdns/rdns.h
+++ b/contrib/librdns/rdns.h
@@ -324,6 +324,18 @@ void rdns_resolver_register_plugin (struct rdns_resolver *resolver,
 		struct rdns_plugin *plugin);
 
 /**
+ * Add a fake reply for a specified name
+ * @param resolver
+ * @param type
+ * @param name
+ * @param reply
+ */
+void rdns_resolver_set_fake_reply (struct rdns_resolver *resolver,
+								   const char *name,
+								   enum dns_rcode rcode,
+								   struct rdns_reply_entry *reply);
+
+/**
  * Init DNS resolver
  * @param resolver
  * @return
diff --git a/contrib/librdns/resolver.c b/contrib/librdns/resolver.c
index 7fa61a391..7afabb1ba 100644
--- a/contrib/librdns/resolver.c
+++ b/contrib/librdns/resolver.c
@@ -484,6 +484,14 @@ rdns_process_retransmit (int fd, void *arg)
 			req->async_event);
 	req->async_event = NULL;
 
+	if (req->state == RDNS_REQUEST_FAKE) {
+		/* Reply is ready */
+		req->func (req->reply, req->arg);
+		REF_RELEASE (req);
+
+		return;
+	}
+
 	r = rdns_send_request (req, fd, false);
 
 	if (r == 0) {
@@ -532,6 +540,7 @@ rdns_make_request_full (
 	size_t olen;
 	const char *cur_name, *last_name = NULL;
 	struct rdns_compression_entry *comp = NULL;
+	struct rdns_fake_reply *fake_rep = NULL;
 
 	if (resolver == NULL || !resolver->initialized) {
 		return NULL;
@@ -569,6 +578,18 @@ rdns_make_request_full (
 	for (i = 0; i < queries * 2; i += 2) {
 		cur = i / 2;
 		cur_name = va_arg (args, const char *);
+
+		if (last_name == NULL) {
+			HASH_FIND_STR (resolver->fake_elts, cur_name, fake_rep);
+
+			if (fake_rep) {
+				/* We actually treat it as a short-circuit */
+				req->reply = rdns_make_reply (req, fake_rep->rcode);
+				req->reply->entries = fake_rep->result;
+				req->state = RDNS_REQUEST_FAKE;
+			}
+		}
+
 		if (cur_name != NULL) {
 			last_name = cur_name;
 			clen = strlen (cur_name);
@@ -585,7 +606,8 @@ rdns_make_request_full (
 			return NULL;
 		}
 
-		if (!rdns_format_dns_name (resolver, last_name, clen,
+		if (req->state != RDNS_REQUEST_FAKE &&
+			!rdns_format_dns_name (resolver, last_name, clen,
 				&req->requested_names[cur].name, &olen)) {
 			rdns_request_free (req);
 			return NULL;
@@ -597,37 +619,39 @@ rdns_make_request_full (
 	}
 	va_end (args);
 
-	rdns_allocate_packet (req, tlen);
-	rdns_make_dns_header (req, queries);
-
-	for (i = 0; i < queries; i ++) {
-		cur_name = req->requested_names[i].name;
-		clen = req->requested_names[i].len;
-		type = req->requested_names[i].type;
-		if (queries > 1) {
-			if (!rdns_add_rr (req, cur_name, clen, type, &comp)) {
-				REF_RELEASE (req);
-				rnds_compression_free (comp);
-				return NULL;
-			}
-		}
-		else {
-			if (!rdns_add_rr (req, cur_name, clen, type, NULL)) {
-				REF_RELEASE (req);
-				rnds_compression_free (comp);
-				return NULL;
+	if (req->state != RDNS_REQUEST_FAKE) {
+		rdns_allocate_packet (req, tlen);
+		rdns_make_dns_header (req, queries);
+
+		for (i = 0; i < queries; i++) {
+			cur_name = req->requested_names[i].name;
+			clen = req->requested_names[i].len;
+			type = req->requested_names[i].type;
+			if (queries > 1) {
+				if (!rdns_add_rr (req, cur_name, clen, type, &comp)) {
+					REF_RELEASE (req);
+					rnds_compression_free (comp);
+					return NULL;
+				}
+			} else {
+				if (!rdns_add_rr (req, cur_name, clen, type, NULL)) {
+					REF_RELEASE (req);
+					rnds_compression_free (comp);
+					return NULL;
+				}
 			}
 		}
-	}
 
-	rnds_compression_free (comp);
+		rnds_compression_free (comp);
 
-	/* Add EDNS RR */
-	rdns_add_edns0 (req);
+		/* Add EDNS RR */
+		rdns_add_edns0 (req);
+
+		req->retransmits = repeats;
+		req->timeout = timeout;
+		req->state = RDNS_REQUEST_NEW;
+	}
 
-	req->retransmits = repeats;
-	req->timeout = timeout;
-	req->state = RDNS_REQUEST_NEW;
 	req->async = resolver->async;
 
 	if (resolver->ups) {
@@ -656,14 +680,21 @@ rdns_make_request_full (
 
 	/* Select random IO channel */
 	req->io = serv->io_channels[ottery_rand_uint32 () % serv->io_cnt];
-	req->io->uses ++;
 
-	/* Now send request to server */
-	r = rdns_send_request (req, req->io->sock, true);
+	if (req->state == RDNS_REQUEST_FAKE) {
+		req->async_event = resolver->async->add_write (resolver->async->data,
+				req->io->sock, req);
+	}
+	else {
+		req->io->uses++;
 
-	if (r == -1) {
-		REF_RELEASE (req);
-		return NULL;
+		/* Now send request to server */
+		r = rdns_send_request (req, req->io->sock, true);
+
+		if (r == -1) {
+			REF_RELEASE (req);
+			return NULL;
+		}
 	}
 
 	REF_RETAIN (req->io);
@@ -889,3 +920,39 @@ rdns_resolver_set_dnssec (struct rdns_resolver *resolver, bool enabled)
 		resolver->enable_dnssec = enabled;
 	}
 }
+
+
+void rdns_resolver_set_fake_reply (struct rdns_resolver *resolver,
+								   const char *name,
+								   enum dns_rcode rcode,
+								   struct rdns_reply_entry *reply)
+{
+	struct rdns_fake_reply *fake_rep;
+
+	HASH_FIND_STR (resolver->fake_elts, name, fake_rep);
+
+	if (fake_rep) {
+		/* Append reply to the existing list */
+		fake_rep->rcode = rcode;
+		DL_APPEND (fake_rep->result, reply);
+	}
+	else {
+		fake_rep = calloc (1, sizeof (*fake_rep));
+
+		if (fake_rep == NULL) {
+			abort ();
+		}
+
+		fake_rep->request = strdup (name);
+
+		if (fake_rep->request == NULL) {
+			abort ();
+		}
+
+		fake_rep->rcode = rcode;
+		DL_APPEND (fake_rep->result, reply);
+
+		HASH_ADD_KEYPTR (hh, resolver->fake_elts, fake_rep->request,
+				strlen (fake_rep->request), fake_rep);
+	}
+}
\ No newline at end of file
diff --git a/contrib/librdns/util.c b/contrib/librdns/util.c
index b03128983..a4018cbd3 100644
--- a/contrib/librdns/util.c
+++ b/contrib/librdns/util.c
@@ -353,34 +353,38 @@ rdns_reply_free (struct rdns_reply *rep)
 {
 	struct rdns_reply_entry *entry, *tmp;
 
-	LL_FOREACH_SAFE (rep->entries, entry, tmp) {
-		switch (entry->type) {
-		case RDNS_REQUEST_PTR:
-			free (entry->content.ptr.name);
-			break;
-		case RDNS_REQUEST_NS:
-			free (entry->content.ns.name);
-			break;
-		case RDNS_REQUEST_MX:
-			free (entry->content.mx.name);
-			break;
-		case RDNS_REQUEST_TXT:
-		case RDNS_REQUEST_SPF:
-			free (entry->content.txt.data);
-			break;
-		case RDNS_REQUEST_SRV:
-			free (entry->content.srv.target);
-			break;
-		case RDNS_REQUEST_TLSA:
-			free (entry->content.tlsa.data);
-			break;
-		case RDNS_REQUEST_SOA:
-			free (entry->content.soa.mname);
-			free (entry->content.soa.admin);
-			break;
+	/* We don't need to free data for faked replies */
+	if (!rep->request || rep->request->state != RDNS_REQUEST_FAKE) {
+		LL_FOREACH_SAFE (rep->entries, entry, tmp) {
+			switch (entry->type) {
+			case RDNS_REQUEST_PTR:
+				free (entry->content.ptr.name);
+				break;
+			case RDNS_REQUEST_NS:
+				free (entry->content.ns.name);
+				break;
+			case RDNS_REQUEST_MX:
+				free (entry->content.mx.name);
+				break;
+			case RDNS_REQUEST_TXT:
+			case RDNS_REQUEST_SPF:
+				free (entry->content.txt.data);
+				break;
+			case RDNS_REQUEST_SRV:
+				free (entry->content.srv.target);
+				break;
+			case RDNS_REQUEST_TLSA:
+				free (entry->content.tlsa.data);
+				break;
+			case RDNS_REQUEST_SOA:
+				free (entry->content.soa.mname);
+				free (entry->content.soa.admin);
+				break;
+			}
+			free (entry);
 		}
-		free (entry);
 	}
+
 	free (rep);
 }
 
@@ -418,7 +422,11 @@ rdns_request_free (struct rdns_request *req)
 				HASH_DEL (req->io->requests, req);
 				req->async_event = NULL;
 			}
-
+			else if (req->state == RDNS_REQUEST_FAKE) {
+				req->async->del_write (req->async->data,
+						req->async_event);
+				req->async_event = NULL;
+			}
 		}
 #ifdef TWEETNACL
 		if (req->curve_plugin_data != NULL) {