Merge pull request #16238 from keszybz/set-handling-more

Fix handling of cases where a duplicate item is added to a set and related cleanups

8 files changed, 45 insertions, 111 deletions

diff --git a/src/resolve/resolved-dns-query.c b/src/resolve/resolved-dns-query.c
index 914f464dd7..906158c5ce 100644
--- a/src/resolve/resolved-dns-query.c
+++ b/src/resolve/resolved-dns-query.c
@@ -94,7 +94,7 @@ static int dns_query_candidate_next_search_domain(DnsQueryCandidate *c) {
 }
 
 static int dns_query_candidate_add_transaction(DnsQueryCandidate *c, DnsResourceKey *key) {
-        DnsTransaction *t;
+        _cleanup_(dns_transaction_gcp) DnsTransaction *t = NULL;
         int r;
 
         assert(c);
@@ -105,39 +105,26 @@ static int dns_query_candidate_add_transaction(DnsQueryCandidate *c, DnsResource
                 r = dns_transaction_new(&t, c->scope, key);
                 if (r < 0)
                         return r;
-        } else {
-                if (set_contains(c->transactions, t))
-                        return 0;
-        }
-
-        r = set_ensure_allocated(&c->transactions, NULL);
-        if (r < 0)
-                goto gc;
-
-        r = set_ensure_allocated(&t->notify_query_candidates, NULL);
-        if (r < 0)
-                goto gc;
+        } else if (set_contains(c->transactions, t))
+                return 0;
 
         r = set_ensure_allocated(&t->notify_query_candidates_done, NULL);
         if (r < 0)
-                goto gc;
+                return r;
 
-        r = set_put(t->notify_query_candidates, c);
+        r = set_ensure_put(&t->notify_query_candidates, NULL, c);
         if (r < 0)
-                goto gc;
+                return r;
 
-        r = set_put(c->transactions, t);
+        r = set_ensure_put(&c->transactions, NULL, t);
         if (r < 0) {
                 (void) set_remove(t->notify_query_candidates, c);
-                goto gc;
+                return r;
         }
 
         t->clamp_ttl = c->query->clamp_ttl;
+        TAKE_PTR(t);
         return 1;
-
-gc:
-        dns_transaction_gc(t);
-        return r;
 }
 
 static int dns_query_candidate_go(DnsQueryCandidate *c) {
diff --git a/src/resolve/resolved-dns-scope.c b/src/resolve/resolved-dns-scope.c
index 1a5fef13dc..764ccee0e0 100644
--- a/src/resolve/resolved-dns-scope.c
+++ b/src/resolve/resolved-dns-scope.c
@@ -1254,11 +1254,7 @@ int dns_scope_announce(DnsScope *scope, bool goodbye) {
                 if (!scope->announced &&
                     dns_resource_key_is_dnssd_ptr(z->rr->key)) {
                         if (!set_contains(types, dns_resource_key_name(z->rr->key))) {
-                                r = set_ensure_allocated(&types, &dns_name_hash_ops);
-                                if (r < 0)
-                                        return log_debug_errno(r, "Failed to allocate set: %m");
-
-                                r = set_put(types, dns_resource_key_name(z->rr->key));
+                                r = set_ensure_put(&types, &dns_name_hash_ops, dns_resource_key_name(z->rr->key));
                                 if (r < 0)
                                         return log_debug_errno(r, "Failed to add item to set: %m");
                         }
diff --git a/src/resolve/resolved-dns-stub.c b/src/resolve/resolved-dns-stub.c
index ce994a7ee0..03edbe26dc 100644
--- a/src/resolve/resolved-dns-stub.c
+++ b/src/resolve/resolved-dns-stub.c
@@ -278,7 +278,7 @@ static int dns_stub_stream_complete(DnsStream *s, int error) {
 }
 
 static void dns_stub_process_query(Manager *m, DnsStream *s, DnsPacket *p) {
-        DnsQuery *q = NULL;
+        _cleanup_(dns_query_freep) DnsQuery *q = NULL;
         int r;
 
         assert(m);
@@ -289,52 +289,52 @@ static void dns_stub_process_query(Manager *m, DnsStream *s, DnsPacket *p) {
             in_addr_is_localhost(p->family, &p->destination) <= 0) {
                 log_error("Got packet on unexpected IP range, refusing.");
                 dns_stub_send_failure(m, s, p, DNS_RCODE_SERVFAIL, false);
-                goto fail;
+                return;
         }
 
         r = dns_packet_extract(p);
         if (r < 0) {
                 log_debug_errno(r, "Failed to extract resources from incoming packet, ignoring packet: %m");
                 dns_stub_send_failure(m, s, p, DNS_RCODE_FORMERR, false);
-                goto fail;
+                return;
         }
 
         if (!DNS_PACKET_VERSION_SUPPORTED(p)) {
                 log_debug("Got EDNS OPT field with unsupported version number.");
                 dns_stub_send_failure(m, s, p, DNS_RCODE_BADVERS, false);
-                goto fail;
+                return;
         }
 
         if (dns_type_is_obsolete(p->question->keys[0]->type)) {
                 log_debug("Got message with obsolete key type, refusing.");
                 dns_stub_send_failure(m, s, p, DNS_RCODE_NOTIMP, false);
-                goto fail;
+                return;
         }
 
         if (dns_type_is_zone_transer(p->question->keys[0]->type)) {
                 log_debug("Got request for zone transfer, refusing.");
                 dns_stub_send_failure(m, s, p, DNS_RCODE_NOTIMP, false);
-                goto fail;
+                return;
         }
 
         if (!DNS_PACKET_RD(p))  {
                 /* If the "rd" bit is off (i.e. recursion was not requested), then refuse operation */
                 log_debug("Got request with recursion disabled, refusing.");
                 dns_stub_send_failure(m, s, p, DNS_RCODE_REFUSED, false);
-                goto fail;
+                return;
         }
 
         if (DNS_PACKET_DO(p) && DNS_PACKET_CD(p)) {
                 log_debug("Got request with DNSSEC CD bit set, refusing.");
                 dns_stub_send_failure(m, s, p, DNS_RCODE_NOTIMP, false);
-                goto fail;
+                return;
         }
 
         r = dns_query_new(m, &q, p->question, p->question, 0, SD_RESOLVED_PROTOCOLS_ALL|SD_RESOLVED_NO_SEARCH);
         if (r < 0) {
                 log_error_errno(r, "Failed to generate query object: %m");
                 dns_stub_send_failure(m, s, p, DNS_RCODE_SERVFAIL, false);
-                goto fail;
+                return;
         }
 
         /* Request that the TTL is corrected by the cached time for this lookup, so that we return vaguely useful TTLs */
@@ -348,30 +348,23 @@ static void dns_stub_process_query(Manager *m, DnsStream *s, DnsPacket *p) {
                 /* Remember which queries belong to this stream, so that we can cancel them when the stream
                  * is disconnected early */
 
-                r = set_ensure_allocated(&s->queries, &trivial_hash_ops);
+                r = set_ensure_put(&s->queries, NULL, q);
                 if (r < 0) {
                         log_oom();
-                        goto fail;
-                }
-
-                if (set_put(s->queries, q) < 0) {
-                        log_oom();
-                        goto fail;
+                        return;
                 }
+                assert(r > 0);
         }
 
         r = dns_query_go(q);
         if (r < 0) {
                 log_error_errno(r, "Failed to start query: %m");
                 dns_stub_send_failure(m, s, p, DNS_RCODE_SERVFAIL, false);
-                goto fail;
+                return;
         }
 
         log_debug("Processing query...");
-        return;
-
-fail:
-        dns_query_free(q);
+        TAKE_PTR(q);
 }
 
 static int on_dns_stub_packet(sd_event_source *s, int fd, uint32_t revents, void *userdata) {
diff --git a/src/resolve/resolved-dns-transaction.c b/src/resolve/resolved-dns-transaction.c
index 5898308d5f..cd5a0e3dd9 100644
--- a/src/resolve/resolved-dns-transaction.c
+++ b/src/resolve/resolved-dns-transaction.c
@@ -1501,11 +1501,7 @@ static int dns_transaction_make_packet_mdns(DnsTransaction *t) {
                 add_known_answers = true;
 
         if (t->key->type == DNS_TYPE_ANY) {
-                r = set_ensure_allocated(&keys, &dns_resource_key_hash_ops);
-                if (r < 0)
-                        return r;
-
-                r = set_put(keys, t->key);
+                r = set_ensure_put(&keys, &dns_resource_key_hash_ops, t->key);
                 if (r < 0)
                         return r;
         }
@@ -1571,11 +1567,7 @@ static int dns_transaction_make_packet_mdns(DnsTransaction *t) {
                         add_known_answers = true;
 
                 if (other->key->type == DNS_TYPE_ANY) {
-                        r = set_ensure_allocated(&keys, &dns_resource_key_hash_ops);
-                        if (r < 0)
-                                return r;
-
-                        r = set_put(keys, other->key);
+                        r = set_ensure_put(&keys, &dns_resource_key_hash_ops, other->key);
                         if (r < 0)
                                 return r;
                 }
@@ -1800,7 +1792,7 @@ static int dns_transaction_find_cyclic(DnsTransaction *t, DnsTransaction *aux) {
 }
 
 static int dns_transaction_add_dnssec_transaction(DnsTransaction *t, DnsResourceKey *key, DnsTransaction **ret) {
-        DnsTransaction *aux;
+        _cleanup_(dns_transaction_gcp) DnsTransaction *aux = NULL;
         int r;
 
         assert(t);
@@ -1833,34 +1825,22 @@ static int dns_transaction_add_dnssec_transaction(DnsTransaction *t, DnsResource
                 }
         }
 
-        r = set_ensure_allocated(&t->dnssec_transactions, NULL);
-        if (r < 0)
-                goto gc;
-
-        r = set_ensure_allocated(&aux->notify_transactions, NULL);
-        if (r < 0)
-                goto gc;
-
         r = set_ensure_allocated(&aux->notify_transactions_done, NULL);
         if (r < 0)
-                goto gc;
+                return r;
 
-        r = set_put(t->dnssec_transactions, aux);
+        r = set_ensure_put(&t->dnssec_transactions, NULL, aux);
         if (r < 0)
-                goto gc;
+                return r;;
 
-        r = set_put(aux->notify_transactions, t);
+        r = set_ensure_put(&aux->notify_transactions, NULL, t);
         if (r < 0) {
                 (void) set_remove(t->dnssec_transactions, aux);
-                goto gc;
+                return r;
         }
 
-        *ret = aux;
+        *ret = TAKE_PTR(aux);
         return 1;
-
-gc:
-        dns_transaction_gc(aux);
-        return r;
 }
 
 static int dns_transaction_request_dnssec_rr(DnsTransaction *t, DnsResourceKey *key) {
diff --git a/src/resolve/resolved-dns-transaction.h b/src/resolve/resolved-dns-transaction.h
index b1d4348409..167541806a 100644
--- a/src/resolve/resolved-dns-transaction.h
+++ b/src/resolve/resolved-dns-transaction.h
@@ -138,6 +138,8 @@ int dns_transaction_new(DnsTransaction **ret, DnsScope *s, DnsResourceKey *key);
 DnsTransaction* dns_transaction_free(DnsTransaction *t);
 
 bool dns_transaction_gc(DnsTransaction *t);
+DEFINE_TRIVIAL_CLEANUP_FUNC(DnsTransaction*, dns_transaction_gc);
+
 int dns_transaction_go(DnsTransaction *t);
 
 void dns_transaction_process_reply(DnsTransaction *t, DnsPacket *p);
diff --git a/src/resolve/resolved-dns-trust-anchor.c b/src/resolve/resolved-dns-trust-anchor.c
index 843f4c0f45..d68d0c3ba1 100644
--- a/src/resolve/resolved-dns-trust-anchor.c
+++ b/src/resolve/resolved-dns-trust-anchor.c
@@ -393,15 +393,9 @@ static int dns_trust_anchor_load_negative(DnsTrustAnchor *d, const char *path, u
                 return -EINVAL;
         }
 
-        r = set_ensure_allocated(&d->negative_by_name, &dns_name_hash_ops);
-        if (r < 0)
-                return log_oom();
-
-        r = set_put(d->negative_by_name, domain);
+        r = set_ensure_consume(&d->negative_by_name, &dns_name_hash_ops, TAKE_PTR(domain));
         if (r < 0)
                 return log_oom();
-        if (r > 0)
-                domain = NULL;
 
         return 0;
 }
@@ -592,11 +586,7 @@ static int dns_trust_anchor_revoked_put(DnsTrustAnchor *d, DnsResourceRecord *rr
 
         assert(d);
 
-        r = set_ensure_allocated(&d->revoked_by_rr, &dns_resource_record_hash_ops);
-        if (r < 0)
-                return r;
-
-        r = set_put(d->revoked_by_rr, rr);
+        r = set_ensure_put(&d->revoked_by_rr, &dns_resource_record_hash_ops, rr);
         if (r < 0)
                 return r;
         if (r > 0)
diff --git a/src/resolve/resolved-dns-zone.c b/src/resolve/resolved-dns-zone.c
index 0ef4c892f7..33879d6142 100644
--- a/src/resolve/resolved-dns-zone.c
+++ b/src/resolve/resolved-dns-zone.c
@@ -162,7 +162,7 @@ static int dns_zone_link_item(DnsZone *z, DnsZoneItem *i) {
 }
 
 static int dns_zone_item_probe_start(DnsZoneItem *i)  {
-        DnsTransaction *t;
+        _cleanup_(dns_transaction_gcp) DnsTransaction *t = NULL;
         int r;
 
         assert(i);
@@ -183,25 +183,20 @@ static int dns_zone_item_probe_start(DnsZoneItem *i)  {
                         return r;
         }
 
-        r = set_ensure_allocated(&t->notify_zone_items, NULL);
-        if (r < 0)
-                goto gc;
-
         r = set_ensure_allocated(&t->notify_zone_items_done, NULL);
         if (r < 0)
-                goto gc;
+                return r;
 
-        r = set_put(t->notify_zone_items, i);
+        r = set_ensure_put(&t->notify_zone_items, NULL, i);
         if (r < 0)
-                goto gc;
+                return r;
 
-        i->probe_transaction = t;
         t->probing = true;
+        i->probe_transaction = TAKE_PTR(t);
 
-        if (t->state == DNS_TRANSACTION_NULL) {
-
+        if (i->probe_transaction->state == DNS_TRANSACTION_NULL) {
                 i->block_ready++;
-                r = dns_transaction_go(t);
+                r = dns_transaction_go(i->probe_transaction);
                 i->block_ready--;
 
                 if (r < 0) {
@@ -212,10 +207,6 @@ static int dns_zone_item_probe_start(DnsZoneItem *i)  {
 
         dns_zone_item_notify(i);
         return 0;
-
-gc:
-        dns_transaction_gc(t);
-        return r;
 }
 
 int dns_zone_put(DnsZone *z, DnsScope *s, DnsResourceRecord *rr, bool probe) {
diff --git a/src/resolve/resolved-etc-hosts.c b/src/resolve/resolved-etc-hosts.c
index 2cb06c098d..6a7f749957 100644
--- a/src/resolve/resolved-etc-hosts.c
+++ b/src/resolve/resolved-etc-hosts.c
@@ -120,15 +120,10 @@ static int parse_line(EtcHosts *hosts, unsigned nr, const char *line) {
                         /* Optimize the case where we don't need to store any addresses, by storing
                          * only the name in a dedicated Set instead of the hashmap */
 
-                        r = set_ensure_allocated(&hosts->no_address, &dns_name_hash_ops);
-                        if (r < 0)
-                                return log_oom();
-
-                        r = set_put(hosts->no_address, name);
+                        r = set_ensure_consume(&hosts->no_address, &dns_name_hash_ops, TAKE_PTR(name));
                         if (r < 0)
                                 return r;
 
-                        TAKE_PTR(name);
                         continue;
                 }