diff --git a/vmm/src/migration_transport.rs b/vmm/src/migration_transport.rs index a33ee503c..cda22dfac 100644 --- a/vmm/src/migration_transport.rs +++ b/vmm/src/migration_transport.rs @@ -5,6 +5,7 @@ use std::io::{self, Read, Write}; use std::net::{TcpListener, TcpStream}; +use std::os::fd::{AsFd, BorrowedFd}; use std::os::unix::io::{AsRawFd, RawFd}; use std::os::unix::net::{UnixListener, UnixStream}; use std::path::PathBuf; @@ -46,6 +47,32 @@ impl ReceiveListener { .map_err(MigratableError::MigrateReceive), } } + + /// Same as [`Self::accept`], but returns `None` if the abort event was signaled. + fn abortable_accept( + &mut self, + abort_event: &impl AsRawFd, + ) -> Result, MigratableError> { + if wait_for_readable(&self, abort_event) + .context("Error while waiting for socket to become readable") + .map_err(MigratableError::MigrateReceive)? + { + // The listener is readable; accept the connection. + Ok(Some(self.accept()?)) + } else { + // The abort event was signaled before any connection arrived. + Ok(None) + } + } +} + +impl AsFd for ReceiveListener { + fn as_fd(&self) -> BorrowedFd<'_> { + match self { + ReceiveListener::Tcp(listener) => listener.as_fd(), + ReceiveListener::Unix(listener) => listener.as_fd(), + } + } } /// Transport-agnostic stream used by the migration protocol. @@ -132,6 +159,55 @@ impl WriteVolatile for SocketStream { } } +// Wait for `fd` to become readable. In this case, we return true. In case +// `abort_event` was signaled, return false. +fn wait_for_readable(fd: &impl AsFd, abort_event: &impl AsRawFd) -> Result { + let fd = fd.as_fd().as_raw_fd(); + let abort_event = abort_event.as_raw_fd(); + + let mut poll_fds = [ + libc::pollfd { + fd: abort_event, + events: libc::POLLIN, + revents: 0, + }, + libc::pollfd { + fd, + events: libc::POLLIN, + revents: 0, + }, + ]; + + loop { + // SAFETY: This is safe, because the file descriptors are valid and the + // poll_fds array is properly initialized. + let ret = unsafe { libc::poll(poll_fds.as_mut_ptr(), poll_fds.len() as libc::nfds_t, -1) }; + + if ret >= 0 { + break; + } + + let err = io::Error::last_os_error(); + if err.raw_os_error() == Some(libc::EINTR) { + continue; + } + + return Err(err); + } + + if poll_fds[0].revents & libc::POLLIN != 0 { + return Ok(false); + } + + if poll_fds[1].revents & libc::POLLIN != 0 { + return Ok(true); + } + + Err(io::Error::other( + "Poll returned, but neither file descriptor is readable?", + )) +} + /// Extract a UNIX socket path from a "unix:" migration URL. fn socket_url_to_path(url: &str) -> Result { url.strip_prefix("unix:")