diff --git a/vmm/src/lib.rs b/vmm/src/lib.rs index f118e93e9..311ec84b8 100644 --- a/vmm/src/lib.rs +++ b/vmm/src/lib.rs @@ -2298,9 +2298,10 @@ impl RequestHandler for Vmm { receive_data_migration.receiver_url ); + let mut listener = + migration_transport::receive_migration_listener(&receive_data_migration.receiver_url)?; // Accept the connection and get the socket - let mut socket = - migration_transport::receive_migration_socket(&receive_data_migration.receiver_url)?; + let mut socket = listener.accept()?; event!("vm", "migration-receive-started"); diff --git a/vmm/src/migration_transport.rs b/vmm/src/migration_transport.rs index 12412b0a0..a33ee503c 100644 --- a/vmm/src/migration_transport.rs +++ b/vmm/src/migration_transport.rs @@ -23,6 +23,31 @@ use vm_migration::{MigratableError, Snapshot}; use crate::{GuestMemoryMmap, VmMigrationConfig}; +/// Transport-agnostic listener used to receive connections. +#[derive(Debug)] +pub(crate) enum ReceiveListener { + Tcp(TcpListener), + Unix(UnixListener), +} + +impl ReceiveListener { + /// Block until a connection is accepted. + pub(crate) fn accept(&mut self) -> Result { + match self { + ReceiveListener::Tcp(listener) => listener + .accept() + .map(|(socket, _)| SocketStream::Tcp(socket)) + .context("Failed to accept TCP migration connection") + .map_err(MigratableError::MigrateReceive), + ReceiveListener::Unix(listener) => listener + .accept() + .map(|(socket, _)| SocketStream::Unix(socket)) + .context("Failed to accept Unix migration connection") + .map_err(MigratableError::MigrateReceive), + } + } +} + /// Transport-agnostic stream used by the migration protocol. pub(crate) enum SocketStream { Unix(UnixStream), @@ -138,35 +163,21 @@ pub(crate) fn send_migration_socket( } } -/// Bind and accept a migration connection for the receiver side. -pub(crate) fn receive_migration_socket( +/// Bind a migration listener for the receiver side. +pub(crate) fn receive_migration_listener( receiver_url: &str, -) -> Result { +) -> Result { if let Some(address) = receiver_url.strip_prefix("tcp:") { - let listener = TcpListener::bind(address).map_err(|e| { - MigratableError::MigrateReceive(anyhow!("Error binding to TCP socket: {e}")) - })?; - - let (socket, _addr) = listener.accept().map_err(|e| { - MigratableError::MigrateReceive(anyhow!( - "Error accepting connection on TCP socket: {e}" - )) - })?; - - Ok(SocketStream::Tcp(socket)) + TcpListener::bind(address) + .map(ReceiveListener::Tcp) + .context("Error binding to TCP socket") + .map_err(MigratableError::MigrateReceive) } else { - let path = socket_url_to_path(receiver_url).map_err(MigratableError::MigrateSend)?; - let listener = UnixListener::bind(&path).map_err(|e| { - MigratableError::MigrateReceive(anyhow!("Error binding to UNIX socket: {e}")) - })?; - - let (socket, _addr) = listener.accept().map_err(|e| { - MigratableError::MigrateReceive(anyhow!( - "Error accepting connection on UNIX socket: {e}" - )) - })?; - - Ok(SocketStream::Unix(socket)) + let path = socket_url_to_path(receiver_url).map_err(MigratableError::MigrateReceive)?; + UnixListener::bind(&path) + .map(ReceiveListener::Unix) + .context("Error binding to UNIX socket") + .map_err(MigratableError::MigrateReceive) } }