From 57bdb749b83e26f7b972d510c2fa09d23bc3617c Mon Sep 17 00:00:00 2001
From: Lennart Poettering <lennart@poettering.net>
Date: Mon, 21 Jan 2019 17:56:34 +0100
Subject: [PATCH 1/8] stream: follow coding style, don't use degrade-to-bool
 for checking numeric value

---
 src/resolve/resolved-dns-stream.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/resolve/resolved-dns-stream.c b/src/resolve/resolved-dns-stream.c
index aee339a4c8..e29c970d51 100644
--- a/src/resolve/resolved-dns-stream.c
+++ b/src/resolve/resolved-dns-stream.c
@@ -41,7 +41,7 @@ static int dns_stream_update_io(DnsStream *s) {
 
 #if ENABLE_DNS_OVER_TLS
         /* For handshake and clean closing purposes, TLS can override requested events */
-        if (s->dnstls_events)
+        if (s->dnstls_events != 0)
                 f = s->dnstls_events;
 #endif
 

From 652ba568c6624bf40d735645f029d83d21bdeaa6 Mon Sep 17 00:00:00 2001
From: Lennart Poettering <lennart@poettering.net>
Date: Mon, 21 Jan 2019 17:57:43 +0100
Subject: [PATCH 2/8] stream: track type of DnsStream object

We use stream objects in four different cases: let's track them.

This in particular allows us to make sure the limit on outgoing streams
cannot be exhausted by having incoming streams as this means we can
neatly separate the counters for all four types.
---
 src/resolve/resolved-dns-stream.c      | 11 ++++++++---
 src/resolve/resolved-dns-stream.h      | 12 +++++++++++-
 src/resolve/resolved-dns-stub.c        |  2 +-
 src/resolve/resolved-dns-transaction.c |  7 +++++--
 src/resolve/resolved-llmnr.c           |  2 +-
 src/resolve/resolved-manager.h         |  2 +-
 6 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/src/resolve/resolved-dns-stream.c b/src/resolve/resolved-dns-stream.c
index e29c970d51..ebafaa529c 100644
--- a/src/resolve/resolved-dns-stream.c
+++ b/src/resolve/resolved-dns-stream.c
@@ -437,7 +437,7 @@ static DnsStream *dns_stream_free(DnsStream *s) {
 
         if (s->manager) {
                 LIST_REMOVE(streams, s->manager->dns_streams, s);
-                s->manager->n_dns_streams--;
+                s->manager->n_dns_streams[s->type]--;
         }
 
 #if ENABLE_DNS_OVER_TLS
@@ -462,6 +462,7 @@ DEFINE_TRIVIAL_REF_UNREF_FUNC(DnsStream, dns_stream, dns_stream_free);
 int dns_stream_new(
                 Manager *m,
                 DnsStream **ret,
+                DnsStreamType type,
                 DnsProtocol protocol,
                 int fd,
                 const union sockaddr_union *tfo_address) {
@@ -471,9 +472,13 @@ int dns_stream_new(
 
         assert(m);
         assert(ret);
+        assert(type >= 0);
+        assert(type < _DNS_STREAM_TYPE_MAX);
+        assert(protocol >= 0);
+        assert(protocol < _DNS_PROTOCOL_MAX);
         assert(fd >= 0);
 
-        if (m->n_dns_streams > DNS_STREAMS_MAX)
+        if (m->n_dns_streams[type] > DNS_STREAMS_MAX)
                 return -EBUSY;
 
         s = new(DnsStream, 1);
@@ -508,7 +513,7 @@ int dns_stream_new(
         (void) sd_event_source_set_description(s->timeout_event_source, "dns-stream-timeout");
 
         LIST_PREPEND(streams, m->dns_streams, s);
-        m->n_dns_streams++;
+        m->n_dns_streams[type]++;
         s->manager = m;
 
         s->fd = fd;
diff --git a/src/resolve/resolved-dns-stream.h b/src/resolve/resolved-dns-stream.h
index f18fc919e7..2c6d9c00b9 100644
--- a/src/resolve/resolved-dns-stream.h
+++ b/src/resolve/resolved-dns-stream.h
@@ -5,6 +5,15 @@
 
 typedef struct DnsStream DnsStream;
 
+typedef enum DnsStreamType {
+        DNS_STREAM_LOOKUP,        /* Outgoing connection to a classic DNS server */
+        DNS_STREAM_LLMNR_SEND,    /* Outgoing LLMNR TCP lookup */
+        DNS_STREAM_LLMNR_RECV,    /* Incoming LLMNR TCP lookup */
+        DNS_STREAM_STUB,          /* Incoming DNS stub connection */
+        _DNS_STREAM_TYPE_MAX,
+        _DNS_STREAM_TYPE_INVALID = -1,
+} DnsStreamType;
+
 #include "resolved-dns-packet.h"
 #include "resolved-dns-transaction.h"
 #include "resolved-manager.h"
@@ -25,6 +34,7 @@ struct DnsStream {
         Manager *manager;
         unsigned n_ref;
 
+        DnsStreamType type;
         DnsProtocol protocol;
 
         int fd;
@@ -66,7 +76,7 @@ struct DnsStream {
         LIST_FIELDS(DnsStream, streams);
 };
 
-int dns_stream_new(Manager *m, DnsStream **s, DnsProtocol protocol, int fd, const union sockaddr_union *tfo_address);
+int dns_stream_new(Manager *m, DnsStream **s, DnsStreamType type, DnsProtocol protocol, int fd, const union sockaddr_union *tfo_address);
 #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 a00716cd85..39ce42d155 100644
--- a/src/resolve/resolved-dns-stub.c
+++ b/src/resolve/resolved-dns-stub.c
@@ -471,7 +471,7 @@ static int on_dns_stub_stream(sd_event_source *s, int fd, uint32_t revents, void
                 return -errno;
         }
 
-        r = dns_stream_new(m, &stream, DNS_PROTOCOL_DNS, cfd, NULL);
+        r = dns_stream_new(m, &stream, DNS_STREAM_STUB, DNS_PROTOCOL_DNS, cfd, NULL);
         if (r < 0) {
                 safe_close(cfd);
                 return r;
diff --git a/src/resolve/resolved-dns-transaction.c b/src/resolve/resolved-dns-transaction.c
index 4a2d2ccbde..075dfac01b 100644
--- a/src/resolve/resolved-dns-transaction.c
+++ b/src/resolve/resolved-dns-transaction.c
@@ -554,9 +554,10 @@ static uint16_t dns_port_for_feature_level(DnsServerFeatureLevel level) {
 }
 
 static int dns_transaction_emit_tcp(DnsTransaction *t) {
-        _cleanup_close_ int fd = -1;
         _cleanup_(dns_stream_unrefp) DnsStream *s = NULL;
+        _cleanup_close_ int fd = -1;
         union sockaddr_union sa;
+        DnsStreamType type;
         int r;
 
         assert(t);
@@ -582,6 +583,7 @@ static int dns_transaction_emit_tcp(DnsTransaction *t) {
                 else
                         fd = dns_scope_socket_tcp(t->scope, AF_UNSPEC, NULL, t->server, dns_port_for_feature_level(t->current_feature_level), &sa);
 
+                type = DNS_STREAM_LOOKUP;
                 break;
 
         case DNS_PROTOCOL_LLMNR:
@@ -607,6 +609,7 @@ static int dns_transaction_emit_tcp(DnsTransaction *t) {
                         fd = dns_scope_socket_tcp(t->scope, family, &address, NULL, LLMNR_PORT, &sa);
                 }
 
+                type = DNS_STREAM_LLMNR_SEND;
                 break;
 
         default:
@@ -617,7 +620,7 @@ static int dns_transaction_emit_tcp(DnsTransaction *t) {
                 if (fd < 0)
                         return fd;
 
-                r = dns_stream_new(t->scope->manager, &s, t->scope->protocol, fd, &sa);
+                r = dns_stream_new(t->scope->manager, &s, type, t->scope->protocol, fd, &sa);
                 if (r < 0)
                         return r;
 
diff --git a/src/resolve/resolved-llmnr.c b/src/resolve/resolved-llmnr.c
index dfa55c577c..d73f86560d 100644
--- a/src/resolve/resolved-llmnr.c
+++ b/src/resolve/resolved-llmnr.c
@@ -295,7 +295,7 @@ 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_PROTOCOL_LLMNR, cfd, NULL);
+        r = dns_stream_new(m, &stream, DNS_STREAM_LLMNR_RECV, DNS_PROTOCOL_LLMNR, cfd, NULL);
         if (r < 0) {
                 safe_close(cfd);
                 return r;
diff --git a/src/resolve/resolved-manager.h b/src/resolve/resolved-manager.h
index 06c76f6014..72171f8c97 100644
--- a/src/resolve/resolved-manager.h
+++ b/src/resolve/resolved-manager.h
@@ -54,7 +54,7 @@ struct Manager {
         unsigned n_dns_queries;
 
         LIST_HEAD(DnsStream, dns_streams);
-        unsigned n_dns_streams;
+        unsigned n_dns_streams[_DNS_STREAM_TYPE_MAX];
 
         /* Unicast dns */
         LIST_HEAD(DnsServer, dns_servers);

From 8227cfa137d8a930519717b752128b8d035b83e5 Mon Sep 17 00:00:00 2001
From: Lennart Poettering <lennart@poettering.net>
Date: Mon, 21 Jan 2019 18:34:00 +0100
Subject: [PATCH 3/8] transaction: simplify handling if we get an unexpected
 DNS packet via TCP

There's no point in calling on_stream_complete() as it doesn't do
anything with the zero argument. Let's hence simplify this and just log.
---
 src/resolve/resolved-dns-transaction.c | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/src/resolve/resolved-dns-transaction.c b/src/resolve/resolved-dns-transaction.c
index 075dfac01b..50c007c307 100644
--- a/src/resolve/resolved-dns-transaction.c
+++ b/src/resolve/resolved-dns-transaction.c
@@ -540,12 +540,8 @@ static int on_stream_packet(DnsStream *s) {
         if (t)
                 return dns_transaction_on_stream_packet(t, p);
 
-        /* Ignore incorrect transaction id as transaction can have been canceled */
-        if (dns_packet_validate_reply(p) <= 0) {
-                log_debug("Invalid TCP reply packet.");
-                on_stream_complete(s, 0);
-        }
-
+        /* Ignore incorrect transaction id as an old transaction can have been canceled. */
+        log_debug("Received unexpected TCP reply packet with id %" PRIu16 ", ignoring.", t->id);
         return 0;
 }
 
@@ -639,8 +635,8 @@ static int dns_transaction_emit_tcp(DnsTransaction *t) {
 
                 if (t->server) {
                         dns_server_unref_stream(t->server);
-                        t->server->stream = dns_stream_ref(s);
                         s->server = dns_server_ref(t->server);
+                        t->server->stream = dns_stream_ref(s);
                 }
 
                 s->complete = on_stream_complete;

From 31f2a5e2b8c72b95c5fd965c6768755d5cba657f Mon Sep 17 00:00:00 2001
From: Lennart Poettering <lennart@poettering.net>
Date: Mon, 21 Jan 2019 18:36:14 +0100
Subject: [PATCH 4/8] llmnr: add comment why we install no complete() handler
 on stream

---
 src/resolve/resolved-llmnr.c | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/resolve/resolved-llmnr.c b/src/resolve/resolved-llmnr.c
index d73f86560d..b7c37f1132 100644
--- a/src/resolve/resolved-llmnr.c
+++ b/src/resolve/resolved-llmnr.c
@@ -302,6 +302,8 @@ static int on_llmnr_stream(sd_event_source *s, int fd, uint32_t revents, void *u
         }
 
         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;
 }
 

From 5971dffd4a420bb10be0d07a3c1b0e1e55554262 Mon Sep 17 00:00:00 2001
From: Lennart Poettering <lennart@poettering.net>
Date: Mon, 21 Jan 2019 19:29:51 +0100
Subject: [PATCH 5/8] resolved: restart stream timeout whenever we managed to
 read or write something

Previously we'd start the timeout once when we allocated the stream.
However, we'd now like to emphasize long-running connections hence let's
rework the timeout logic, and restart it whenever we see action ont the
stream. Thus, idle streams are eventually closed down, but those where
we read or write from are not.
---
 src/resolve/resolved-dns-stream.c | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/src/resolve/resolved-dns-stream.c b/src/resolve/resolved-dns-stream.c
index ebafaa529c..ecc7e9fa09 100644
--- a/src/resolve/resolved-dns-stream.c
+++ b/src/resolve/resolved-dns-stream.c
@@ -281,6 +281,7 @@ static int on_stream_timeout(sd_event_source *es, usec_t usec, void *userdata) {
 
 static int on_stream_io(sd_event_source *es, int fd, uint32_t revents, void *userdata) {
         _cleanup_(dns_stream_unrefp) DnsStream *s = dns_stream_ref(userdata); /* Protect stream while we process it */
+        bool progressed = false;
         int r;
 
         assert(s);
@@ -324,8 +325,10 @@ static int on_stream_io(sd_event_source *es, int fd, uint32_t revents, void *use
                 if (ss < 0) {
                         if (!IN_SET(-ss, EINTR, EAGAIN))
                                 return dns_stream_complete(s, -ss);
-                } else
+                } else {
+                        progressed = true;
                         s->n_written += ss;
+                }
 
                 /* Are we done? If so, disable the event source for EPOLLOUT */
                 if (s->n_written >= sizeof(s->write_size) + s->write_packet->size) {
@@ -348,8 +351,10 @@ static int on_stream_io(sd_event_source *es, int fd, uint32_t revents, void *use
                                         return dns_stream_complete(s, -ss);
                         } else if (ss == 0)
                                 return dns_stream_complete(s, ECONNRESET);
-                        else
+                        else {
+                                progressed = true;
                                 s->n_read += ss;
+                        }
                 }
 
                 if (s->n_read >= sizeof(s->read_size)) {
@@ -424,6 +429,13 @@ static int on_stream_io(sd_event_source *es, int fd, uint32_t revents, void *use
             (s->read_packet && s->n_read >= sizeof(s->read_size) + s->read_packet->size))
                 return dns_stream_complete(s, 0);
 
+        /* If we did something, let's restart the timeout event source */
+        if (progressed && s->timeout_event_source) {
+                r = sd_event_source_set_time(s->timeout_event_source, now(clock_boottime_or_monotonic()) + DNS_STREAM_TIMEOUT_USEC);
+                if (r < 0)
+                        log_warning_errno(errno, "Couldn't restart TCP connection timeout, ignoring: %m");
+        }
+
         return 0;
 }
 

From 9c9e017067a6889458e1df557a8d9e1c599fce33 Mon Sep 17 00:00:00 2001
From: Lennart Poettering <lennart@poettering.net>
Date: Mon, 21 Jan 2019 19:32:32 +0100
Subject: [PATCH 6/8] resolved: only call complete() with zero argument in
 LLMNR client cases

In all other cases (i.e. classic DNS connection towards an upstream
server, or incoming stub connection, or incoming LMMNR connection) we
want long-running connections, hence keep the connection open for good.
Only in the LLMNR client case let's close the stream as soon as we are
done.
---
 src/resolve/resolved-dns-stream.c | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/resolve/resolved-dns-stream.c b/src/resolve/resolved-dns-stream.c
index ecc7e9fa09..45b06eb732 100644
--- a/src/resolve/resolved-dns-stream.c
+++ b/src/resolve/resolved-dns-stream.c
@@ -425,7 +425,11 @@ static int on_stream_io(sd_event_source *es, int fd, uint32_t revents, void *use
                 }
         }
 
-        if ((s->write_packet && s->n_written >= sizeof(s->write_size) + s->write_packet->size) &&
+        /* 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);
 

From f447d9e3765f33a0bae8abca0e7b673b8e429dd8 Mon Sep 17 00:00:00 2001
From: Lennart Poettering <lennart@poettering.net>
Date: Mon, 21 Jan 2019 19:44:01 +0100
Subject: [PATCH 7/8] resolved: add comment to dns_stream_complete() about its
 'error' argument

---
 src/resolve/resolved-dns-stream.c | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/resolve/resolved-dns-stream.c b/src/resolve/resolved-dns-stream.c
index 45b06eb732..3fd056bdb7 100644
--- a/src/resolve/resolved-dns-stream.c
+++ b/src/resolve/resolved-dns-stream.c
@@ -52,6 +52,10 @@ static int dns_stream_complete(DnsStream *s, int error) {
         _cleanup_(dns_stream_unrefp) _unused_ DnsStream *ref = dns_stream_ref(s); /* Protect stream while we process it */
 
         assert(s);
+        assert(error >= 0);
+
+        /* Error is > 0 when the connection failed for some reason in the network stack. It's == 0 if we sent
+         * and receieved exactly one packet each (in the LLMNR client case). */
 
 #if ENABLE_DNS_OVER_TLS
         if (s->encrypted) {

From b412af57a31c9cba56b1fcd0ea4dd626620cb543 Mon Sep 17 00:00:00 2001
From: Lennart Poettering <lennart@poettering.net>
Date: Mon, 21 Jan 2019 19:44:30 +0100
Subject: [PATCH 8/8] resolved: keep stub stream connections up for as long as
 client wants

This enables pipelining of queries from clients to our stub server.

Fixes: #11332
---
 src/resolve/resolved-dns-query.c  |  6 +--
 src/resolve/resolved-dns-stream.c |  8 +++-
 src/resolve/resolved-dns-stream.h |  2 +-
 src/resolve/resolved-dns-stub.c   | 63 ++++++++++++++-----------------
 4 files changed, 38 insertions(+), 41 deletions(-)

diff --git a/src/resolve/resolved-dns-query.c b/src/resolve/resolved-dns-query.c
index 535ef4e776..4a41921cf3 100644
--- a/src/resolve/resolved-dns-query.c
+++ b/src/resolve/resolved-dns-query.c
@@ -387,10 +387,8 @@ DnsQuery *dns_query_free(DnsQuery *q) {
 
         if (q->request_dns_stream) {
                 /* Detach the stream from our query, in case something else keeps a reference to it. */
-                q->request_dns_stream->complete = NULL;
-                q->request_dns_stream->on_packet = NULL;
-                q->request_dns_stream->query = NULL;
-                dns_stream_unref(q->request_dns_stream);
+                (void) set_remove(q->request_dns_stream->queries, q);
+                q->request_dns_stream = dns_stream_unref(q->request_dns_stream);
         }
 
         free(q->request_address_string);
diff --git a/src/resolve/resolved-dns-stream.c b/src/resolve/resolved-dns-stream.c
index 3fd056bdb7..cb7b186fc9 100644
--- a/src/resolve/resolved-dns-stream.c
+++ b/src/resolve/resolved-dns-stream.c
@@ -11,6 +11,8 @@
 #define DNS_STREAM_TIMEOUT_USEC (10 * USEC_PER_SEC)
 #define DNS_STREAMS_MAX 128
 
+#define DNS_QUERIES_PER_STREAM 32
+
 static void dns_stream_stop(DnsStream *s) {
         assert(s);
 
@@ -36,7 +38,11 @@ static int dns_stream_update_io(DnsStream *s) {
                 s->n_written = 0;
                 f |= EPOLLOUT;
         }
-        if (!s->read_packet || s->n_read < sizeof(s->read_size) + s->read_packet->size)
+
+        /* Let's read a packet if we haven't queued any yet. Except if we already hit a limit of parallel
+         * queries for this connection. */
+        if ((!s->read_packet || s->n_read < sizeof(s->read_size) + s->read_packet->size) &&
+                set_size(s->queries) < DNS_QUERIES_PER_STREAM)
                 f |= EPOLLIN;
 
 #if ENABLE_DNS_OVER_TLS
diff --git a/src/resolve/resolved-dns-stream.h b/src/resolve/resolved-dns-stream.h
index 2c6d9c00b9..780051b38a 100644
--- a/src/resolve/resolved-dns-stream.h
+++ b/src/resolve/resolved-dns-stream.h
@@ -68,7 +68,7 @@ struct DnsStream {
 
         LIST_HEAD(DnsTransaction, transactions); /* when used by the transaction logic */
         DnsServer *server;                       /* when used by the transaction logic */
-        DnsQuery *query;                         /* when used by the DNS stub logic */
+        Set *queries;                            /* when used by the DNS stub logic */
 
         /* used when DNS-over-TLS is enabled */
         bool encrypted:1;
diff --git a/src/resolve/resolved-dns-stub.c b/src/resolve/resolved-dns-stub.c
index 39ce42d155..906bdc4bbb 100644
--- a/src/resolve/resolved-dns-stub.c
+++ b/src/resolve/resolved-dns-stub.c
@@ -126,14 +126,6 @@ static int dns_stub_finish_reply_packet(
         return 0;
 }
 
-static void dns_stub_detach_stream(DnsStream *s) {
-        assert(s);
-
-        s->complete = NULL;
-        s->on_packet = NULL;
-        s->query = NULL;
-}
-
 static int dns_stub_send(Manager *m, DnsStream *s, DnsPacket *p, DnsPacket *reply) {
         int r;
 
@@ -257,27 +249,27 @@ static void dns_stub_query_complete(DnsQuery *q) {
                 assert_not_reached("Impossible state");
         }
 
-        /* If there's a packet to write set, let's leave the stream around */
-        if (q->request_dns_stream && DNS_STREAM_QUEUED(q->request_dns_stream)) {
-
-                /* Detach the stream from our query (make it an orphan), but do not drop the reference to it. The
-                 * default completion action of the stream will drop the reference. */
-
-                dns_stub_detach_stream(q->request_dns_stream);
-                q->request_dns_stream = NULL;
-        }
-
         dns_query_free(q);
 }
 
 static int dns_stub_stream_complete(DnsStream *s, int error) {
         assert(s);
 
-        log_debug_errno(error, "DNS TCP connection terminated, destroying query: %m");
+        log_debug_errno(error, "DNS TCP connection terminated, destroying queries: %m");
 
-        assert(s->query);
-        dns_query_free(s->query);
+        for (;;) {
+                DnsQuery *q;
 
+                q = set_first(s->queries);
+                if (!q)
+                        break;
+
+                dns_query_free(q);
+        }
+
+        /* This drops the implicit ref we keep around since it was allocated, as incoming stub connections
+         * should be kept as long as the client wants to. */
+        dns_stream_unref(s);
         return 0;
 }
 
@@ -289,8 +281,6 @@ static void dns_stub_process_query(Manager *m, DnsStream *s, DnsPacket *p) {
         assert(p);
         assert(p->protocol == DNS_PROTOCOL_DNS);
 
-        /* Takes ownership of the *s stream object */
-
         if (in_addr_is_localhost(p->family, &p->sender) <= 0 ||
             in_addr_is_localhost(p->family, &p->destination) <= 0) {
                 log_error("Got packet on unexpected IP range, refusing.");
@@ -351,9 +341,19 @@ static void dns_stub_process_query(Manager *m, DnsStream *s, DnsPacket *p) {
         q->complete = dns_stub_query_complete;
 
         if (s) {
-                s->on_packet = NULL;
-                s->complete = dns_stub_stream_complete;
-                s->query = q;
+                /* 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);
+                if (r < 0) {
+                        log_oom();
+                        goto fail;
+                }
+
+                if (set_put(s->queries, q) < 0) {
+                        log_oom();
+                        goto fail;
+                }
         }
 
         r = dns_query_go(q);
@@ -367,9 +367,6 @@ static void dns_stub_process_query(Manager *m, DnsStream *s, DnsPacket *p) {
         return;
 
 fail:
-        if (s && DNS_STREAM_QUEUED(s))
-                dns_stub_detach_stream(s);
-
         dns_query_free(q);
 }
 
@@ -451,10 +448,6 @@ static int on_dns_stub_stream_packet(DnsStream *s) {
         } else
                 log_debug("Invalid DNS stub TCP packet, ignoring.");
 
-        /* Drop the reference to the stream. Either a query was created and added its own reference to the stream now,
-         * or that didn't happen in which case we want to free the stream */
-        dns_stream_unref(s);
-
         return 0;
 }
 
@@ -478,9 +471,9 @@ static int on_dns_stub_stream(sd_event_source *s, int fd, uint32_t revents, void
         }
 
         stream->on_packet = on_dns_stub_stream_packet;
+        stream->complete = dns_stub_stream_complete;
 
-        /* We let the reference to the stream dangling here, it will either be dropped by the default "complete" action
-         * of the stream, or by our packet callback, or when the manager is shut down. */
+        /* We let the reference to the stream dangle here, it will be dropped later by the complete callback. */
 
         return 0;
 }