Merge pull request #11512 from poettering/resolved-stub-pipeline

support incoming pipelined TCP stream queries in DNS stub server

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);

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,12 +38,16 @@ 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
         /* 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
 
@@ -52,6 +58,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) {
@@ -281,6 +291,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 +335,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 +361,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)) {
@@ -420,10 +435,21 @@ 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);
 
+        /* 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;
 }
 
@@ -437,7 +463,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 +488,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 +498,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 +539,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;

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;
@@ -58,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;
@@ -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

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;
 }
 
@@ -471,16 +464,16 @@ 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;
         }
 
         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;
 }

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;
 }
 
@@ -554,9 +550,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 +579,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 +605,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 +616,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;
 
@@ -636,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;

src/resolve/resolved-llmnr.c

@@ -295,13 +295,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_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;
         }
 
         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;
 }
 

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);