resolved: Make event flags logic robust for DoT

Since when handling a DNS over TLS stream, the TLS library can override the
requested events through dnstls_events for handshake/shutdown purposes,
obtaining the event flags through sd_event_source_get_io_events and checking
for EPOLLIN or EPOLLOUT does not really tell us whether we want to read/write
a packet. Instead, it could just be OpenSSL/GnuTLS doing something else.

To make the logic more robust (and simpler), save the flags that tell us
whether we want to read/write a packet, and check them instead of the IO flags.

(& use uint32_t for the flags like in sd_event_source_set_io_events prototype)

2 files changed, 11 insertions, 19 deletions

diff --git a/src/resolve/resolved-dns-stream.c b/src/resolve/resolved-dns-stream.c
index cf9d1a9d5e..290c28ed65 100644
--- a/src/resolve/resolved-dns-stream.c
+++ b/src/resolve/resolved-dns-stream.c
@@ -27,7 +27,7 @@ static void dns_stream_stop(DnsStream *s) {
 }
 
 static int dns_stream_update_io(DnsStream *s) {
-        int f = 0;
+        uint32_t f = 0;
 
         assert(s);
 
@@ -47,6 +47,8 @@ static int dns_stream_update_io(DnsStream *s) {
                 set_size(s->queries) < DNS_QUERIES_PER_STREAM)
                 f |= EPOLLIN;
 
+        s->requested_events = f;
+
 #if ENABLE_DNS_OVER_TLS
         /* For handshake and clean closing purposes, TLS can override requested events */
         if (s->dnstls_events != 0)
@@ -452,19 +454,11 @@ static int on_stream_io_impl(DnsStream *s, uint32_t revents) {
                 }
         }
 
-        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);
-        }
+        /* Complete the stream if finished reading and writing one packet, and there's nothing
+         * else left to write. */
+        if (s->type == DNS_STREAM_LLMNR_SEND && s->packet_received &&
+            !FLAGS_SET(s->requested_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) {
@@ -499,10 +493,7 @@ static int on_stream_io(sd_event_source *es, int fd, uint32_t revents, void *use
                 uint32_t events;
 
                 /* Make sure the stream still wants to process more data... */
-                r = sd_event_source_get_io_events(s->io_event_source, &events);
-                if (r < 0)
-                        return r;
-                if (!FLAGS_SET(events, EPOLLIN))
+                if (!FLAGS_SET(s->requested_events, EPOLLIN))
                         break;
 
                 r = on_stream_io_impl(s, EPOLLIN);
diff --git a/src/resolve/resolved-dns-stream.h b/src/resolve/resolved-dns-stream.h
index 1c606365cd..ba4a59e41c 100644
--- a/src/resolve/resolved-dns-stream.h
+++ b/src/resolve/resolved-dns-stream.h
@@ -61,6 +61,7 @@ struct DnsStream {
         uint32_t ttl;
         bool identified;
         bool packet_received; /* At least one packet is received. Used by LLMNR. */
+        uint32_t requested_events;
 
         /* only when using TCP fast open */
         union sockaddr_union tfo_address;
@@ -68,7 +69,7 @@ struct DnsStream {
 
 #if ENABLE_DNS_OVER_TLS
         DnsTlsStreamData dnstls_data;
-        int dnstls_events;
+        uint32_t dnstls_events;
 #endif
 
         sd_event_source *io_event_source;