vmm: move SocketStream into the migration_transport module

This is mainly to clean up the lib.rs a bit more.

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

vmm/src/lib.rs

@@ -6,9 +6,7 @@
 use std::collections::HashMap;
 use std::fs::File;
 use std::io::{Read, Write, stdout};
-use std::net::TcpStream;
 use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
-use std::os::unix::net::UnixStream;
 use std::panic::AssertUnwindSafe;
 #[cfg(feature = "guest_debug")]
 use std::path::PathBuf;
@@ -38,8 +36,8 @@ use serde::{Deserialize, Serialize};
 use signal_hook::iterator::{Handle, Signals};
 use thiserror::Error;
 use tracer::trace_scoped;
-use vm_memory::bitmap::{AtomicBitmap, BitmapSlice};
-use vm_memory::{ReadVolatile, VolatileMemoryError, VolatileSlice, WriteVolatile};
+use vm_memory::ReadVolatile;
+use vm_memory::bitmap::AtomicBitmap;
 use vm_migration::protocol::*;
 use vm_migration::{
     MemoryMigrationContext, Migratable, MigratableError, Pausable, Snapshot, Snapshottable,
@@ -61,6 +59,7 @@ use crate::memory_manager::MemoryManager;
 #[cfg(all(feature = "kvm", target_arch = "x86_64"))]
 use crate::migration::get_vm_snapshot;
 use crate::migration::{recv_vm_config, recv_vm_state};
+use crate::migration_transport::SocketStream;
 use crate::seccomp_filters::{Thread, get_seccomp_filter};
 use crate::vm::{Error as VmError, Vm, VmState};
 use crate::vm_config::{
@@ -265,89 +264,6 @@ impl From<u64> for EpollDispatch {
     }
 }
 
-enum SocketStream {
-    Unix(UnixStream),
-    Tcp(TcpStream),
-}
-
-impl Read for SocketStream {
-    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
-        match self {
-            SocketStream::Unix(stream) => stream.read(buf),
-            SocketStream::Tcp(stream) => stream.read(buf),
-        }
-    }
-}
-
-impl Write for SocketStream {
-    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
-        match self {
-            SocketStream::Unix(stream) => stream.write(buf),
-            SocketStream::Tcp(stream) => stream.write(buf),
-        }
-    }
-
-    fn flush(&mut self) -> std::io::Result<()> {
-        match self {
-            SocketStream::Unix(stream) => stream.flush(),
-            SocketStream::Tcp(stream) => stream.flush(),
-        }
-    }
-}
-
-impl AsRawFd for SocketStream {
-    fn as_raw_fd(&self) -> RawFd {
-        match self {
-            SocketStream::Unix(s) => s.as_raw_fd(),
-            SocketStream::Tcp(s) => s.as_raw_fd(),
-        }
-    }
-}
-
-impl ReadVolatile for SocketStream {
-    fn read_volatile<B: BitmapSlice>(
-        &mut self,
-        buf: &mut VolatileSlice<B>,
-    ) -> std::result::Result<usize, VolatileMemoryError> {
-        match self {
-            SocketStream::Unix(s) => s.read_volatile(buf),
-            SocketStream::Tcp(s) => s.read_volatile(buf),
-        }
-    }
-
-    fn read_exact_volatile<B: BitmapSlice>(
-        &mut self,
-        buf: &mut VolatileSlice<B>,
-    ) -> std::result::Result<(), VolatileMemoryError> {
-        match self {
-            SocketStream::Unix(s) => s.read_exact_volatile(buf),
-            SocketStream::Tcp(s) => s.read_exact_volatile(buf),
-        }
-    }
-}
-
-impl WriteVolatile for SocketStream {
-    fn write_volatile<B: BitmapSlice>(
-        &mut self,
-        buf: &VolatileSlice<B>,
-    ) -> std::result::Result<usize, VolatileMemoryError> {
-        match self {
-            SocketStream::Unix(s) => s.write_volatile(buf),
-            SocketStream::Tcp(s) => s.write_volatile(buf),
-        }
-    }
-
-    fn write_all_volatile<B: BitmapSlice>(
-        &mut self,
-        buf: &VolatileSlice<B>,
-    ) -> std::result::Result<(), VolatileMemoryError> {
-        match self {
-            SocketStream::Unix(s) => s.write_all_volatile(buf),
-            SocketStream::Tcp(s) => s.write_all_volatile(buf),
-        }
-    }
-}
-
 pub struct EpollContext {
     epoll_file: File,
 }

vmm/src/migration_transport.rs

@@ -3,8 +3,9 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use std::io::Write;
+use std::io::{self, Read, Write};
 use std::net::{TcpListener, TcpStream};
+use std::os::unix::io::{AsRawFd, RawFd};
 use std::os::unix::net::{UnixListener, UnixStream};
 use std::path::PathBuf;
 use std::result::Result;
@@ -12,11 +13,99 @@ use std::result::Result;
 use anyhow::{Context, anyhow};
 use log::info;
 use serde_json;
-use vm_memory::{Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic};
+use vm_memory::bitmap::BitmapSlice;
+use vm_memory::{
+    Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic, ReadVolatile, VolatileMemoryError,
+    VolatileSlice, WriteVolatile,
+};
 use vm_migration::protocol::{MemoryRangeTable, Request, Response};
 use vm_migration::{MigratableError, Snapshot};
 
-use crate::{GuestMemoryMmap, SocketStream, VmMigrationConfig};
+use crate::{GuestMemoryMmap, VmMigrationConfig};
+
+/// Transport-agnostic stream used by the migration protocol.
+pub(crate) enum SocketStream {
+    Unix(UnixStream),
+    Tcp(TcpStream),
+}
+
+impl Read for SocketStream {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        match self {
+            SocketStream::Unix(stream) => stream.read(buf),
+            SocketStream::Tcp(stream) => stream.read(buf),
+        }
+    }
+}
+
+impl Write for SocketStream {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        match self {
+            SocketStream::Unix(stream) => stream.write(buf),
+            SocketStream::Tcp(stream) => stream.write(buf),
+        }
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        match self {
+            SocketStream::Unix(stream) => stream.flush(),
+            SocketStream::Tcp(stream) => stream.flush(),
+        }
+    }
+}
+
+impl AsRawFd for SocketStream {
+    fn as_raw_fd(&self) -> RawFd {
+        match self {
+            SocketStream::Unix(s) => s.as_raw_fd(),
+            SocketStream::Tcp(s) => s.as_raw_fd(),
+        }
+    }
+}
+
+impl ReadVolatile for SocketStream {
+    fn read_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &mut VolatileSlice<B>,
+    ) -> Result<usize, VolatileMemoryError> {
+        match self {
+            SocketStream::Unix(s) => s.read_volatile(buf),
+            SocketStream::Tcp(s) => s.read_volatile(buf),
+        }
+    }
+
+    fn read_exact_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &mut VolatileSlice<B>,
+    ) -> Result<(), VolatileMemoryError> {
+        match self {
+            SocketStream::Unix(s) => s.read_exact_volatile(buf),
+            SocketStream::Tcp(s) => s.read_exact_volatile(buf),
+        }
+    }
+}
+
+impl WriteVolatile for SocketStream {
+    fn write_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &VolatileSlice<B>,
+    ) -> Result<usize, VolatileMemoryError> {
+        match self {
+            SocketStream::Unix(s) => s.write_volatile(buf),
+            SocketStream::Tcp(s) => s.write_volatile(buf),
+        }
+    }
+
+    fn write_all_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &VolatileSlice<B>,
+    ) -> Result<(), VolatileMemoryError> {
+        match self {
+            SocketStream::Unix(s) => s.write_all_volatile(buf),
+            SocketStream::Tcp(s) => s.write_all_volatile(buf),
+        }
+    }
+}
 
 /// Extract a UNIX socket path from a "unix:" migration URL.
 fn socket_url_to_path(url: &str) -> Result<PathBuf, anyhow::Error> {