From 18230451c03a6d20141efbc85341b6a5c6809077 Mon Sep 17 00:00:00 2001
From: Yu Watanabe <watanabe.yu+github@gmail.com>
Date: Fri, 28 Jan 2022 08:30:36 +0900
Subject: [PATCH 1/4] resolve: make dns_stream_new() take on_packet and
 complete callbacks

And make on_packet callback mandatory.
---
 src/resolve/resolved-dns-stream.c      | 20 +++++++++++---------
 src/resolve/resolved-dns-stream.h      | 11 ++++++++++-
 src/resolve/resolved-dns-stub.c        |  5 ++---
 src/resolve/resolved-dns-transaction.c |  6 ++----
 src/resolve/resolved-llmnr.c           |  8 ++++----
 src/resolve/test-resolved-stream.c     |  4 ++--
 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 = {

From 624f907ea9a42930bffb343dd44fbb0e34746cb0 Mon Sep 17 00:00:00 2001
From: Yu Watanabe <watanabe.yu+github@gmail.com>
Date: Fri, 28 Jan 2022 08:57:05 +0900
Subject: [PATCH 2/4] resolve: call dns_stream_take_read_packet() in
 on_stream_io()

As dns_stream_take_read_packet() is called only in on_packet callbacks,
and all on_packet callbacks call it.
---
 src/resolve/resolved-dns-stream.c      | 39 +++++++++++++-------------
 src/resolve/resolved-dns-stream.h      |  6 ++--
 src/resolve/resolved-dns-stub.c        |  7 ++---
 src/resolve/resolved-dns-transaction.c |  8 ++----
 src/resolve/resolved-llmnr.c           |  6 ++--
 src/resolve/test-resolved-stream.c     |  4 +--
 6 files changed, 31 insertions(+), 39 deletions(-)

diff --git a/src/resolve/resolved-dns-stream.c b/src/resolve/resolved-dns-stream.c
index bdf46170d18..1b2db512129 100644
--- a/src/resolve/resolved-dns-stream.c
+++ b/src/resolve/resolved-dns-stream.c
@@ -281,6 +281,22 @@ static int on_stream_timeout(sd_event_source *es, usec_t usec, void *userdata) {
         return dns_stream_complete(s, ETIMEDOUT);
 }
 
+static DnsPacket *dns_stream_take_read_packet(DnsStream *s) {
+        assert(s);
+
+        if (!s->read_packet)
+                return NULL;
+
+        if (s->n_read < sizeof(s->read_size))
+                return NULL;
+
+        if (s->n_read < sizeof(s->read_size) + be16toh(s->read_size))
+                return NULL;
+
+        s->n_read = 0;
+        return TAKE_PTR(s->read_packet);
+}
+
 static int on_stream_io_impl(DnsStream *s, uint32_t revents) {
         bool progressed = false;
         int r;
@@ -413,9 +429,10 @@ static int on_stream_io_impl(DnsStream *s, uint32_t revents) {
 
                         /* 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)) {
+                        _cleanup_(dns_packet_unrefp) DnsPacket *p = dns_stream_take_read_packet(s);
+                        if (p) {
                                 assert(s->on_packet);
-                                r = s->on_packet(s);
+                                r = s->on_packet(s, p);
                                 if (r < 0)
                                         return r;
 
@@ -520,7 +537,7 @@ int dns_stream_new(
                 DnsProtocol protocol,
                 int fd,
                 const union sockaddr_union *tfo_address,
-                int (on_packet)(DnsStream*),
+                int (on_packet)(DnsStream*, DnsPacket*),
                 int (complete)(DnsStream*, int), /* optional */
                 usec_t connect_timeout_usec) {
 
@@ -604,22 +621,6 @@ int dns_stream_write_packet(DnsStream *s, DnsPacket *p) {
         return dns_stream_update_io(s);
 }
 
-DnsPacket *dns_stream_take_read_packet(DnsStream *s) {
-        assert(s);
-
-        if (!s->read_packet)
-                return NULL;
-
-        if (s->n_read < sizeof(s->read_size))
-                return NULL;
-
-        if (s->n_read < sizeof(s->read_size) + be16toh(s->read_size))
-                return NULL;
-
-        s->n_read = 0;
-        return TAKE_PTR(s->read_packet);
-}
-
 void dns_stream_detach(DnsStream *s) {
         assert(s);
 
diff --git a/src/resolve/resolved-dns-stream.h b/src/resolve/resolved-dns-stream.h
index 548b2edc9ef..fedbab2da2c 100644
--- a/src/resolve/resolved-dns-stream.h
+++ b/src/resolve/resolved-dns-stream.h
@@ -78,7 +78,7 @@ struct DnsStream {
         size_t n_written, n_read;
         OrderedSet *write_queue;
 
-        int (*on_packet)(DnsStream *s);
+        int (*on_packet)(DnsStream *s, DnsPacket *p);
         int (*complete)(DnsStream *s, int error);
 
         LIST_HEAD(DnsTransaction, transactions); /* when used by the transaction logic */
@@ -100,7 +100,7 @@ int dns_stream_new(
                 DnsProtocol protocol,
                 int fd,
                 const union sockaddr_union *tfo_address,
-                int (on_packet)(DnsStream*),
+                int (on_packet)(DnsStream*, DnsPacket*),
                 int (complete)(DnsStream*, int), /* optional */
                 usec_t connect_timeout_usec);
 #if ENABLE_DNS_OVER_TLS
@@ -123,6 +123,4 @@ static inline bool DNS_STREAM_QUEUED(DnsStream *s) {
         return !!s->write_packet;
 }
 
-DnsPacket *dns_stream_take_read_packet(DnsStream *s);
-
 void dns_stream_detach(DnsStream *s);
diff --git a/src/resolve/resolved-dns-stub.c b/src/resolve/resolved-dns-stub.c
index 7eb93f11747..992ae19bbc7 100644
--- a/src/resolve/resolved-dns-stub.c
+++ b/src/resolve/resolved-dns-stub.c
@@ -1044,12 +1044,9 @@ static int on_dns_stub_packet_extra(sd_event_source *s, int fd, uint32_t revents
         return on_dns_stub_packet_internal(s, fd, revents, l->manager, l);
 }
 
-static int on_dns_stub_stream_packet(DnsStream *s) {
-        _cleanup_(dns_packet_unrefp) DnsPacket *p = NULL;
-
+static int on_dns_stub_stream_packet(DnsStream *s, DnsPacket *p) {
         assert(s);
-
-        p = dns_stream_take_read_packet(s);
+        assert(s->manager);
         assert(p);
 
         if (dns_packet_validate_query(p) > 0) {
diff --git a/src/resolve/resolved-dns-transaction.c b/src/resolve/resolved-dns-transaction.c
index 20d257bbf3b..f937f9f7b59 100644
--- a/src/resolve/resolved-dns-transaction.c
+++ b/src/resolve/resolved-dns-transaction.c
@@ -644,14 +644,12 @@ static int on_stream_complete(DnsStream *s, int error) {
         return 0;
 }
 
-static int on_stream_packet(DnsStream *s) {
-        _cleanup_(dns_packet_unrefp) DnsPacket *p = NULL;
+static int on_stream_packet(DnsStream *s, DnsPacket *p) {
         DnsTransaction *t;
 
         assert(s);
-
-        /* Take ownership of packet to be able to receive new packets */
-        assert_se(p = dns_stream_take_read_packet(s));
+        assert(s->manager);
+        assert(p);
 
         t = hashmap_get(s->manager->dns_transactions, UINT_TO_PTR(DNS_PACKET_ID(p)));
         if (t && t->stream == s) /* Validate that the stream we got this on actually is the stream the
diff --git a/src/resolve/resolved-llmnr.c b/src/resolve/resolved-llmnr.c
index 150cbab1863..b4e551c219d 100644
--- a/src/resolve/resolved-llmnr.c
+++ b/src/resolve/resolved-llmnr.c
@@ -277,13 +277,11 @@ int manager_llmnr_ipv6_udp_fd(Manager *m) {
         return m->llmnr_ipv6_udp_fd = TAKE_FD(s);
 }
 
-static int on_llmnr_stream_packet(DnsStream *s) {
-        _cleanup_(dns_packet_unrefp) DnsPacket *p = NULL;
+static int on_llmnr_stream_packet(DnsStream *s, DnsPacket *p) {
         DnsScope *scope;
 
         assert(s);
-
-        p = dns_stream_take_read_packet(s);
+        assert(s->manager);
         assert(p);
 
         scope = manager_find_scope(s->manager, p);
diff --git a/src/resolve/test-resolved-stream.c b/src/resolve/test-resolved-stream.c
index 76467629fbd..8a01460a0ee 100644
--- a/src/resolve/test-resolved-stream.c
+++ b/src/resolve/test-resolved-stream.c
@@ -194,9 +194,9 @@ static const size_t MAX_RECEIVED_PACKETS = 2;
 static DnsPacket *received_packets[2] = {};
 static size_t n_received_packets = 0;
 
-static int on_stream_packet(DnsStream *stream) {
+static int on_stream_packet(DnsStream *stream, DnsPacket *p) {
         assert_se(n_received_packets < MAX_RECEIVED_PACKETS);
-        assert_se(received_packets[n_received_packets++] = dns_stream_take_read_packet(stream));
+        assert_se(received_packets[n_received_packets++] = dns_packet_ref(p));
         return 0;
 }
 

From 4aa6129897d2e8de9b275b44270c1c9da745de0e Mon Sep 17 00:00:00 2001
From: Yu Watanabe <watanabe.yu+github@gmail.com>
Date: Fri, 28 Jan 2022 09:01:07 +0900
Subject: [PATCH 3/4] resolve: mention that dns_stream_update() needs to be
 called after dns_stream_take_read_packet()

Based on the analysis by Joan Bruguera <joanbrugueram@gmail.com>.
See https://github.com/systemd/systemd/pull/22132#discussion_r793951650.
---
 src/resolve/resolved-dns-stream.c | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/src/resolve/resolved-dns-stream.c b/src/resolve/resolved-dns-stream.c
index 1b2db512129..d16ea95d434 100644
--- a/src/resolve/resolved-dns-stream.c
+++ b/src/resolve/resolved-dns-stream.c
@@ -284,6 +284,13 @@ static int on_stream_timeout(sd_event_source *es, usec_t usec, void *userdata) {
 static DnsPacket *dns_stream_take_read_packet(DnsStream *s) {
         assert(s);
 
+        /* Note, dns_stream_update() should be called after this is called. When this is called, the
+         * stream may be already full and the EPOLLIN flag is dropped from the stream IO event source.
+         * Even this makes a room to read in the stream, this does not call dns_stream_update(), hence
+         * EPOLLIN flag is not set automatically. So, to read further packets from the stream,
+         * dns_stream_update() must be called explicitly. Currently, this is only called from
+         * on_stream_io_impl(), and there dns_stream_update() is called. */
+
         if (!s->read_packet)
                 return NULL;
 

From a5e2a488e83fabf6d8ade7621c2fc3574a8faaa7 Mon Sep 17 00:00:00 2001
From: Yu Watanabe <watanabe.yu+github@gmail.com>
Date: Fri, 28 Jan 2022 09:29:59 +0900
Subject: [PATCH 4/4] resolve: llmnr: fix never hit condition

Previously, the condition in on_stream_io_impl() never hit, as the
read packet is always taken from the stream in the few lines above.

Instead of the dns_stream_complete() under the condition, the stream
is unref()ed in the on_packet callback for LLMNR stream, unlike the
other on_packet callbacks.

That's quite tricky. Also, potentially, the stream may still have
queued packets to write.

This fix the condition, and drops the unref() in the on_packet callback.

C.f. https://github.com/systemd/systemd/pull/22274#issuecomment-1023708449.

Closes #22266.
---
 src/resolve/resolved-dns-stream.c | 22 +++++++++++++++-------
 src/resolve/resolved-dns-stream.h |  1 +
 src/resolve/resolved-llmnr.c      |  4 +---
 3 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/src/resolve/resolved-dns-stream.c b/src/resolve/resolved-dns-stream.c
index d16ea95d434..cf9d1a9d5e5 100644
--- a/src/resolve/resolved-dns-stream.c
+++ b/src/resolve/resolved-dns-stream.c
@@ -446,17 +446,25 @@ static int on_stream_io_impl(DnsStream *s, uint32_t revents) {
                                 r = dns_stream_update_io(s);
                                 if (r < 0)
                                         return dns_stream_complete(s, -r);
+
+                                s->packet_received = true;
                         }
                 }
         }
 
-        /* Call "complete" callback if finished reading and writing one packet, and there's nothing else left
-         * to write. */
-        if (s->type == DNS_STREAM_LLMNR_SEND &&
-            (s->write_packet && s->n_written >= sizeof(s->write_size) + s->write_packet->size) &&
-            ordered_set_isempty(s->write_queue) &&
-            (s->read_packet && s->n_read >= sizeof(s->read_size) + s->read_packet->size))
-                return dns_stream_complete(s, 0);
+        if (s->type == DNS_STREAM_LLMNR_SEND && s->packet_received) {
+                uint32_t events;
+
+                /* Complete the stream if finished reading and writing one packet, and there's nothing
+                 * else left to write. */
+
+                r = sd_event_source_get_io_events(s->io_event_source, &events);
+                if (r < 0)
+                        return r;
+
+                if (!FLAGS_SET(events, EPOLLOUT))
+                        return dns_stream_complete(s, 0);
+        }
 
         /* If we did something, let's restart the timeout event source */
         if (progressed && s->timeout_event_source) {
diff --git a/src/resolve/resolved-dns-stream.h b/src/resolve/resolved-dns-stream.h
index fedbab2da2c..1c606365cdc 100644
--- a/src/resolve/resolved-dns-stream.h
+++ b/src/resolve/resolved-dns-stream.h
@@ -60,6 +60,7 @@ struct DnsStream {
         int ifindex;
         uint32_t ttl;
         bool identified;
+        bool packet_received; /* At least one packet is received. Used by LLMNR. */
 
         /* only when using TCP fast open */
         union sockaddr_union tfo_address;
diff --git a/src/resolve/resolved-llmnr.c b/src/resolve/resolved-llmnr.c
index b4e551c219d..76e42940f45 100644
--- a/src/resolve/resolved-llmnr.c
+++ b/src/resolve/resolved-llmnr.c
@@ -294,7 +294,6 @@ static int on_llmnr_stream_packet(DnsStream *s, DnsPacket *p) {
         } else
                 log_debug("Invalid LLMNR TCP packet, ignoring.");
 
-        dns_stream_unref(s);
         return 0;
 }
 
@@ -311,8 +310,7 @@ static int on_llmnr_stream(sd_event_source *s, int fd, uint32_t revents, void *u
                 return -errno;
         }
 
-        /* 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 */
+        /* We don't configure a "complete" handler here, we rely on the default handler, 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) {