diff --git a/main/dns_query_set.c b/main/dns_query_set.c index c7a4eb18ebefaf0f831d3b0ceaf71691efa25000..8dfc5eaebe5d387cdfd6bb6aa8a6480f64da1a43 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; diff --git a/tests/test_dns_query_set.c b/tests/test_dns_query_set.c index 08829f59e9b5a9af6eb01d895353ec5d766d3c63..9acf1bf811ab6022b3d1bf9aa09e7db7c5e7a829 100644 --- a/tests/test_dns_query_set.c +++ b/tests/test_dns_query_set.c @@ -131,9 +131,14 @@ static int query_set_resolve(struct ast_dns_query *query) static int query_set_cancel(struct ast_dns_query *query) { struct ast_dns_query_set *query_set = ast_dns_query_get_data(query); - struct query_set_data *qsdata = query_set->user_data; + struct query_set_data *qsdata; int res = -1; + if (!query_set) { + return -1; + } + qsdata = query_set->user_data; + if (qsdata->cancel++ < qsdata->cancel_allowed) { res = 0; }