Merge pull request #23361 from keszybz/resolved-helpers

Add some ref-unref helpers for resolved

11 files changed, 77 insertions, 64 deletions

diff --git a/src/basic/macro.h b/src/basic/macro.h
index 3bf982803d..e6f89608f4 100644
--- a/src/basic/macro.h
+++ b/src/basic/macro.h
@@ -457,8 +457,15 @@ static inline int __coverity_check_and_return__(int condition) {
                 _copy;                                                  \
         })
 
+#define saturate_add(x, y, limit)                                       \
+        ({                                                              \
+                typeof(limit) _x = (x);                                 \
+                typeof(limit) _y = (y);                                 \
+                _x > (limit) || _y >= (limit) - _x ? (limit) : _x + _y; \
+        })
+
 static inline size_t size_add(size_t x, size_t y) {
-        return y >= SIZE_MAX - x ? SIZE_MAX : x + y;
+        return saturate_add(x, y, SIZE_MAX);
 }
 
 typedef struct {
diff --git a/src/resolve/resolved-dns-answer.c b/src/resolve/resolved-dns-answer.c
index 96ae9bc005..08520d9424 100644
--- a/src/resolve/resolved-dns-answer.c
+++ b/src/resolve/resolved-dns-answer.c
@@ -60,10 +60,7 @@ static int dns_answer_reserve_internal(DnsAnswer *a, size_t n) {
         m = ordered_set_size(a->items);
         assert(m <= UINT16_MAX); /* We can only place 64K RRs in an answer at max */
 
-        if (n > UINT16_MAX - m)
-                n = UINT16_MAX;
-        else
-                n += m;
+        n = saturate_add(m, n, UINT16_MAX);
 
         /* Higher multipliers give slightly higher efficiency through hash collisions, but the gains
          * quickly drop off after 2. */
@@ -192,16 +189,11 @@ int dns_answer_add(
 
                 /* Entry already exists, keep the entry with the higher TTL. */
                 if (rr->ttl > exist->rr->ttl) {
-                        dns_resource_record_ref(rr);
-                        dns_resource_record_unref(exist->rr);
-                        exist->rr = rr;
+                        DNS_RR_REPLACE(exist->rr, dns_resource_record_ref(rr));
 
                         /* Update RRSIG and RR at the same time */
-                        if (rrsig) {
-                                dns_resource_record_ref(rrsig);
-                                dns_resource_record_unref(exist->rrsig);
-                                exist->rrsig = rrsig;
-                        }
+                        if (rrsig)
+                                DNS_RR_REPLACE(exist->rrsig, dns_resource_record_ref(rrsig));
                 }
 
                 exist->flags |= flags;
@@ -493,9 +485,7 @@ int dns_answer_extend(DnsAnswer **a, DnsAnswer *b) {
         if (r < 0)
                 return r;
 
-        dns_answer_unref(*a);
-        *a = merged;
-
+        DNS_ANSWER_REPLACE(*a, merged);
         return 0;
 }
 
@@ -578,8 +568,7 @@ int dns_answer_remove_by_answer_keys(DnsAnswer **a, DnsAnswer *b) {
 
                 /* Let's remember this entry's RR key, to optimize the loop a bit: if we have an RRset with
                  * more than one item then we don't need to remove the key multiple times */
-                dns_resource_key_unref(prev);
-                prev = dns_resource_key_ref(item->rr->key);
+                DNS_RESOURCE_KEY_REPLACE(prev, dns_resource_key_ref(item->rr->key));
         }
 
         return 0;
@@ -707,10 +696,7 @@ int dns_answer_reserve_or_clone(DnsAnswer **a, size_t n_free) {
         ns = dns_answer_size(*a);
         assert(ns <= UINT16_MAX); /* Maximum number of RRs we can stick into a DNS packet section */
 
-        if (n_free > UINT16_MAX - ns) /* overflow check */
-                ns = UINT16_MAX;
-        else
-                ns += n_free;
+        ns = saturate_add(ns, n_free, UINT16_MAX);
 
         n = dns_answer_new(ns);
         if (!n)
@@ -720,8 +706,7 @@ int dns_answer_reserve_or_clone(DnsAnswer **a, size_t n_free) {
         if (r < 0)
                 return r;
 
-        dns_answer_unref(*a);
-        *a = TAKE_PTR(n);
+        DNS_ANSWER_REPLACE(*a, TAKE_PTR(n));
         return 0;
 }
 
diff --git a/src/resolve/resolved-dns-answer.h b/src/resolve/resolved-dns-answer.h
index 414c03192a..a249f06273 100644
--- a/src/resolve/resolved-dns-answer.h
+++ b/src/resolve/resolved-dns-answer.h
@@ -45,6 +45,14 @@ DnsAnswer *dns_answer_new(size_t n);
 DnsAnswer *dns_answer_ref(DnsAnswer *a);
 DnsAnswer *dns_answer_unref(DnsAnswer *a);
 
+#define DNS_ANSWER_REPLACE(a, b)                \
+        do {                                    \
+                typeof(a)* _a = &(a);           \
+                typeof(b) _b = (b);             \
+                dns_answer_unref(*_a);          \
+                *_a = _b;                       \
+        } while(0)
+
 int dns_answer_add(DnsAnswer *a, DnsResourceRecord *rr, int ifindex, DnsAnswerFlags flags, DnsResourceRecord *rrsig);
 int dns_answer_add_extend(DnsAnswer **a, DnsResourceRecord *rr, int ifindex, DnsAnswerFlags flags, DnsResourceRecord *rrsig);
 int dns_answer_add_soa(DnsAnswer *a, const char *name, uint32_t ttl, int ifindex);
diff --git a/src/resolve/resolved-dns-cache.c b/src/resolve/resolved-dns-cache.c
index 45bcbfd9d7..e7ab4e5569 100644
--- a/src/resolve/resolved-dns-cache.c
+++ b/src/resolve/resolved-dns-cache.c
@@ -378,20 +378,13 @@ static void dns_cache_item_update_positive(
 
                 assert_se(hashmap_replace(c->by_key, rr->key, i) >= 0);
 
-        dns_resource_record_ref(rr);
-        dns_resource_record_unref(i->rr);
-        i->rr = rr;
+        DNS_RR_REPLACE(i->rr, dns_resource_record_ref(rr));
 
-        dns_resource_key_unref(i->key);
-        i->key = dns_resource_key_ref(rr->key);
+        DNS_RESOURCE_KEY_REPLACE(i->key, dns_resource_key_ref(rr->key));
 
-        dns_answer_ref(answer);
-        dns_answer_unref(i->answer);
-        i->answer = answer;
+        DNS_ANSWER_REPLACE(i->answer, dns_answer_ref(answer));
 
-        dns_packet_ref(full_packet);
-        dns_packet_unref(i->full_packet);
-        i->full_packet = full_packet;
+        DNS_PACKET_REPLACE(i->full_packet, dns_packet_ref(full_packet));
 
         i->until = calculate_until(rr, min_ttl, UINT32_MAX, timestamp, false);
         i->query_flags = query_flags & CACHEABLE_QUERY_FLAGS;
diff --git a/src/resolve/resolved-dns-packet.c b/src/resolve/resolved-dns-packet.c
index 50785a6823..d6fb4880b0 100644
--- a/src/resolve/resolved-dns-packet.c
+++ b/src/resolve/resolved-dns-packet.c
@@ -2359,8 +2359,7 @@ static int dns_packet_extract_answer(DnsPacket *p, DnsAnswer **ret_answer) {
                 /* Remember this RR, so that we can potentially merge its ->key object with the
                  * next RR. Note that we only do this if we actually decided to keep the RR around.
                  */
-                dns_resource_record_unref(previous);
-                previous = dns_resource_record_ref(rr);
+                DNS_RR_REPLACE(previous, dns_resource_record_ref(rr));
         }
 
         if (bad_opt) {
diff --git a/src/resolve/resolved-dns-packet.h b/src/resolve/resolved-dns-packet.h
index 0b797ecb1a..95b0b506ea 100644
--- a/src/resolve/resolved-dns-packet.h
+++ b/src/resolve/resolved-dns-packet.h
@@ -201,6 +201,14 @@ DnsPacket *dns_packet_unref(DnsPacket *p);
 
 DEFINE_TRIVIAL_CLEANUP_FUNC(DnsPacket*, dns_packet_unref);
 
+#define DNS_PACKET_REPLACE(a, b)                \
+        do {                                    \
+                typeof(a)* _a = &(a);           \
+                typeof(b) _b = (b);             \
+                dns_packet_unref(*_a);          \
+                *_a = _b;                       \
+        } while(0)
+
 int dns_packet_validate(DnsPacket *p);
 int dns_packet_validate_reply(DnsPacket *p);
 int dns_packet_validate_query(DnsPacket *p);
diff --git a/src/resolve/resolved-dns-query.c b/src/resolve/resolved-dns-query.c
index b62c5c349f..175fbe2296 100644
--- a/src/resolve/resolved-dns-query.c
+++ b/src/resolve/resolved-dns-query.c
@@ -850,17 +850,14 @@ static void dns_query_accept(DnsQuery *q, DnsQueryCandidate *c) {
                                 q->answer_query_flags |= dns_transaction_source_to_query_flags(t->answer_source);
                         } else {
                                 /* Override non-successful previous answers */
-                                dns_answer_unref(q->answer);
-                                q->answer = dns_answer_ref(t->answer);
-
+                                DNS_ANSWER_REPLACE(q->answer, dns_answer_ref(t->answer));
                                 q->answer_query_flags = dns_transaction_source_to_query_flags(t->answer_source);
                         }
 
                         q->answer_rcode = t->answer_rcode;
                         q->answer_errno = 0;
 
-                        dns_packet_unref(q->answer_full_packet);
-                        q->answer_full_packet = dns_packet_ref(t->received);
+                        DNS_PACKET_REPLACE(q->answer_full_packet, dns_packet_ref(t->received));
 
                         if (FLAGS_SET(t->answer_query_flags, SD_RESOLVED_AUTHENTICATED)) {
                                 has_authenticated = true;
@@ -896,14 +893,12 @@ static void dns_query_accept(DnsQuery *q, DnsQueryCandidate *c) {
                             !FLAGS_SET(t->answer_query_flags, SD_RESOLVED_AUTHENTICATED))
                                 continue;
 
-                        dns_answer_unref(q->answer);
-                        q->answer = dns_answer_ref(t->answer);
+                        DNS_ANSWER_REPLACE(q->answer, dns_answer_ref(t->answer));
                         q->answer_rcode = t->answer_rcode;
                         q->answer_dnssec_result = t->answer_dnssec_result;
                         q->answer_query_flags = t->answer_query_flags | dns_transaction_source_to_query_flags(t->answer_source);
                         q->answer_errno = t->answer_errno;
-                        dns_packet_unref(q->answer_full_packet);
-                        q->answer_full_packet = dns_packet_ref(t->received);
+                        DNS_PACKET_REPLACE(q->answer_full_packet, dns_packet_ref(t->received));
 
                         state = t->state;
                         break;
diff --git a/src/resolve/resolved-dns-rr.c b/src/resolve/resolved-dns-rr.c
index ff271486d7..92245770bc 100644
--- a/src/resolve/resolved-dns-rr.c
+++ b/src/resolve/resolved-dns-rr.c
@@ -360,13 +360,10 @@ bool dns_resource_key_reduce(DnsResourceKey **a, DnsResourceKey **b) {
                 return false;
 
         /* Keep the one which already has more references. */
-        if ((*a)->n_ref > (*b)->n_ref) {
-                dns_resource_key_unref(*b);
-                *b = dns_resource_key_ref(*a);
-        } else {
-                dns_resource_key_unref(*a);
-                *a = dns_resource_key_ref(*b);
-        }
+        if ((*a)->n_ref > (*b)->n_ref)
+                DNS_RESOURCE_KEY_REPLACE(*b, dns_resource_key_ref(*a));
+        else
+                DNS_RESOURCE_KEY_REPLACE(*a, dns_resource_key_ref(*b));
 
         return true;
 }
@@ -1708,9 +1705,7 @@ int dns_resource_record_clamp_ttl(DnsResourceRecord **rr, uint32_t max_ttl) {
 
         new_rr->ttl = new_ttl;
 
-        dns_resource_record_unref(*rr);
-        *rr = new_rr;
-
+        DNS_RR_REPLACE(*rr, new_rr);
         return 1;
 }
 
diff --git a/src/resolve/resolved-dns-rr.h b/src/resolve/resolved-dns-rr.h
index 66aa10b82c..91b1276f72 100644
--- a/src/resolve/resolved-dns-rr.h
+++ b/src/resolve/resolved-dns-rr.h
@@ -292,6 +292,15 @@ int dns_resource_key_new_append_suffix(DnsResourceKey **ret, DnsResourceKey *key
 DnsResourceKey* dns_resource_key_new_consume(uint16_t class, uint16_t type, char *name);
 DnsResourceKey* dns_resource_key_ref(DnsResourceKey *key);
 DnsResourceKey* dns_resource_key_unref(DnsResourceKey *key);
+
+#define DNS_RESOURCE_KEY_REPLACE(a, b)          \
+        do {                                    \
+                typeof(a)* _a = &(a);           \
+                typeof(b) _b = (b);             \
+                dns_resource_key_unref(*_a);    \
+                *_a = _b;                       \
+        } while(0)
+
 const char* dns_resource_key_name(const DnsResourceKey *key);
 bool dns_resource_key_is_address(const DnsResourceKey *key);
 bool dns_resource_key_is_dnssd_ptr(const DnsResourceKey *key);
@@ -319,6 +328,15 @@ DnsResourceRecord* dns_resource_record_new(DnsResourceKey *key);
 DnsResourceRecord* dns_resource_record_new_full(uint16_t class, uint16_t type, const char *name);
 DnsResourceRecord* dns_resource_record_ref(DnsResourceRecord *rr);
 DnsResourceRecord* dns_resource_record_unref(DnsResourceRecord *rr);
+
+#define DNS_RR_REPLACE(a, b)                    \
+        do {                                    \
+                typeof(a)* _a = &(a);           \
+                typeof(b) _b = (b);             \
+                dns_resource_record_unref(*_a); \
+                *_a = _b;                       \
+        } while(0)
+
 int dns_resource_record_new_reverse(DnsResourceRecord **ret, int family, const union in_addr_union *address, const char *name);
 int dns_resource_record_new_address(DnsResourceRecord **ret, int family, const union in_addr_union *address, const char *name);
 int dns_resource_record_equal(const DnsResourceRecord *a, const DnsResourceRecord *b);
diff --git a/src/resolve/resolved-dns-transaction.c b/src/resolve/resolved-dns-transaction.c
index fad462b37d..ab0f318c7a 100644
--- a/src/resolve/resolved-dns-transaction.c
+++ b/src/resolve/resolved-dns-transaction.c
@@ -1097,10 +1097,8 @@ void dns_transaction_process_reply(DnsTransaction *t, DnsPacket *p, bool encrypt
                 assert_not_reached();
         }
 
-        if (t->received != p) {
-                dns_packet_unref(t->received);
-                t->received = dns_packet_ref(p);
-        }
+        if (t->received != p)
+                DNS_PACKET_REPLACE(t->received, dns_packet_ref(p));
 
         t->answer_source = DNS_TRANSACTION_NETWORK;
 
@@ -1365,8 +1363,7 @@ void dns_transaction_process_reply(DnsTransaction *t, DnsPacket *p, bool encrypt
          * field is later replaced by the DNSSEC validated subset. The 'answer_auxiliary' field carries the
          * original complete record set, including RRSIG and friends. We use this when passing data to
          * clients that ask for DNSSEC metadata. */
-        dns_answer_unref(t->answer);
-        t->answer = dns_answer_ref(p->answer);
+        DNS_ANSWER_REPLACE(t->answer, dns_answer_ref(p->answer));
         t->answer_rcode = DNS_PACKET_RCODE(p);
         t->answer_dnssec_result = _DNSSEC_RESULT_INVALID;
         SET_FLAG(t->answer_query_flags, SD_RESOLVED_AUTHENTICATED, false);
@@ -3453,8 +3450,7 @@ int dns_transaction_validate_dnssec(DnsTransaction *t) {
                 break;
         }
 
-        dns_answer_unref(t->answer);
-        t->answer = TAKE_PTR(validated);
+        DNS_ANSWER_REPLACE(t->answer, TAKE_PTR(validated));
 
         /* At this point the answer only contains validated
          * RRsets. Now, let's see if it actually answers the question
diff --git a/src/test/test-macro.c b/src/test/test-macro.c
index ba319953cd..c39f64b385 100644
--- a/src/test/test-macro.c
+++ b/src/test/test-macro.c
@@ -6,6 +6,15 @@
 #include "macro.h"
 #include "tests.h"
 
+TEST(saturate_add) {
+        assert_se(saturate_add(1, 2, UINT8_MAX) == 3);
+        assert_se(saturate_add(1, UINT8_MAX-2, UINT8_MAX) == UINT8_MAX-1);
+        assert_se(saturate_add(1, UINT8_MAX-1, UINT8_MAX) == UINT8_MAX);
+        assert_se(saturate_add(1, UINT8_MAX, UINT8_MAX) == UINT8_MAX);
+        assert_se(saturate_add(2, UINT8_MAX, UINT8_MAX) == UINT8_MAX);
+        assert_se(saturate_add(60, 60, 50) == 50);
+}
+
 TEST(align_power2) {
         unsigned long i, p2;