dns: Make query sets hold on to queries for their lifetime.

The query set documentation states that upon completion queries can be
retrieved for the lifetime of the query set. This is a reasonable
expectation but does not currently occur. This was originally done
to resolve a circular reference between queries and query sets, but
in practice the query can be kept.

This change makes it so a query does not have a reference to the
query set until it begins resolving. It also makes it so that the
reference is given up upon the query being completed. This allows
the queries to remain for the lifetime of the query set. As the
query set on the query is only useful to the query set functionality
and only for the lifetime that the query is resolving this is safe
to do.

ASTERISK-24994 #close
Reported by: Joshua Colp

Change-Id: I54e09c0cb45475896654e7835394524e816d1aa0

1 files changed, 15 insertions, 12 deletions

diff --git a/main/dns_query_set.c b/main/dns_query_set.c
index c7a4eb18e..8dfc5eaeb 100644
--- a/main/dns_query_set.c
+++ b/main/dns_query_set.c
@@ -43,9 +43,10 @@ ASTERISK_REGISTER_FILE()
 /*! \brief The default number of expected queries to be added to the query set */
 #define DNS_QUERY_SET_EXPECTED_QUERY_COUNT 5
 
-/*! \brief Release all queries held in a query set */
-static void dns_query_set_release(struct ast_dns_query_set *query_set)
+/*! \brief Destructor for DNS query set */
+static void dns_query_set_destroy(void *data)
 {
+	struct ast_dns_query_set *query_set = data;
 	int idx;
 
 	for (idx = 0; idx < AST_VECTOR_SIZE(&query_set->queries); ++idx) {
@@ -53,16 +54,8 @@ static void dns_query_set_release(struct ast_dns_query_set *query_set)
 
 		ao2_ref(query->query, -1);
 	}
-
 	AST_VECTOR_FREE(&query_set->queries);
-}
-
-/*! \brief Destructor for DNS query set */
-static void dns_query_set_destroy(void *data)
-{
-	struct ast_dns_query_set *query_set = data;
 
-	dns_query_set_release(query_set);
 	ao2_cleanup(query_set->user_data);
 }
 
@@ -88,7 +81,15 @@ static void dns_query_set_callback(const struct ast_dns_query *query)
 {
 	struct ast_dns_query_set *query_set = ast_dns_query_get_data(query);
 
+	/* The reference count of the query set is bumped here in case this query holds the last reference */
+	ao2_ref(query_set, +1);
+
+	/* Drop the query set from the query so the query set can be destroyed if this is the last one */
+	ao2_cleanup(((struct ast_dns_query *)query)->user_data);
+	((struct ast_dns_query *)query)->user_data = NULL;
+
 	if (ast_atomic_fetchadd_int(&query_set->queries_completed, +1) != (AST_VECTOR_SIZE(&query_set->queries) - 1)) {
+		ao2_ref(query_set, -1);
 		return;
 	}
 
@@ -100,7 +101,7 @@ static void dns_query_set_callback(const struct ast_dns_query *query)
 	ao2_cleanup(query_set->user_data);
 	query_set->user_data = NULL;
 
-	dns_query_set_release(query_set);
+	ao2_ref(query_set, -1);
 }
 
 int ast_dns_query_set_add(struct ast_dns_query_set *query_set, const char *name, int rr_type, int rr_class)
@@ -116,7 +117,7 @@ int ast_dns_query_set_add(struct ast_dns_query_set *query_set, const char *name,
 		return -1;
 	}
 
-	query.query = dns_query_alloc(name, rr_type, rr_class, dns_query_set_callback, query_set);
+	query.query = dns_query_alloc(name, rr_type, rr_class, dns_query_set_callback, NULL);
 	if (!query.query) {
 		return -1;
 	}
@@ -169,6 +170,8 @@ void ast_dns_query_set_resolve_async(struct ast_dns_query_set *query_set, ast_dn
 	for (idx = 0; idx < AST_VECTOR_SIZE(&query_set->queries); ++idx) {
 		struct dns_query_set_query *query = AST_VECTOR_GET_ADDR(&query_set->queries, idx);
 
+		query->query->user_data = ao2_bump(query_set);
+
 		if (!query->query->resolver->resolve(query->query)) {
 			query->started = 1;
 			continue;