resolved: add dns_packet_set_flags()

We need to support the TC bit in queries in case known answers exceed the
maximum packet size. Factor out the flags compilation to
dns_packet_set_flags() and make it externally available.

2 files changed, 44 insertions, 21 deletions

diff --git a/src/resolve/resolved-dns-packet.c b/src/resolve/resolved-dns-packet.c
index 9bd08eeec2..e5de5a1cd6 100644
--- a/src/resolve/resolved-dns-packet.c
+++ b/src/resolve/resolved-dns-packet.c
@@ -65,40 +65,44 @@ int dns_packet_new(DnsPacket **ret, DnsProtocol protocol, size_t mtu) {
         return 0;
 }
 
-int dns_packet_new_query(DnsPacket **ret, DnsProtocol protocol, size_t mtu, bool dnssec_checking_disabled) {
-        DnsPacket *p;
-        DnsPacketHeader *h;
-        int r;
+void dns_packet_set_flags(DnsPacket *p, bool dnssec_checking_disabled, bool truncated) {
 
-        assert(ret);
+        DnsPacketHeader *h;
 
-        r = dns_packet_new(&p, protocol, mtu);
-        if (r < 0)
-                return r;
+        assert(p);
 
         h = DNS_PACKET_HEADER(p);
 
-        if (protocol == DNS_PROTOCOL_LLMNR)
+        switch(p->protocol) {
+        case DNS_PROTOCOL_LLMNR:
+                assert(!truncated);
+
                 h->flags = htobe16(DNS_PACKET_MAKE_FLAGS(0 /* qr */,
                                                          0 /* opcode */,
                                                          0 /* c */,
-                                                         0 /* tc */,
+                                                         0/* tc */,
                                                          0 /* t */,
                                                          0 /* ra */,
                                                          0 /* ad */,
                                                          0 /* cd */,
                                                          0 /* rcode */));
-        else if (protocol == DNS_PROTOCOL_MDNS)
-                h->flags = htobe16(DNS_PACKET_MAKE_FLAGS(0 /* qr */,
-                                                         0 /* opcode */,
-                                                         0 /* aa */,
-                                                         0 /* tc */,
-                                                         0 /* rd (ask for recursion) */,
-                                                         0 /* ra */,
-                                                         0 /* ad */,
-                                                         0 /* cd */,
-                                                         0 /* rcode */));
-        else
+                break;
+
+        case DNS_PROTOCOL_MDNS:
+                h->flags = htobe16(DNS_PACKET_MAKE_FLAGS(0         /* qr */,
+                                                         0         /* opcode */,
+                                                         0         /* aa */,
+                                                         truncated /* tc */,
+                                                         0         /* rd (ask for recursion) */,
+                                                         0         /* ra */,
+                                                         0         /* ad */,
+                                                         0         /* cd */,
+                                                         0         /* rcode */));
+                break;
+
+        default:
+                assert(!truncated);
+
                 h->flags = htobe16(DNS_PACKET_MAKE_FLAGS(0 /* qr */,
                                                          0 /* opcode */,
                                                          0 /* aa */,
@@ -108,6 +112,23 @@ int dns_packet_new_query(DnsPacket **ret, DnsProtocol protocol, size_t mtu, bool
                                                          0 /* ad */,
                                                          dnssec_checking_disabled /* cd */,
                                                          0 /* rcode */));
+        }
+}
+
+int dns_packet_new_query(DnsPacket **ret, DnsProtocol protocol, size_t mtu, bool dnssec_checking_disabled) {
+        DnsPacket *p;
+        int r;
+
+        assert(ret);
+
+        r = dns_packet_new(&p, protocol, mtu);
+        if (r < 0)
+                return r;
+
+        /* Always set the TC bit to 0 initially.
+         * If there are multiple packets later, we'll update the bit shortly before sending.
+         */
+        dns_packet_set_flags(p, dnssec_checking_disabled, false);
 
         *ret = p;
         return 0;
diff --git a/src/resolve/resolved-dns-packet.h b/src/resolve/resolved-dns-packet.h
index 48b3572cb4..405838e60d 100644
--- a/src/resolve/resolved-dns-packet.h
+++ b/src/resolve/resolved-dns-packet.h
@@ -146,6 +146,8 @@ static inline unsigned DNS_PACKET_RRCOUNT(DnsPacket *p) {
 int dns_packet_new(DnsPacket **p, DnsProtocol protocol, size_t mtu);
 int dns_packet_new_query(DnsPacket **p, DnsProtocol protocol, size_t mtu, bool dnssec_checking_disabled);
 
+void dns_packet_set_flags(DnsPacket *p, bool dnssec_checking_disabled, bool truncated);
+
 DnsPacket *dns_packet_ref(DnsPacket *p);
 DnsPacket *dns_packet_unref(DnsPacket *p);