vmspawn: rework how AF_VSOCK/SOCK_STREAM notifications are read

Stream sockets are stream sockets, i.e. they won#t give us the full data
right-away, we must buffer locally and read until we hit EOF. Hence do
so.

moreover, make sure to close the fd once we are done, otherwise the
sender might block on us.

1 files changed, 97 insertions, 42 deletions

diff --git a/src/vmspawn/vmspawn.c b/src/vmspawn/vmspawn.c
index 584e8062b0..77100f59c2 100644
--- a/src/vmspawn/vmspawn.c
+++ b/src/vmspawn/vmspawn.c
@@ -605,71 +605,100 @@ static int open_vsock(void) {
         return TAKE_FD(vsock_fd);
 }
 
-static int vmspawn_dispatch_notify_fd(sd_event_source *source, int fd, uint32_t revents, void *userdata) {
-        char buf[NOTIFY_BUFFER_MAX+1];
-        const char *p = NULL;
-        struct iovec iovec = {
-                .iov_base = buf,
-                .iov_len = sizeof(buf)-1,
-        };
-        struct msghdr msghdr = {
-                .msg_iov = &iovec,
-                .msg_iovlen = 1,
-        };
-        ssize_t n;
-        _cleanup_strv_free_ char **tags = NULL;
-        int r, *exit_status = ASSERT_PTR(userdata);
+typedef struct NotifyConnectionData {
+        char buffer[NOTIFY_BUFFER_MAX+1];
+        size_t full;
+        int *exit_status;
+} NotifyConnectionData;
 
-        n = recvmsg_safe(fd, &msghdr, MSG_DONTWAIT);
-        if (ERRNO_IS_NEG_TRANSIENT(n))
-                return 0;
-        if (n == -EXFULL) {
-                log_warning_errno(n, "Got message with truncated control data, ignoring: %m");
-                return 0;
-        }
-        if (n < 0)
-                return log_warning_errno(n, "Couldn't read notification socket: %m");
+static int read_vsock_notify(NotifyConnectionData *d, int fd) {
+        int r;
 
-        if ((size_t) n >= sizeof(buf)) {
-                log_warning("Received notify message exceeded maximum size. Ignoring.");
-                return 0;
+        assert(d);
+        assert(fd >= 0);
+
+        for (;;) {
+                assert(d->full < sizeof(d->buffer));
+
+                ssize_t n = read(fd, d->buffer + d->full, sizeof(d->buffer) - d->full);
+                if (n < 0) {
+                        if (ERRNO_IS_TRANSIENT(errno))
+                                return 0;
+
+                        return log_error_errno(errno, "Failed to read notification message: %m");
+                }
+                if (n == 0) /* We hit EOF! Let's parse this */
+                        break;
+
+                d->full += n;
+
+                if (d->full >= sizeof(d->buffer))
+                        return log_error_errno(SYNTHETIC_ERRNO(EBADMSG), "Received notify message exceeded maximum size.");
         }
 
-        buf[n] = 0;
-        tags = strv_split(buf, "\n\r");
+        /* We reached EOF, now parse the thing */
+        assert(d->full < sizeof(d->buffer));
+        d->buffer[d->full] = 0;
+
+        _cleanup_strv_free_ char **tags = strv_split(d->buffer, "\n\r");
         if (!tags)
                 return log_oom();
 
-        STRV_FOREACH(s, tags)
-                log_debug("Received tag %s from notify socket", *s);
+        if (DEBUG_LOGGING) {
+                _cleanup_free_ char *j = strv_join(tags, " ");
+                log_debug("Received notification message with tags: %s", strnull(j));
+        }
 
         if (strv_contains(tags, "READY=1")) {
-                r = sd_notify(false, "READY=1\n");
+                r = sd_notify(false, "READY=1");
                 if (r < 0)
                         log_warning_errno(r, "Failed to send readiness notification, ignoring: %m");
         }
 
-        p = strv_find_startswith(tags, "STATUS=");
+        const char *p = strv_find_startswith(tags, "STATUS=");
         if (p)
                 (void) sd_notifyf(false, "STATUS=VM running: %s", p);
 
         p = strv_find_startswith(tags, "EXIT_STATUS=");
         if (p) {
-                r = safe_atoi(p, exit_status);
+                r = safe_atoi(p, d->exit_status);
                 if (r < 0)
                         log_warning_errno(r, "Failed to parse exit status from %s, ignoring: %m", p);
         }
 
-        /* we will only receive one message from each connection so disable this source once one is received */
-        source = sd_event_source_disable_unref(source);
+        return 1; /* done */
+}
+
+static int vmspawn_dispatch_notify_fd(sd_event_source *source, int fd, uint32_t revents, void *userdata) {
+        NotifyConnectionData *d = ASSERT_PTR(userdata);
+        int r;
+
+        assert(source);
+        assert(fd >= 0);
+
+        r = read_vsock_notify(d, fd);
+        if (r != 0) {
+                int q;
+
+                /* If we are done or are seeing an error we'll turn off floating mode, which means the event
+                 * loop itself won't keep the event source pinned anymore, and since noone else (hopefully!)
+                 * keeps a reference anymore the whole thing will be released once we exit from this handler
+                 * here. */
+
+                q = sd_event_source_set_floating(source, false);
+                if (q < 0)
+                        log_warning_errno(q, "Failed to disable floating mode of event source, ignoring: %m");
+
+                return r;
+        }
 
         return 0;
 }
 
 static int vmspawn_dispatch_vsock_connections(sd_event_source *source, int fd, uint32_t revents, void *userdata) {
-        int r;
-        sd_event *event;
         _cleanup_close_ int conn_fd = -EBADF;
+        sd_event *event;
+        int r;
 
         assert(userdata);
 
@@ -680,7 +709,10 @@ static int vmspawn_dispatch_vsock_connections(sd_event_source *source, int fd, u
 
         conn_fd = accept4(fd, NULL, NULL, SOCK_CLOEXEC|SOCK_NONBLOCK);
         if (conn_fd < 0) {
-                log_warning_errno(errno, "Failed to accept connection from VSOCK fd (%m), ignoring...");
+                if (ERRNO_IS_TRANSIENT(errno))
+                        return 0;
+
+                log_warning_errno(errno, "Failed to accept connection from VSOCK connection, ignoring: %m");
                 return 0;
         }
 
@@ -688,13 +720,36 @@ static int vmspawn_dispatch_vsock_connections(sd_event_source *source, int fd, u
         if (!event)
                 return log_error_errno(SYNTHETIC_ERRNO(ENOENT), "Failed to retrieve event from event source, exiting task");
 
+        _cleanup_free_ NotifyConnectionData *d = new(NotifyConnectionData, 1);
+        if (!d)
+                return log_oom();
+
+        *d = (NotifyConnectionData) {
+                .exit_status = userdata,
+        };
+
         /* add a new floating task to read from the connection */
-        r = sd_event_add_io(event, NULL, conn_fd, revents, vmspawn_dispatch_notify_fd, userdata);
+        _cleanup_(sd_event_source_unrefp) sd_event_source *s = NULL;
+        r = sd_event_add_io(event, &s, conn_fd, EPOLLIN, vmspawn_dispatch_notify_fd, d);
         if (r < 0)
                 return log_error_errno(r, "Failed to allocate notify connection event source: %m");
 
-        /* conn_fd is now owned by the event loop so don't clean it up */
-        TAKE_FD(conn_fd);
+        r = sd_event_source_set_io_fd_own(s, true);
+        if (r < 0)
+                return log_error_errno(r, "Failed to pass ownership of notify to event source: %m");
+        TAKE_FD(conn_fd); /* conn_fd is now owned by the event loop so don't clean it up */
+
+        r = sd_event_source_set_destroy_callback(s, free);
+        if (r < 0)
+                return log_error_errno(r, "Failed to set destroy callback on event source: %m");
+        TAKE_PTR(d); /* The data object will now automatically be freed by the event source when it goes away */
+
+        /* Finally, make sure the event loop pins the event source */
+        r = sd_event_source_set_floating(s, true);
+        if (r < 0)
+                return log_error_errno(r, "Failed to set event source to floating mode: %m");
+
+        (void) sd_event_source_set_description(s, "vmspawn-notify-socket-connection");
 
         return 0;
 }
@@ -711,7 +766,7 @@ static int setup_notify_parent(sd_event *event, int fd, int *exit_status, sd_eve
         if (r < 0)
                 return log_error_errno(r, "Failed to allocate notify socket event source: %m");
 
-        (void) sd_event_source_set_description(*ret_notify_event_source, "vmspawn-notify-sock");
+        (void) sd_event_source_set_description(*ret_notify_event_source, "vmspawn-notify-socket-listen");
 
         return 0;
 }