vmm: extract small request/response helpers to reduce boilerplate

On-behalf-of: SAP sebastian.eydam@sap.com
Signed-off-by: Sebastian Eydam <sebastian.eydam@cyberus-technology.de>

vmm/src/lib.rs

@@ -1176,7 +1176,7 @@ impl Vmm {
         table.write_to(socket)?;
         // And then the memory itself
         vm.send_memory_regions(table, socket)?;
-        Response::read_from(socket)?.ok_or_abandon(
+        migration_transport::expect_ok_response(
             socket,
             MigratableError::MigrateSend(anyhow!("Error during dirty memory migration")),
         )?;
@@ -1379,9 +1379,9 @@ impl Vmm {
             migration_transport::send_migration_socket(&send_data_migration.destination_url)?;
 
         // Start the migration
-        Request::start().write_to(&mut socket)?;
+        migration_transport::send_request_expect_ok(
-        Response::read_from(&mut socket)?.ok_or_abandon(
             &mut socket,
+            Request::start(),
             MigratableError::MigrateSend(anyhow!("Error starting migration")),
         )?;
 
@@ -1439,7 +1439,7 @@ impl Vmm {
         socket
             .write_all(&config_data)
             .map_err(MigratableError::MigrateSocket)?;
-        Response::read_from(&mut socket)?.ok_or_abandon(
+        migration_transport::expect_ok_response(
             &mut socket,
             MigratableError::MigrateSend(anyhow!("Error during config migration")),
         )?;
@@ -1466,15 +1466,15 @@ impl Vmm {
         socket
             .write_all(&snapshot_data)
             .map_err(MigratableError::MigrateSocket)?;
-        Response::read_from(&mut socket)?.ok_or_abandon(
+        migration_transport::expect_ok_response(
             &mut socket,
             MigratableError::MigrateSend(anyhow!("Error during state migration")),
         )?;
         // Complete the migration
         // At this step, the receiving VMM will acquire disk locks again.
-        Request::complete().write_to(&mut socket)?;
+        migration_transport::send_request_expect_ok(
-        Response::read_from(&mut socket)?.ok_or_abandon(
             &mut socket,
+            Request::complete(),
             MigratableError::MigrateSend(anyhow!("Error completing migration")),
         )?;
 

vmm/src/migration_transport.rs

@@ -11,6 +11,7 @@ use std::result::Result;
 use anyhow::anyhow;
 use log::info;
 use vm_migration::MigratableError;
+use vm_migration::protocol::{Request, Response};
 
 use crate::SocketStream;
 
@@ -76,3 +77,23 @@ pub(crate) fn receive_migration_socket(
         Ok(SocketStream::Unix(socket))
     }
 }
+
+/// Read a response and return Ok(()) if it was a [`Response::Ok`].
+pub(crate) fn expect_ok_response(
+    socket: &mut SocketStream,
+    error: MigratableError,
+) -> Result<(), MigratableError> {
+    Response::read_from(socket)?
+        .ok_or_abandon(socket, error)
+        .map(|_| ())
+}
+
+/// Send a request and validate that the peer responds with OK.
+pub(crate) fn send_request_expect_ok(
+    socket: &mut SocketStream,
+    request: Request,
+    error: MigratableError,
+) -> Result<(), MigratableError> {
+    request.write_to(socket)?;
+    expect_ok_response(socket, error)
+}