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

@@ -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;
+                (void) set_remove(q->request_dns_stream->queries, q);
-                q->request_dns_stream->on_packet = NULL;
+                q->request_dns_stream = dns_stream_unref(q->request_dns_stream);
-                q->request_dns_stream->query = NULL;
-                dns_stream_unref(q->request_dns_stream);
         }
 
         free(q->request_address_string);

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

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;

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);
+        for (;;) {
-        dns_query_free(s->query);
+                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;
+                /* Remember which queries belong to this stream, so that we can cancel them when the stream
-                s->complete = dns_stub_stream_complete;
+                 * is disconnected early */
-                s->query = q;
+
+                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
+        /* We let the reference to the stream dangle here, it will be dropped later by the complete callback. */
-         * of the stream, or by our packet callback, or when the manager is shut down. */
 
         return 0;
 }