resolve: make dns_stream_new() take on_packet and complete callbacks

And make on_packet callback mandatory.

6 files changed, 31 insertions, 23 deletions

diff --git a/src/resolve/resolved-dns-stream.c b/src/resolve/resolved-dns-stream.c
index 51ffa6b4b05..bdf46170d18 100644
--- a/src/resolve/resolved-dns-stream.c
+++ b/src/resolve/resolved-dns-stream.c
@@ -411,16 +411,13 @@ static int on_stream_io_impl(DnsStream *s, uint32_t revents) {
                                         s->n_read += ss;
                         }
 
-                        /* Are we done? If so, disable the event source for EPOLLIN */
+                        /* Are we done? If so, call the packet handler and re-enable EPOLLIN for the
+                         * event source if necessary. */
                         if (s->n_read >= sizeof(s->read_size) + be16toh(s->read_size)) {
-                                /* If there's a packet handler
-                                 * installed, call that. Note that
-                                 * this is optional... */
-                                if (s->on_packet) {
-                                        r = s->on_packet(s);
-                                        if (r < 0)
-                                                return r;
-                                }
+                                assert(s->on_packet);
+                                r = s->on_packet(s);
+                                if (r < 0)
+                                        return r;
 
                                 r = dns_stream_update_io(s);
                                 if (r < 0)
@@ -523,6 +520,8 @@ int dns_stream_new(
                 DnsProtocol protocol,
                 int fd,
                 const union sockaddr_union *tfo_address,
+                int (on_packet)(DnsStream*),
+                int (complete)(DnsStream*, int), /* optional */
                 usec_t connect_timeout_usec) {
 
         _cleanup_(dns_stream_unrefp) DnsStream *s = NULL;
@@ -535,6 +534,7 @@ int dns_stream_new(
         assert(protocol >= 0);
         assert(protocol < _DNS_PROTOCOL_MAX);
         assert(fd >= 0);
+        assert(on_packet);
 
         if (m->n_dns_streams[type] > DNS_STREAMS_MAX)
                 return -EBUSY;
@@ -576,6 +576,8 @@ int dns_stream_new(
         s->manager = m;
 
         s->fd = fd;
+        s->on_packet = on_packet;
+        s->complete = complete;
 
         if (tfo_address) {
                 s->tfo_address = *tfo_address;
diff --git a/src/resolve/resolved-dns-stream.h b/src/resolve/resolved-dns-stream.h
index 96b977f6280..548b2edc9ef 100644
--- a/src/resolve/resolved-dns-stream.h
+++ b/src/resolve/resolved-dns-stream.h
@@ -93,7 +93,16 @@ struct DnsStream {
         LIST_FIELDS(DnsStream, streams);
 };
 
-int dns_stream_new(Manager *m, DnsStream **s, DnsStreamType type, DnsProtocol protocol, int fd, const union sockaddr_union *tfo_address, usec_t timeout);
+int dns_stream_new(
+                Manager *m,
+                DnsStream **ret,
+                DnsStreamType type,
+                DnsProtocol protocol,
+                int fd,
+                const union sockaddr_union *tfo_address,
+                int (on_packet)(DnsStream*),
+                int (complete)(DnsStream*, int), /* optional */
+                usec_t connect_timeout_usec);
 #if ENABLE_DNS_OVER_TLS
 int dns_stream_connect_tls(DnsStream *s, void *tls_session);
 #endif
diff --git a/src/resolve/resolved-dns-stub.c b/src/resolve/resolved-dns-stub.c
index 73590e3f9bd..7eb93f11747 100644
--- a/src/resolve/resolved-dns-stub.c
+++ b/src/resolve/resolved-dns-stub.c
@@ -1074,15 +1074,14 @@ static int on_dns_stub_stream_internal(sd_event_source *s, int fd, uint32_t reve
                 return -errno;
         }
 
-        r = dns_stream_new(m, &stream, DNS_STREAM_STUB, DNS_PROTOCOL_DNS, cfd, NULL, DNS_STREAM_STUB_TIMEOUT_USEC);
+        r = dns_stream_new(m, &stream, DNS_STREAM_STUB, DNS_PROTOCOL_DNS, cfd, NULL,
+                           on_dns_stub_stream_packet, dns_stub_stream_complete, DNS_STREAM_STUB_TIMEOUT_USEC);
         if (r < 0) {
                 safe_close(cfd);
                 return r;
         }
 
         stream->stub_listener_extra = l;
-        stream->on_packet = on_dns_stub_stream_packet;
-        stream->complete = dns_stub_stream_complete;
 
         /* We let the reference to the stream dangle here, it will be dropped later by the complete callback. */
 
diff --git a/src/resolve/resolved-dns-transaction.c b/src/resolve/resolved-dns-transaction.c
index 0cf99127124..20d257bbf3b 100644
--- a/src/resolve/resolved-dns-transaction.c
+++ b/src/resolve/resolved-dns-transaction.c
@@ -754,7 +754,8 @@ static int dns_transaction_emit_tcp(DnsTransaction *t) {
                 if (fd < 0)
                         return fd;
 
-                r = dns_stream_new(t->scope->manager, &s, type, t->scope->protocol, fd, &sa, stream_timeout_usec);
+                r = dns_stream_new(t->scope->manager, &s, type, t->scope->protocol, fd, &sa,
+                                   on_stream_packet, on_stream_complete, stream_timeout_usec);
                 if (r < 0)
                         return r;
 
@@ -777,9 +778,6 @@ static int dns_transaction_emit_tcp(DnsTransaction *t) {
                         t->server->stream = dns_stream_ref(s);
                 }
 
-                s->complete = on_stream_complete;
-                s->on_packet = on_stream_packet;
-
                 /* The interface index is difficult to determine if we are
                  * connecting to the local host, hence fill this in right away
                  * instead of determining it from the socket */
diff --git a/src/resolve/resolved-llmnr.c b/src/resolve/resolved-llmnr.c
index 32483006b1e..150cbab1863 100644
--- a/src/resolve/resolved-llmnr.c
+++ b/src/resolve/resolved-llmnr.c
@@ -313,15 +313,15 @@ static int on_llmnr_stream(sd_event_source *s, int fd, uint32_t revents, void *u
                 return -errno;
         }
 
-        r = dns_stream_new(m, &stream, DNS_STREAM_LLMNR_RECV, DNS_PROTOCOL_LLMNR, cfd, NULL, DNS_STREAM_DEFAULT_TIMEOUT_USEC);
+        /* We don't configure a "complete" handler here, we rely on the default handler than simply drops the
+         * reference to the stream, thus freeing it */
+        r = dns_stream_new(m, &stream, DNS_STREAM_LLMNR_RECV, DNS_PROTOCOL_LLMNR, cfd, NULL,
+                           on_llmnr_stream_packet, NULL, DNS_STREAM_DEFAULT_TIMEOUT_USEC);
         if (r < 0) {
                 safe_close(cfd);
                 return r;
         }
 
-        stream->on_packet = on_llmnr_stream_packet;
-        /* We don't configure a "complete" handler here, we rely on the default handler than simply drops the
-         * reference to the stream, thus freeing it */
         return 0;
 }
 
diff --git a/src/resolve/test-resolved-stream.c b/src/resolve/test-resolved-stream.c
index fd7ade19d1e..76467629fbd 100644
--- a/src/resolve/test-resolved-stream.c
+++ b/src/resolve/test-resolved-stream.c
@@ -253,8 +253,8 @@ static void test_dns_stream(bool tls) {
 
         /* Initialize DNS stream */
         assert_se(dns_stream_new(&manager, &stream, DNS_STREAM_LOOKUP, DNS_PROTOCOL_DNS,
-                                 TAKE_FD(clientfd), NULL, DNS_STREAM_DEFAULT_TIMEOUT_USEC) >= 0);
-        stream->on_packet = on_stream_packet;
+                                 TAKE_FD(clientfd), NULL, on_stream_packet, NULL,
+                                 DNS_STREAM_DEFAULT_TIMEOUT_USEC) >= 0);
 #if ENABLE_DNS_OVER_TLS
         if (tls) {
                 DnsServer server = {