8 files changed, 525 insertions, 467 deletions

diff --git a/include/asterisk/dns_internal.h b/include/asterisk/dns_internal.h
index e42e3f7d8..d518f9066 100644
--- a/include/asterisk/dns_internal.h
+++ b/include/asterisk/dns_internal.h
@@ -195,12 +195,49 @@ void dns_naptr_sort(struct ast_dns_result *result);
  * \retval non-NULL success
  * \retval NULL failure
  */
-struct ast_dns_record *ast_dns_srv_alloc(struct ast_dns_query *query, const char *data, const size_t size);
+struct ast_dns_record *dns_srv_alloc(struct ast_dns_query *query, const char *data, const size_t size);
 
 /*!
  * \brief Sort the SRV records on a result
  *
  * \param result The DNS result
  */
-void ast_dns_srv_sort(struct ast_dns_result *result);
+void dns_srv_sort(struct ast_dns_result *result);
 
+/*!
+ * \brief Find the location of a DNS record within the entire DNS answer
+ *
+ * The DNS record that has been returned by the resolver may be a copy of the record that was
+ * found in the complete DNS response. If so, then some DNS record types (specifically those that
+ * parse domains) will need to locate the DNS record within the complete DNS response. This is so
+ * that if the domain contains pointers to other sections of the DNS response, then the referenced
+ * domains may be located.
+ *
+ * \param record The DNS record returned by a resolver implementation
+ * \param record_size The size of the DNS record in bytes
+ * \param response The complete DNS answer
+ * \param response_size The size of the complete DNS response
+ */
+char *dns_find_record(const char *record, size_t record_size, const char *response, size_t response_size);
+
+/*!
+ * \brief Parse a 16-bit unsigned value from a DNS record
+ *
+ * \param cur Pointer to the location of the 16-bit value in the DNS record
+ * \param[out] val The parsed 16-bit unsigned integer
+ * \return The number of bytes consumed while parsing
+ */
+int dns_parse_short(unsigned char *cur, uint16_t *val);
+
+/*!
+ * \brief Parse a DNS string from a DNS record
+ *
+ * A DNS string consists of an 8-bit size, followed by the
+ * string value (not NULL-terminated).
+ *
+ * \param cur Pointer to the location of the DNS string
+ * \param[out] size The parsed size of the DNS string
+ * \param[out] val The contained string (not NULL-terminated)
+ * \return The number of bytes consumed while parsing
+ */
+int dns_parse_string(char *cur, uint8_t *size, char **val);
diff --git a/include/asterisk/dns_test.h b/include/asterisk/dns_test.h
new file mode 100644
index 000000000..0858bab33
--- /dev/null
+++ b/include/asterisk/dns_test.h
@@ -0,0 +1,109 @@
+/*
+ * Asterisk -- An open source telephony toolkit.
+ *
+ * Copyright (C) 2015, Digium, Inc.
+ *
+ * Mark Michelson <mmichelson@digium.com>
+ *
+ * Includes code and algorithms from the Zapata library.
+ *
+ * See http://www.asterisk.org for more information about
+ * the Asterisk project. Please do not directly contact
+ * any of the maintainers of this project for assistance;
+ * the project provides a web site, mailing lists and IRC
+ * channels for your use.
+ *
+ * This program is free software, distributed under the terms of
+ * the GNU General Public License Version 2. See the LICENSE file
+ * at the top of the source tree.
+ */
+
+#ifndef DNS_TEST_H
+#define DNS_TEST_H
+
+/*!
+ * \brief Representation of a string in DNS
+ *
+ * In DNS, a string has a byte to indicate the length,
+ * followed by a series of bytes representing the string.
+ * DNS does not NULL-terminate its strings. However, the
+ * string stored in this structure is expected to be NULL-
+ * terminated.
+ */
+struct ast_dns_test_string {
+	uint8_t len;
+	const char *val;
+};
+
+/*!
+ * \brief Write a DNS string to a buffer
+ *
+ * This writes the DNS string to the buffer and returns the total
+ * number of bytes written to the buffer.
+ *
+ * There is no buffer size passed to this function. Tests are expected to
+ * use a buffer that is sufficiently large for their tests.
+ *
+ * \param string The string to write
+ * \param buf The buffer to write the string into
+ * \return The number of bytes written to the buffer
+ */
+int ast_dns_test_write_string(const struct ast_dns_test_string *string, char *buf);
+
+/*!
+ * \brief Write a DNS domain to a buffer
+ *
+ * A DNS domain consists of a series of labels separated
+ * by dots. Each of these labels gets written as a DNS
+ * string. A DNS domain ends with a NULL label, which is
+ * essentially a zero-length DNS string.
+ *
+ * There is no buffer size passed to this function. Tests are expected to
+ * use a buffer that is sufficiently large for their tests.
+ *
+ * \param string The DNS domain to write
+ * \param buf The buffer to write the domain into
+ * \return The number of bytes written to the buffer
+ */
+int ast_dns_test_write_domain(const char *string, char *buf);
+
+/*!
+ * \brief Callback to write specific DNS record to an answer
+ *
+ * When generating a DNS result, the type of DNS record being generated
+ * will need to be performed by individual test cases. This is a callback
+ * that tests can define to write a specific type of DNS record to the
+ * provided buffer.
+ *
+ * There is no buffer size passed to this function. Tests are expected to
+ * use a buffer that is sufficiently large for their tests.
+ *
+ * \param record Pointer to test-specific DNS record data
+ * \param buf The buffer into which to write the DNS record
+ * \return The number of bytes written to the buffer
+ */
+typedef int (*record_fn)(void *record, char *buf);
+
+/*!
+ * \brief Generate a full DNS response for the given DNS records.
+ *
+ * This function takes care of generating the DNS header, question, and
+ * answer sections of a DNS response. In order to place test-specific
+ * record data into the DNS answers, a callback is provided as a parameter
+ * to this function so that the necessary records can be encoded properly
+ * by the tests.
+ *
+ * There is no buffer size passed to this function. Tests are expected to
+ * use a buffer that is sufficiently large for their tests.
+ *
+ * \param query The DNS query that is being processed
+ * \param records An array of test-specific representations of DNS records
+ * \param num_records The number of elements in the records array
+ * \param record_size The size of each element in the records array
+ * \param generate The test-specific encoder for DNS records
+ * \param buffer The buffer into which to write the DNS response
+ */
+int ast_dns_test_generate_result(struct ast_dns_query *query, void *records, size_t num_records,
+		size_t record_size, record_fn generate, char *buffer);
+
+#endif /* DNS_TEST_H */
diff --git a/main/dns_core.c b/main/dns_core.c
index 228147bbf..53ea1d09e 100644
--- a/main/dns_core.c
+++ b/main/dns_core.c
@@ -426,6 +426,20 @@ static struct ast_dns_record *generic_record_alloc(struct ast_dns_query *query,
 	return record;
 }
 
+typedef struct ast_dns_record *(*dns_alloc_fn)(struct ast_dns_query *query, const char *data, const size_t size);
+
+static dns_alloc_fn dns_alloc_table [] = {
+	[ns_t_naptr] = dns_naptr_alloc,
+	[ns_t_srv] = dns_srv_alloc,
+};
+
+static struct ast_dns_record *allocate_dns_record(int rr_type, struct ast_dns_query *query, const char *data, const size_t size)
+{
+	dns_alloc_fn allocator = dns_alloc_table[rr_type] ?: generic_record_alloc;
+
+	return allocator(query, data, size);
+}
+
 int ast_dns_resolver_add_record(struct ast_dns_query *query, int rr_type, int rr_class, int ttl, const char *data, const size_t size)
 {
 	struct ast_dns_record *record;
@@ -460,14 +474,7 @@ int ast_dns_resolver_add_record(struct ast_dns_query *query, int rr_type, int rr
 		return -1;
 	}
 
-	if (rr_type == ns_t_naptr) {
-		record = dns_naptr_alloc(query, data, size);
-	} else if (rr_type == ns_t_srv) {
-		record = ast_dns_srv_alloc(query, data, size);
-	} else {
-		record = generic_record_alloc(query, data, size);
-	}
-
+	record = allocate_dns_record(rr_type, query, data, size);
 	if (!record) {
 		return -1;
 	}
@@ -483,13 +490,23 @@ int ast_dns_resolver_add_record(struct ast_dns_query *query, int rr_type, int rr
 	return 0;
 }
 
-void ast_dns_resolver_completed(struct ast_dns_query *query)
+typedef void (*dns_sort_fn)(struct ast_dns_result *result);
+
+static dns_sort_fn dns_sort_table [] = {
+	[ns_t_naptr] = dns_naptr_sort,
+	[ns_t_srv] = dns_srv_sort,
+};
+
+static void sort_result(int rr_type, struct ast_dns_result *result)
 {
-	if (ast_dns_query_get_rr_type(query) == ns_t_naptr) {
-		dns_naptr_sort(query->result);
-	} else if (ast_dns_query_get_rr_type(query) == ns_t_srv) {
-		ast_dns_srv_sort(query->result);
+	if (dns_sort_table[rr_type]) {
+		dns_sort_table[rr_type](result);
 	}
+}
+
+void ast_dns_resolver_completed(struct ast_dns_query *query)
+{
+	sort_result(ast_dns_query_get_rr_type(query), query->result);
 
 	query->callback(query);
 }
@@ -593,3 +610,43 @@ void ast_dns_resolver_unregister(struct ast_dns_resolver *resolver)
 
 	ast_verb(2, "Unregistered DNS resolver '%s'\n", resolver->name);
 }
+
+char *dns_find_record(const char *record, size_t record_size, const char *response, size_t response_size)
+{
+	size_t remaining_size = response_size;
+	const char *search_base = response;
+	char *record_offset;
+
+	while (1) {
+		record_offset = memchr(search_base, record[0], remaining_size);
+
+		ast_assert(record_offset != NULL);
+		ast_assert(search_base + remaining_size - record_offset >= record_size);
+
+		if (!memcmp(record_offset, record, record_size)) {
+			return record_offset;
+		}
+
+		remaining_size -= record_offset - search_base;
+		search_base = record_offset + 1;
+	}
+}
+
+int dns_parse_short(unsigned char *cur, uint16_t *val)
+{
+	/* This assignment takes a big-endian 16-bit value and stores it in the
+	 * machine's native byte order. Using this method allows us to avoid potential
+	 * alignment issues in case the order is not on a short-addressable boundary.
+	 * See http://commandcenter.blogspot.com/2012/04/byte-order-fallacy.html for
+	 * more information
+	 */
+	*val = (cur[1] << 0) | (cur[0] << 8);
+	return sizeof(*val);
+}
+
+int dns_parse_string(char *cur, uint8_t *size, char **val)
+{
+	*size = *cur++;
+	*val = cur;
+	return *size + 1;
+}
diff --git a/main/dns_naptr.c b/main/dns_naptr.c
index 7e870002f..72f28dc81 100644
--- a/main/dns_naptr.c
+++ b/main/dns_naptr.c
@@ -392,54 +392,10 @@ struct ast_dns_record *dns_naptr_alloc(struct ast_dns_query *query, const char *
 	char *regexp;
 	char replacement[256] = "";
 	int replacement_size;
-	char *naptr_offset;
-	char *naptr_search_base = (char *)query->result->answer;
-	size_t remaining_size = query->result->answer_size;
-	char *end_of_record;
+	const char *end_of_record;
 	enum flags_result flags_res;
 
-	/*
-	 * This is bordering on the hackiest thing I've ever written.
-	 * Part of parsing a NAPTR record is to parse a potential replacement
-	 * domain name. Decoding this domain name requires the use of the
-	 * dn_expand() function. This function requires that the domain you
-	 * pass in be a pointer to within the full DNS answer. Unfortunately,
-	 * libunbound gives its RRs back as copies of data from the DNS answer
-	 * instead of pointers to within the DNS answer. This means that in order
-	 * to be able to parse the domain name correctly, I need to find the
-	 * current NAPTR record inside the DNS answer and operate on it. This
-	 * loop is designed to find the current NAPTR record within the full
-	 * DNS answer and set the "ptr" variable to the beginning of the
-	 * NAPTR RDATA
-	 */
-	while (1) {
-		naptr_offset = memchr(naptr_search_base, data[0], remaining_size);
-
-		/* Since the NAPTR record we have been given came from the DNS answer,
-		 * we should never run into a situation where we can't find ourself
-		 * in the answer
-		 */
-		ast_assert(naptr_offset != NULL);
-		ast_assert(naptr_search_base + remaining_size - naptr_offset >= size);
-
-		/* ... but just to be on the safe side, let's be sure we can break
-		 * out if the assertion doesn't hold
-		 */
-		if (!naptr_offset || naptr_search_base + remaining_size - naptr_offset < size) {
-			ast_log(LOG_ERROR, "Failed to locate NAPTR record within DNS result\n");
-			return NULL;
-		}
-
-		if (!memcmp(naptr_offset, data, size)) {
-			/* BAM! FOUND IT! */
-			ptr = naptr_offset;
-			break;
-		}
-		/* Data didn't match us, so keep looking */
-		remaining_size -= naptr_offset - naptr_search_base;
-		naptr_search_base = naptr_offset + 1;
-	}
-
+	ptr = dns_find_record(data, size, query->result->answer, query->result->answer_size);
 	ast_assert(ptr != NULL);
 
 	end_of_record = ptr + size;
@@ -451,53 +407,31 @@ struct ast_dns_record *dns_naptr_alloc(struct ast_dns_query *query, const char *
 	 * See http://commandcenter.blogspot.com/2012/04/byte-order-fallacy.html for
 	 * more information
 	 */
-	order = ((unsigned char)(ptr[1]) << 0) | ((unsigned char)(ptr[0]) << 8);
-	ptr += 2;
-
+	ptr += dns_parse_short((unsigned char *) ptr, &order);
 	if (PAST_END_OF_RECORD) {
 		return NULL;
 	}
 
 	/* PREFERENCE */
-	preference = ((unsigned char) (ptr[1]) << 0) | ((unsigned char)(ptr[0]) << 8);
-	ptr += 2;
-
+	ptr += dns_parse_short((unsigned char *) ptr, &preference);
 	if (PAST_END_OF_RECORD) {
 		return NULL;
 	}
 
 	/* FLAGS */
-	flags_size = *ptr;
-	++ptr;
-	if (PAST_END_OF_RECORD) {
-		return NULL;
-	}
-	flags = ptr;
-	ptr += flags_size;
+	ptr += dns_parse_string(ptr, &flags_size, &flags);
 	if (PAST_END_OF_RECORD) {
 		return NULL;
 	}
 
 	/* SERVICES */
-	services_size = *ptr;
-	++ptr;
-	if (PAST_END_OF_RECORD) {
-		return NULL;
-	}
-	services = ptr;
-	ptr += services_size;
+	ptr += dns_parse_string(ptr, &services_size, &services);
 	if (PAST_END_OF_RECORD) {
 		return NULL;
 	}
 
 	/* REGEXP */
-	regexp_size = *ptr;
-	++ptr;
-	if (PAST_END_OF_RECORD) {
-		return NULL;
-	}
-	regexp = ptr;
-	ptr += regexp_size;
+	ptr += dns_parse_string(ptr, &regexp_size, &regexp);
 	if (PAST_END_OF_RECORD) {
 		return NULL;
 	}
diff --git a/main/dns_srv.c b/main/dns_srv.c
index 7895073de..a617ede4d 100644
--- a/main/dns_srv.c
+++ b/main/dns_srv.c
@@ -41,59 +41,36 @@ ASTERISK_FILE_VERSION(__FILE__, "$Revision$")
 #include "asterisk/dns_internal.h"
 #include "asterisk/utils.h"
 
-struct ast_dns_record *ast_dns_srv_alloc(struct ast_dns_query *query, const char *data, const size_t size)
+struct ast_dns_record *dns_srv_alloc(struct ast_dns_query *query, const char *data, const size_t size)
 {
 	uint16_t priority;
 	uint16_t weight;
 	uint16_t port;
 	const char *ptr;
-	char *srv_offset;
-	char *srv_search_base = (char *)query->result->answer;
-	size_t remaining_size = query->result->answer_size;
 	const char *end_of_record;
 	struct ast_dns_srv_record *srv;
 	int host_size;
 	char host[NI_MAXHOST] = "";
 
-	while (1) {
-		srv_offset = memchr(srv_search_base, data[0], remaining_size);
-
-		ast_assert(srv_offset != NULL);
-		ast_assert(srv_search_base + remaining_size - srv_offset >= size);
-
-		if (!memcmp(srv_offset, data, size)) {
-			ptr = srv_offset;
-			break;
-		}
-
-		remaining_size -= srv_offset - srv_search_base;
-		srv_search_base = srv_offset + 1;
-	}
-
+	ptr = dns_find_record(data, size, query->result->answer, query->result->answer_size);
 	ast_assert(ptr != NULL);
 
 	end_of_record = ptr + size;
 
 	/* PRIORITY */
-	priority = ((unsigned char)(ptr[1]) << 0) | ((unsigned char)(ptr[0]) << 8);
-	ptr += 2;
-
+	ptr += dns_parse_short((unsigned char *) ptr, &priority);
 	if (ptr >= end_of_record) {
 		return NULL;
 	}
 
 	/* WEIGHT */
-	weight = ((unsigned char)(ptr[1]) << 0) | ((unsigned char)(ptr[0]) << 8);
-	ptr += 2;
-
+	ptr += dns_parse_short((unsigned char *) ptr, &weight);
 	if (ptr >= end_of_record) {
 		return NULL;
 	}
 
 	/* PORT */
-	port = ((unsigned char)(ptr[1]) << 0) | ((unsigned char)(ptr[0]) << 8);
-	ptr += 2;
-
+	ptr += dns_parse_short((unsigned char *) ptr, &port);
 	if (ptr >= end_of_record) {
 		return NULL;
 	}
@@ -129,7 +106,7 @@ struct ast_dns_record *ast_dns_srv_alloc(struct ast_dns_query *query, const char
 /* This implementation was taken from the existing srv.c which, after reading the RFC, implements it
  * as it should.
  */
-void ast_dns_srv_sort(struct ast_dns_result *result)
+void dns_srv_sort(struct ast_dns_result *result)
 {
 	struct ast_dns_record *current;
 	struct dns_records newlist = AST_LIST_HEAD_NOLOCK_INIT_VALUE;
diff --git a/main/dns_test.c b/main/dns_test.c
new file mode 100644
index 000000000..3d6169259
--- /dev/null
+++ b/main/dns_test.c
@@ -0,0 +1,265 @@
+/*
+ * Asterisk -- An open source telephony toolkit.
+ *
+ * Copyright (C) 2015, Digium, Inc.
+ *
+ * Mark Michelson <mmichelson@digium.com>
+ *
+ * Includes code and algorithms from the Zapata library.
+ *
+ * See http://www.asterisk.org for more information about
+ * the Asterisk project. Please do not directly contact
+ * any of the maintainers of this project for assistance;
+ * the project provides a web site, mailing lists and IRC
+ * channels for your use.
+ *
+ * This program is free software, distributed under the terms of
+ * the GNU General Public License Version 2. See the LICENSE file
+ * at the top of the source tree.
+ */
+
+/*** MODULEINFO
+	<support_level>core</support_level>
+ ***/
+
+#include "asterisk.h"
+#include "asterisk/dns_core.h"
+#include "asterisk/dns_test.h"
+#include "asterisk/utils.h"
+
+#ifdef TEST_FRAMEWORK
+
+const char DNS_HEADER[] = {
+	/* ID  == 0 */
+	0x00, 0x00,
+	/* QR == 1, Opcode == 0, AA == 1, TC == 0, RD == 1 */
+	0x85,
+	/* RA == 1, Z == 0, RCODE == 0 */
+	0x80,
+	/* QDCOUNT == 1 */
+	0x00, 0x01,
+	/* ANCOUNT == 1 */
+	0x00, 0x00,
+	/* NSCOUNT == 0 */
+	0x00, 0x00,
+	/* ARCOUNT == 0 */
+	0x00, 0x00,
+};
+
+/*!
+ * \brief Generate a DNS header and write it to a buffer
+ *
+ * The DNS header is the first part of a DNS request or response. In our
+ * case, the only part of the header that a test can affect is the number
+ * of answers. The rest of the DNS header is based on hard-coded values.
+ *
+ * There is no buffer size passed to this function since we provide
+ * the data ourselves and have sized the buffer to be way larger
+ * than necessary for the tests.
+ *
+ * \param num_records The number of DNS records in this DNS response
+ * \param buf The buffer to write the header into
+ * \retval The number of bytes written to the buffer
+ */
+static int generate_dns_header(unsigned short num_records, char *buf)
+{
+	unsigned short net_num_records = htons(num_records);
+
+	memcpy(buf, DNS_HEADER, ARRAY_LEN(DNS_HEADER));
+	/* Overwrite the ANCOUNT with the actual number of answers */
+	memcpy(&buf[6], &net_num_records, sizeof(num_records));
+
+	return ARRAY_LEN(DNS_HEADER);
+}
+
+const char DNS_QUESTION [] = {
+	/* goose */
+	0x05, 0x67, 0x6f, 0x6f, 0x73, 0x65,
+	/* feathers */
+	0x08, 0x66, 0x65, 0x61, 0x74, 0x68, 0x65, 0x72, 0x73,
+	/* end label */
+	0x00,
+	/* NAPTR type */
+	0x00, 0x23,
+	/* IN class */
+	0x00, 0x01,
+};
+
+/*!
+ * \brief Generate a DNS question and write it to a buffer
+ *
+ * The DNS question is the second part of a DNS request or response.
+ * All DNS questions in this file are for the same domain and thus
+ * the DNS question is a hard-coded value.
+ *
+ * There is no buffer size passed to this function since we provide
+ * the data ourselves and have sized the buffer to be way larger
+ * than necessary for the tests.
+ *
+ * \param buf The buffer to write the question into
+ * \retval The number of bytes written to the buffer
+ */
+static int generate_dns_question(char *buf)
+{
+	memcpy(buf, DNS_QUESTION, ARRAY_LEN(DNS_QUESTION));
+	return ARRAY_LEN(DNS_QUESTION);
+}
+
+const char NAPTR_ANSWER [] = {
+	/* Domain points to name from question */
+	0xc0, 0x0c,
+	/* NAPTR type */
+	0x00, 0x23,
+	/* IN Class */
+	0x00, 0x01,
+	/* TTL (12345 by default) */
+	0x00, 0x00, 0x30, 0x39,
+};
+
+/*!
+ * \brief Generate a DNS answer and write it to a buffer
+ *
+ * The DNS answer is the third (and in our case final) part of a
+ * DNS response. The DNS answer generated here is only partial.
+ * The record-specific data is generated by a separate function.
+ * DNS answers in our tests may have variable TTLs, but the rest
+ * is hard-coded.
+ *
+ * There is no buffer size passed to this function since we provide
+ * the data ourselves and have sized the buffer to be way larger
+ * than necessary for the tests.
+ *
+ * \param buf The buffer to write the answer into
+ * \retval The number of bytes written to the buffer
+ */
+static int generate_dns_answer(int ttl, char *buf)
+{
+	int net_ttl = htonl(ttl);
+
+	memcpy(buf, NAPTR_ANSWER, ARRAY_LEN(NAPTR_ANSWER));
+	/* Overwrite TTL if one is provided */
+	if (ttl) {
+		memcpy(&buf[6], &net_ttl, sizeof(int));
+	}
+
+	return ARRAY_LEN(NAPTR_ANSWER);
+}
+
+/*!
+ * \brief Write a DNS string to a buffer
+ *
+ * This writes the DNS string to the buffer and returns the total
+ * number of bytes written to the buffer.
+ *
+ * There is no buffer size passed to this function since we provide
+ * the data ourselves and have sized the buffer to be way larger
+ * than necessary for the tests.
+ *
+ * \param string The string to write
+ * \param buf The buffer to write the string into
+ * \return The number of bytes written to the buffer
+ */
+int ast_dns_test_write_string(const struct ast_dns_test_string *string, char *buf)
+{
+	uint8_t len = string->len;
+	size_t actual_len = strlen(string->val);
+	buf[0] = len;
+	/*
+	 * We use the actual length of the string instead of
+	 * the stated value since sometimes we're going to lie about
+	 * the length of the string
+	 */
+	if (actual_len) {
+		memcpy(&buf[1], string->val, strlen(string->val));
+	}
+
+	return actual_len + 1;
+}
+
+/*!
+ * \brief Write a DNS domain to a buffer
+ *
+ * A DNS domain consists of a series of labels separated
+ * by dots. Each of these labels gets written as a DNS
+ * string. A DNS domain ends with a NULL label, which is
+ * essentially a zero-length DNS string.
+ *
+ *
+ * There is no buffer size passed to this function since we provide
+ * the data ourselves and have sized the buffer to be way larger
+ * than necessary for the tests.
+ *
+ * \param string The DNS domain to write
+ * \param buf The buffer to write the domain into
+ * \return The number of bytes written to the buffer
+ */
+int ast_dns_test_write_domain(const char *string, char *buf)
+{
+	char *copy = ast_strdupa(string);
+	char *part;
+	char *ptr = buf;
+	static const struct ast_dns_test_string null_label = {
+		.len = 0,
+		.val = "",
+	};
+
+	while (1) {
+		struct ast_dns_test_string dns_str;
+		part = strsep(&copy, ".");
+		if (ast_strlen_zero(part)) {
+			break;
+		}
+		dns_str.len = strlen(part);
+		dns_str.val = part;
+
+		ptr += ast_dns_test_write_string(&dns_str, ptr);
+	}
+	ptr += ast_dns_test_write_string(&null_label, ptr);
+
+	return ptr - buf;
+}
+
+int ast_dns_test_generate_result(struct ast_dns_query *query, void *records, size_t num_records,
+		size_t record_size, record_fn generate, char *buffer)
+{
+	char *ptr = buffer;
+	char *record_iter;
+
+	ptr += generate_dns_header(num_records, ptr);
+	ptr += generate_dns_question(ptr);
+
+	for (record_iter = records; record_iter < (char *) records + num_records * record_size; record_iter += record_size) {
+		unsigned short rdlength;
+		unsigned short net_rdlength;
+
+		/* XXX Do we even want to override TTL? */
+		ptr += generate_dns_answer(0, ptr);
+		rdlength = generate(record_iter, ptr + 2);
+		net_rdlength = htons(rdlength);
+		memcpy(ptr, &net_rdlength, 2);
+		ptr += 2;
+		ptr += rdlength;
+	}
+
+	return ptr - buffer;
+}
+
+#else /* TEST_FRAMEWORK */
+
+int ast_dns_test_write_string(struct ast_dns_test_string *string, char *buf)
+{
+	return 0;
+}
+
+int ast_dns_test_write_domain(const char *string, char *buf)
+{
+	return 0;
+}
+
+int ast_dns_test_generate_result(struct ast_dns_query *query, void *records, size_t num_records,
+		size_t record_size, record_fn generate, char *buffer)
+{
+	return 0;
+}
+
+#endif
diff --git a/tests/test_dns_naptr.c b/tests/test_dns_naptr.c
index 1aaac674b..c7c7eb364 100644
--- a/tests/test_dns_naptr.c
+++ b/tests/test_dns_naptr.c
@@ -30,217 +30,15 @@
 #include "asterisk/dns_core.h"
 #include "asterisk/dns_resolver.h"
 #include "asterisk/dns_naptr.h"
-
-#define DNS_HEADER_SIZE 96
-
-const char DNS_HEADER[] = {
-	/* ID  == 0 */
-	0x00, 0x00,
-	/* QR == 1, Opcode == 0, AA == 1, TC == 0, RD == 1 */
-	0x85,
-	/* RA == 1, Z == 0, RCODE == 0 */
-	0x80,
-	/* QDCOUNT == 1 */
-	0x00, 0x01,
-	/* ANCOUNT == 1 */
-	0x00, 0x00,
-	/* NSCOUNT == 0 */
-	0x00, 0x00,
-	/* ARCOUNT == 0 */
-	0x00, 0x00,
-};
-
-/*!
- * \brief Generate a DNS header and write it to a buffer
- *
- * The DNS header is the first part of a DNS request or response. In our
- * case, the only part of the header that a test can affect is the number
- * of answers. The rest of the DNS header is based on hard-coded values.
- *
- * There is no buffer size passed to this function since we provide
- * the data ourselves and have sized the buffer to be way larger
- * than necessary for the tests.
- *
- * \param num_records The number of DNS records in this DNS response
- * \param buf The buffer to write the header into
- * \retval The number of bytes written to the buffer
- */
-static int generate_dns_header(unsigned short num_records, char *buf)
-{
-	unsigned short net_num_records = htons(num_records);
-
-	memcpy(buf, DNS_HEADER, ARRAY_LEN(DNS_HEADER));
-	/* Overwrite the ANCOUNT with the actual number of answers */
-	memcpy(&buf[6], &net_num_records, sizeof(num_records));
-
-	return ARRAY_LEN(DNS_HEADER);
-}
-
-const char DNS_QUESTION [] = {
-	/* goose */
-	0x05, 0x67, 0x6f, 0x6f, 0x73, 0x65,
-	/* feathers */
-	0x08, 0x66, 0x65, 0x61, 0x74, 0x68, 0x65, 0x72, 0x73,
-	/* end label */
-	0x00,
-	/* NAPTR type */
-	0x00, 0x23,
-	/* IN class */
-	0x00, 0x01,
-};
-
-/*!
- * \brief Generate a DNS question and write it to a buffer
- *
- * The DNS question is the second part of a DNS request or response.
- * All DNS questions in this file are for the same domain and thus
- * the DNS question is a hard-coded value.
- *
- * There is no buffer size passed to this function since we provide
- * the data ourselves and have sized the buffer to be way larger
- * than necessary for the tests.
- *
- * \param buf The buffer to write the question into
- * \retval The number of bytes written to the buffer
- */
-static int generate_dns_question(char *buf)
-{
-	memcpy(buf, DNS_QUESTION, ARRAY_LEN(DNS_QUESTION));
-	return ARRAY_LEN(DNS_QUESTION);
-}
-
-const char NAPTR_ANSWER [] = {
-	/* Domain points to name from question */
-	0xc0, 0x0c,
-	/* NAPTR type */
-	0x00, 0x23,
-	/* IN Class */
-	0x00, 0x01,
-	/* TTL (12345 by default) */
-	0x00, 0x00, 0x30, 0x39,
-};
-
-/*!
- * \brief Generate a DNS answer and write it to a buffer
- *
- * The DNS answer is the third (and in our case final) part of a
- * DNS response. The DNS answer generated here is only partial.
- * The record-specific data is generated by a separate function.
- * DNS answers in our tests may have variable TTLs, but the rest
- * is hard-coded.
- *
- * There is no buffer size passed to this function since we provide
- * the data ourselves and have sized the buffer to be way larger
- * than necessary for the tests.
- *
- * \param buf The buffer to write the answer into
- * \retval The number of bytes written to the buffer
- */
-static int generate_dns_answer(int ttl, char *buf)
-{
-	int net_ttl = htonl(ttl);
-
-	memcpy(buf, NAPTR_ANSWER, ARRAY_LEN(NAPTR_ANSWER));
-	/* Overwrite TTL if one is provided */
-	if (ttl) {
-		memcpy(&buf[6], &net_ttl, sizeof(int));
-	}
-
-	return ARRAY_LEN(NAPTR_ANSWER);
-}
-
-/*!
- * \brief Representation of a string in DNS
- *
- * In DNS, a string has a byte to indicate the length,
- * followed by a series of bytes representing the string.
- * DNS does not NULL-terminate its strings.
- */
-struct dns_string {
-	uint8_t len;
-	const char *val;
-};
-
-/*!
- * \brief Write a DNS string to a buffer
- *
- * This writes the DNS string to the buffer and returns the total
- * number of bytes written to the buffer.
- *
- * There is no buffer size passed to this function since we provide
- * the data ourselves and have sized the buffer to be way larger
- * than necessary for the tests.
- *
- * \param string The string to write
- * \param buf The buffer to write the string into
- * \return The number of bytes written to the buffer
- */
-static int write_dns_string(const struct dns_string *string, char *buf)
-{
-	uint8_t len = string->len;
-	buf[0] = len;
-	/*
-	 * We use the actual length of the string instead of
-	 * the stated value since sometimes we're going to lie about
-	 * the length of the string
-	 */
-	if (strlen(string->val)) {
-		memcpy(&buf[1], string->val, strlen(string->val));
-	}
-
-	return strlen(string->val) + 1;
-}
-
-/*!
- * \brief Write a DNS domain to a buffer
- *
- * A DNS domain consists of a series of labels separated
- * by dots. Each of these labels gets written as a DNS
- * string. A DNS domain ends with a NULL label, which is
- * essentially a zero-length DNS string.
- *
- *
- * There is no buffer size passed to this function since we provide
- * the data ourselves and have sized the buffer to be way larger
- * than necessary for the tests.
- *
- * \param string The DNS domain to write
- * \param buf The buffer to write the domain into
- * \return The number of bytes written to the buffer
- */
-static int write_dns_domain(const char *string, char *buf)
-{
-	char *copy = ast_strdupa(string);
-	char *part;
-	char *ptr = buf;
-	static const struct dns_string null_label = {
-		.len = 0,
-		.val = "",
-	};
-
-	while (1) {
-		struct dns_string dns_str;
-		part = strsep(&copy, ".");
-		if (ast_strlen_zero(part)) {
-			break;
-		}
-		dns_str.len = strlen(part);
-		dns_str.val = part;
-
-		ptr += write_dns_string(&dns_str, ptr);
-	}
-	ptr += write_dns_string(&null_label, ptr);
-
-	return ptr - buf;
-}
+#include "asterisk/dns_test.h"
 
 struct naptr_record {
 	uint16_t order;
 	uint16_t preference;
-	struct dns_string flags;
-	struct dns_string services;
-	struct dns_string regexp;
-	const char * replacement;
+	struct ast_dns_test_string flags;
+	struct ast_dns_test_string services;
+	struct ast_dns_test_string regexp;
+	const char *replacement;
 };
 
 /*!
@@ -257,8 +55,9 @@ struct naptr_record {
  * \param buf The buffer to write the record into
  * \return The number of bytes written to the buffer
  */
-static int generate_naptr_record(struct naptr_record *record, char *buf)
+static int generate_naptr_record(void *dns_record, char *buf)
 {
+	struct naptr_record *record = dns_record;
 	uint16_t net_order = htons(record->order);
 	uint16_t net_preference = htons(record->preference);
 	char *ptr = buf;
@@ -269,10 +68,10 @@ static int generate_naptr_record(struct naptr_record *record, char *buf)
 	memcpy(ptr, &net_preference, sizeof(net_preference));
 	ptr += sizeof(net_preference);
 
-	ptr += write_dns_string(&record->flags, ptr);
-	ptr += write_dns_string(&record->services, ptr);
-	ptr += write_dns_string(&record->regexp, ptr);
-	ptr += write_dns_domain(record->replacement, ptr);
+	ptr += ast_dns_test_write_string(&record->flags, ptr);
+	ptr += ast_dns_test_write_string(&record->services, ptr);
+	ptr += ast_dns_test_write_string(&record->regexp, ptr);
+	ptr += ast_dns_test_write_domain(record->replacement, ptr);
 
 	return ptr - buf;
 }
@@ -312,31 +111,19 @@ static void *naptr_thread(void *dns_query)
 {
 	struct ast_dns_query *query = dns_query;
 	int i;
-	char *ptr = ans_buffer;
-
-	ptr += generate_dns_header(num_test_records, ptr);
-	ptr += generate_dns_question(ptr);
+	int ans_size;
 
-	for (i = 0; i < num_test_records; ++i) {
-		unsigned short rdlength;
-		unsigned short net_rdlength;
-
-		ptr += generate_dns_answer(0, ptr);
-		rdlength = generate_naptr_record(&test_records[i], ptr + 2);
-		net_rdlength = htons(rdlength);
-		memcpy(ptr, &net_rdlength, 2);
-		ptr += 2;
-		ptr += rdlength;
-	}
+	ans_size = ast_dns_test_generate_result(query, test_records, num_test_records,
+			sizeof(struct naptr_record), generate_naptr_record, ans_buffer);
 
-	ast_dns_resolver_set_result(query, 0, 0, ns_r_noerror, "goose.feathers", ans_buffer, ptr - ans_buffer);
+	ast_dns_resolver_set_result(query, 0, 0, ns_r_noerror, "goose.feathers", ans_buffer, ans_size);
 
 	for (i = 0; i < num_test_records; ++i) {
 		char record[128];
-		ptr = record;
+		int naptr_size;
 
-		ptr += generate_naptr_record(&test_records[i], ptr);
-		ast_dns_resolver_add_record(query, ns_t_naptr, ns_c_in, 12345, record, ptr - record);
+		naptr_size = generate_naptr_record(&test_records[i], record);
+		ast_dns_resolver_add_record(query, ns_t_naptr, ns_c_in, 12345, record, naptr_size);
 	}
 
 	ast_dns_resolver_completed(query);
diff --git a/tests/test_dns_srv.c b/tests/test_dns_srv.c
index de79be1d9..c5145bb2e 100644
--- a/tests/test_dns_srv.c
+++ b/tests/test_dns_srv.c
@@ -31,104 +31,7 @@
 #include "asterisk/dns_core.h"
 #include "asterisk/dns_resolver.h"
 #include "asterisk/dns_srv.h"
-
-#define DNS_HEADER_SIZE 96
-
-const char DNS_HEADER[] = {
-	/* ID  == 0 */
-	0x00, 0x00,
-	/* QR == 1, Opcode == 0, AA == 1, TC == 0, RD == 1 */
-	0x85,
-	/* RA == 1, Z == 0, RCODE == 0 */
-	0x80,
-	/* QDCOUNT == 1 */
-	0x00, 0x01,
-	/* ANCOUNT == 1 */
-	0x00, 0x00,
-	/* NSCOUNT == 0 */
-	0x00, 0x00,
-	/* ARCOUNT == 0 */
-	0x00, 0x00,
-};
-
-static int generate_dns_header(unsigned short num_records, char *buf)
-{
-	unsigned short net_num_records = htons(num_records);
-
-	memcpy(buf, DNS_HEADER, ARRAY_LEN(DNS_HEADER));
-	/* Overwrite the ANCOUNT with the actual number of answers */
-	memcpy(&buf[6], &net_num_records, sizeof(num_records));
-
-	return ARRAY_LEN(DNS_HEADER);
-}
-
-const char DNS_QUESTION [] = {
-	/* goose */
-	0x05, 0x67, 0x6f, 0x6f, 0x73, 0x65,
-	/* feathers */
-	0x08, 0x66, 0x65, 0x61, 0x74, 0x68, 0x65, 0x72, 0x73,
-	/* end label */
-	0x00,
-	/* SRV type */
-	0x00, 0x23,
-	/* IN class */
-	0x00, 0x01,
-};
-
-static int generate_dns_question(char *buf)
-{
-	memcpy(buf, DNS_QUESTION, ARRAY_LEN(DNS_QUESTION));
-	return ARRAY_LEN(DNS_QUESTION);
-}
-
-const char SRV_ANSWER [] = {
-	/* Domain points to name from question */
-	0xc0, 0x0c,
-	/* NAPTR type */
-	0x00, 0x23,
-	/* IN Class */
-	0x00, 0x01,
-	/* TTL (12345 by default) */
-	0x00, 0x00, 0x30, 0x39,
-};
-
-static int generate_dns_answer(int ttl, char *buf)
-{
-	int net_ttl = htonl(ttl);
-
-	memcpy(buf, SRV_ANSWER, ARRAY_LEN(SRV_ANSWER));
-	/* Overwrite TTL if one is provided */
-	if (ttl) {
-		memcpy(&buf[6], &net_ttl, sizeof(int));
-	}
-
-	return ARRAY_LEN(SRV_ANSWER);
-}
-
-static int write_dns_string(const char *string, char *buf)
-{
-	uint8_t len = strlen(string);
-	buf[0] = len;
-	if (len) {
-		memcpy(&buf[1], string, len);
-	}
-
-	return len + 1;
-}
-
-static int write_dns_domain(const char *string, char *buf)
-{
-	char *copy = ast_strdupa(string);
-	char *part;
-	char *ptr = buf;
-
-	while ((part = strsep(&copy, "."))) {
-		ptr += write_dns_string(part, ptr);
-	}
-	ptr += write_dns_string("", ptr);
-
-	return ptr - buf;
-}
+#include "asterisk/dns_test.h"
 
 struct srv_record {
 	uint16_t priority;
@@ -141,8 +44,9 @@ struct srv_record {
 	unsigned int ignore_host;
 };
 
-static int generate_srv_record(struct srv_record *record, char *buf)
+static int generate_srv_record(void *dns_record, char *buf)
 {
+	struct srv_record *record = dns_record;
 	uint16_t priority = htons(record->priority);
 	uint16_t weight = htons(record->weight);
 	uint16_t port = htons(record->port);
@@ -164,7 +68,7 @@ static int generate_srv_record(struct srv_record *record, char *buf)
 	}
 
 	if (!record->ignore_host) {
-		ptr += write_dns_domain(record->host, ptr);
+		ptr += ast_dns_test_write_domain(record->host, ptr);
 	}
 
 	return ptr - buf;
@@ -178,31 +82,19 @@ static void *srv_thread(void *dns_query)
 {
 	struct ast_dns_query *query = dns_query;
 	int i;
-	char *ptr = ans_buffer;
-
-	ptr += generate_dns_header(num_test_records, ptr);
-	ptr += generate_dns_question(ptr);
+	int ans_size;
 
-	for (i = 0; i < num_test_records; ++i) {
-		unsigned short rdlength;
-		unsigned short net_rdlength;
-
-		ptr += generate_dns_answer(0, ptr);
-		rdlength = generate_srv_record(&test_records[i], ptr + 2);
-		net_rdlength = htons(rdlength);
-		memcpy(ptr, &net_rdlength, 2);
-		ptr += 2;
-		ptr += rdlength;
-	}
+	ans_size = ast_dns_test_generate_result(query, test_records, num_test_records,
+			sizeof(struct srv_record), generate_srv_record, ans_buffer);
 
-	ast_dns_resolver_set_result(query, 0, 0, ns_r_noerror, "goose.feathers", ans_buffer, ptr - ans_buffer);
+	ast_dns_resolver_set_result(query, 0, 0, ns_r_noerror, "goose.feathers", ans_buffer, ans_size);
 
 	for (i = 0; i < num_test_records; ++i) {
 		char record[128];
-		ptr = record;
+		int srv_size;
 
-		ptr += generate_srv_record(&test_records[i], ptr);
-		ast_dns_resolver_add_record(query, ns_t_srv, ns_c_in, 12345, record, ptr - record);
+		srv_size = generate_srv_record(&test_records[i], record);
+		ast_dns_resolver_add_record(query, ns_t_srv, ns_c_in, 12345, record, srv_size);
 	}
 
 	ast_dns_resolver_completed(query);