vhost-user: Add support for GET_SHMEM_CONFIG message

Add support for GET_SHMEM_CONFIG message to retrieve
VirtIO Shared Memory Regions configuration.

This is useful when the frontend is unaware of specific
backend type and configuration of the memory layout.

Based on the patch [1] which is just waiting for being
merged.

[1] -
https://lore.kernel.org/all/[email protected]/

Signed-off-by: Albert Esteve <[email protected]>

Author: aesteve-rh

vhost-user-backend/CHANGELOG.md

@@ -3,6 +3,8 @@
 ## [Unreleased]
 
 ### Added
+- [[#339]](https://github.com/rust-vmm/vhost/pull/339) Add support for `GET_SHMEM_CONFIG` message
+
 ### Changed
 ### Deprecated
 ### Fixed

vhost-user-backend/src/backend.rs

@@ -25,7 +25,7 @@ use std::sync::{Arc, Mutex, RwLock};
 
 use vhost::vhost_user::message::{
     VhostTransferStateDirection, VhostTransferStatePhase, VhostUserProtocolFeatures,
-    VhostUserSharedMsg,
+    VhostUserShMemConfig, VhostUserSharedMsg,
 };
 use vhost::vhost_user::Backend;
 use vm_memory::bitmap::Bitmap;
@@ -180,6 +180,13 @@ pub trait VhostUserBackend: Send + Sync {
             "back end does not support state transfer",
         ))
     }
+
+    fn get_shmem_config(&self) -> Result<VhostUserShMemConfig> {
+        Err(std::io::Error::new(
+            std::io::ErrorKind::Unsupported,
+            "back end does not support shared memory regions",
+        ))
+    }
 }
 
 /// Trait without interior mutability for vhost user backend servers to implement concrete services.
@@ -322,6 +329,13 @@ pub trait VhostUserBackendMut: Send + Sync {
             "back end does not support state transfer",
         ))
     }
+
+    fn get_shmem_config(&self) -> Result<VhostUserShMemConfig> {
+        Err(std::io::Error::new(
+            std::io::ErrorKind::Unsupported,
+            "back end does not support shared memory regions",
+        ))
+    }
 }
 
 impl<T: VhostUserBackend> VhostUserBackend for Arc<T> {
@@ -411,6 +425,10 @@ impl<T: VhostUserBackend> VhostUserBackend for Arc<T> {
     fn check_device_state(&self) -> Result<()> {
         self.deref().check_device_state()
     }
+
+    fn get_shmem_config(&self) -> Result<VhostUserShMemConfig> {
+        self.deref().get_shmem_config()
+    }
 }
 
 impl<T: VhostUserBackendMut> VhostUserBackend for Mutex<T> {
@@ -503,6 +521,10 @@ impl<T: VhostUserBackendMut> VhostUserBackend for Mutex<T> {
     fn check_device_state(&self) -> Result<()> {
         self.lock().unwrap().check_device_state()
     }
+
+    fn get_shmem_config(&self) -> Result<VhostUserShMemConfig> {
+        self.lock().unwrap().get_shmem_config()
+    }
 }
 
 impl<T: VhostUserBackendMut> VhostUserBackend for RwLock<T> {
@@ -595,6 +617,10 @@ impl<T: VhostUserBackendMut> VhostUserBackend for RwLock<T> {
     fn check_device_state(&self) -> Result<()> {
         self.read().unwrap().check_device_state()
     }
+
+    fn get_shmem_config(&self) -> Result<VhostUserShMemConfig> {
+        self.read().unwrap().get_shmem_config()
+    }
 }
 
 #[cfg(test)]

vhost-user-backend/src/handler.rs

@@ -18,7 +18,7 @@ use crate::bitmap::{BitmapReplace, MemRegionBitmap, MmapLogReg};
 use userfaultfd::{Uffd, UffdBuilder};
 use vhost::vhost_user::message::{
     VhostTransferStateDirection, VhostTransferStatePhase, VhostUserConfigFlags, VhostUserLog,
-    VhostUserMemoryRegion, VhostUserProtocolFeatures, VhostUserSharedMsg,
+    VhostUserMemoryRegion, VhostUserProtocolFeatures, VhostUserShMemConfig, VhostUserSharedMsg,
     VhostUserSingleMemoryRegion, VhostUserVirtioFeatures, VhostUserVringAddrFlags,
     VhostUserVringState,
 };
@@ -677,6 +677,12 @@ where
             .map_err(VhostUserError::ReqHandlerError)
     }
 
+    fn get_shmem_config(&self) -> VhostUserResult<VhostUserShMemConfig> {
+        self.backend
+            .get_shmem_config()
+            .map_err(VhostUserError::ReqHandlerError)
+    }
+
     #[cfg(feature = "postcopy")]
     fn postcopy_advice(&mut self) -> VhostUserResult<File> {
         let mut uffd_builder = UffdBuilder::new();

vhost/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Added
 - [[#251]](https://github.com/rust-vmm/vhost/pull/251) Add `SHMEM_MAP` and `SHMEM_UNMAP` support
+- [[#339]](https://github.com/rust-vmm/vhost/pull/339) Add support for GET_SHMEM_CONFIG message
 
 ### Changed
 ### Deprecated

vhost/src/vhost_user/backend_req_handler.rs

@@ -81,6 +81,7 @@ pub trait VhostUserBackendReqHandler {
         fd: File,
     ) -> Result<Option<File>>;
     fn check_device_state(&self) -> Result<()>;
+    fn get_shmem_config(&self) -> Result<VhostUserShMemConfig>;
     #[cfg(feature = "postcopy")]
     fn postcopy_advice(&self) -> Result<File>;
     #[cfg(feature = "postcopy")]
@@ -146,6 +147,7 @@ pub trait VhostUserBackendReqHandlerMut {
         fd: File,
     ) -> Result<Option<File>>;
     fn check_device_state(&mut self) -> Result<()>;
+    fn get_shmem_config(&self) -> Result<VhostUserShMemConfig>;
     #[cfg(feature = "postcopy")]
     fn postcopy_advice(&mut self) -> Result<File>;
     #[cfg(feature = "postcopy")]
@@ -289,6 +291,10 @@ impl<T: VhostUserBackendReqHandlerMut> VhostUserBackendReqHandler for Mutex<T> {
         self.lock().unwrap().check_device_state()
     }
 
+    fn get_shmem_config(&self) -> Result<VhostUserShMemConfig> {
+        self.lock().unwrap().get_shmem_config()
+    }
+
     #[cfg(feature = "postcopy")]
     fn postcopy_advice(&self) -> Result<File> {
         self.lock().unwrap().postcopy_advice()
@@ -679,6 +685,11 @@ impl<S: VhostUserBackendReqHandler> BackendReqHandler<S> {
                 };
                 self.send_reply_message(&hdr, &msg)?;
             }
+            Ok(FrontendReq::GET_SHMEM_CONFIG) => {
+                self.check_proto_feature(VhostUserProtocolFeatures::SHMEM)?;
+                let msg = self.backend.get_shmem_config().unwrap_or_default();
+                self.send_reply_message(&hdr, &msg)?;
+            }
             #[cfg(feature = "postcopy")]
             Ok(FrontendReq::POSTCOPY_ADVISE) => {
                 self.check_proto_feature(VhostUserProtocolFeatures::PAGEFAULT)?;
@@ -1038,4 +1049,111 @@ mod tests {
         handler.check_state().unwrap_err();
         assert!(handler.as_raw_fd() >= 0);
     }
+
+    #[test]
+    fn test_get_shmem_config_max_regions() {
+        // Create a configuration with maximum number of regions (8)
+        let memory_sizes = [
+            0x1000, 0x2000, 0x3000, 0x4000, 0x5000, 0x6000, 0x7000, 0x8000,
+        ];
+        let config = VhostUserShMemConfig::new(8, &memory_sizes);
+
+        let (p1, p2) = UnixStream::pair().unwrap();
+        let mut dummy_backend = DummyBackendReqHandler::new();
+        dummy_backend.set_shmem_config(config);
+        let backend = Arc::new(Mutex::new(dummy_backend));
+        let mut handler = BackendReqHandler::new(
+            Endpoint::<VhostUserMsgHeader<FrontendReq>>::from_stream(p1),
+            backend,
+        );
+        handler.acked_protocol_features = VhostUserProtocolFeatures::SHMEM.bits();
+
+        let mut frontend_endpoint = Endpoint::<VhostUserMsgHeader<FrontendReq>>::from_stream(p2);
+
+        let handle = std::thread::spawn(move || {
+            let hdr = VhostUserMsgHeader::new(FrontendReq::GET_SHMEM_CONFIG, 0, 0);
+            frontend_endpoint
+                .send_message(&hdr, &VhostUserEmpty, None)
+                .unwrap();
+
+            let (reply_hdr, reply_config, rfds) = frontend_endpoint
+                .recv_body::<VhostUserShMemConfig>()
+                .unwrap();
+            assert_eq!(reply_hdr.get_code().unwrap(), FrontendReq::GET_SHMEM_CONFIG);
+            assert!(reply_hdr.is_reply());
+            assert!(rfds.is_none());
+            reply_config
+        });
+
+        handler.handle_request().unwrap();
+
+        let reply_config = handle.join().unwrap();
+        assert_eq!(reply_config.nregions, 8);
+        for i in 0..8 {
+            assert_eq!(reply_config.memory_sizes[i], (i as u64 + 1) * 0x1000);
+        }
+    }
+
+    #[test]
+    fn test_get_shmem_config_non_continuous_regions() {
+        // Create a configuration with non-continuous regions
+        let memory_sizes = [0x10000, 0, 0x20000, 0, 0, 0, 0, 0];
+        let config = VhostUserShMemConfig::new(2, &memory_sizes);
+
+        let (p1, p2) = UnixStream::pair().unwrap();
+        let mut dummy_backend = DummyBackendReqHandler::new();
+        dummy_backend.set_shmem_config(config);
+        let backend = Arc::new(Mutex::new(dummy_backend));
+        let mut handler = BackendReqHandler::new(
+            Endpoint::<VhostUserMsgHeader<FrontendReq>>::from_stream(p1),
+            backend,
+        );
+        handler.acked_protocol_features = VhostUserProtocolFeatures::SHMEM.bits();
+
+        let mut frontend_endpoint = Endpoint::<VhostUserMsgHeader<FrontendReq>>::from_stream(p2);
+
+        let handle = std::thread::spawn(move || {
+            let hdr = VhostUserMsgHeader::new(FrontendReq::GET_SHMEM_CONFIG, 0, 0);
+            frontend_endpoint
+                .send_message(&hdr, &VhostUserEmpty, None)
+                .unwrap();
+
+            let (reply_hdr, reply_config, rfds) = frontend_endpoint
+                .recv_body::<VhostUserShMemConfig>()
+                .unwrap();
+            assert_eq!(reply_hdr.get_code().unwrap(), FrontendReq::GET_SHMEM_CONFIG);
+            assert!(reply_hdr.is_reply());
+            assert!(rfds.is_none());
+            reply_config
+        });
+
+        handler.handle_request().unwrap();
+
+        let reply_config = handle.join().unwrap();
+        assert_eq!(reply_config.nregions, 2);
+        assert_eq!(reply_config.memory_sizes[0], 0x10000);
+        assert_eq!(reply_config.memory_sizes[1], 0);
+        assert_eq!(reply_config.memory_sizes[2], 0x20000);
+        for i in 3..8 {
+            assert_eq!(reply_config.memory_sizes[i], 0);
+        }
+    }
+
+    #[test]
+    fn test_get_shmem_config_feature_not_negotiated() {
+        // Test that the request fails when SHMEM protocol feature is not negotiated
+        let (p1, p2) = UnixStream::pair().unwrap();
+        let backend = Arc::new(Mutex::new(DummyBackendReqHandler::new()));
+        let mut handler = BackendReqHandler::new(
+            Endpoint::<VhostUserMsgHeader<FrontendReq>>::from_stream(p1),
+            backend,
+        );
+        let mut frontend_endpoint = Endpoint::<VhostUserMsgHeader<FrontendReq>>::from_stream(p2);
+
+        std::thread::spawn(move || {
+            let hdr = VhostUserMsgHeader::new(FrontendReq::GET_SHMEM_CONFIG, 0, 0);
+            let _ = frontend_endpoint.send_message(&hdr, &VhostUserEmpty, None);
+        });
+        assert!(handler.handle_request().is_err());
+    }
 }

vhost/src/vhost_user/dummy_backend.rs

@@ -27,6 +27,7 @@ pub struct DummyBackendReqHandler {
     pub vring_enabled: [bool; MAX_QUEUE_NUM],
     pub inflight_file: Option<File>,
     pub shared_file: Option<File>,
+    pub shmem_config: Option<VhostUserShMemConfig>,
 }
 
 impl DummyBackendReqHandler {
@@ -37,6 +38,12 @@ impl DummyBackendReqHandler {
         }
     }
 
+    /// Set the shared memory configuration to be returned by `get_shmem_config`
+    pub fn set_shmem_config(&mut self, config: VhostUserShMemConfig) {
+        self.acked_protocol_features |= VhostUserProtocolFeatures::SHMEM.bits();
+        self.shmem_config = Some(config);
+    }
+
     /// Helper to check if VirtioFeature enabled
     fn check_feature(&self, feat: VhostUserVirtioFeatures) -> Result<()> {
         if self.acked_features & feat.bits() != 0 {
@@ -329,6 +336,15 @@ impl VhostUserBackendReqHandlerMut for DummyBackendReqHandler {
         )))
     }
 
+    fn get_shmem_config(&self) -> Result<VhostUserShMemConfig> {
+        self.shmem_config.ok_or_else(|| {
+            Error::ReqHandlerError(std::io::Error::new(
+                std::io::ErrorKind::Unsupported,
+                "dummy back end does not support shared memory regions",
+            ))
+        })
+    }
+
     #[cfg(feature = "postcopy")]
     fn postcopy_advice(&mut self) -> Result<File> {
         let file = tempfile::tempfile().unwrap();

vhost/src/vhost_user/message.rs

@@ -169,6 +169,8 @@ enum_value! {
         /// After transferring state, check the backend for any errors that may have
         /// occurred during the transfer
         CHECK_DEVICE_STATE = 43,
+        /// Get shared memory regions configuration from the backend.
+        GET_SHMEM_CONFIG = 44,
     }
 }
 
@@ -688,6 +690,34 @@ impl VhostUserSingleMemoryRegion {
 unsafe impl ByteValued for VhostUserSingleMemoryRegion {}
 impl VhostUserMsgValidator for VhostUserSingleMemoryRegion {}
 
+/// Get shared memory regions configuration.
+#[repr(C)]
+#[derive(Debug, Default, Clone, Copy)]
+pub struct VhostUserShMemConfig {
+    /// Total number of shared memory regions
+    pub nregions: u32,
+    /// Padding for correct alignment
+    padding: u32,
+    /// Size of each memory region
+    pub memory_sizes: [u64; 8],
+}
+
+impl VhostUserShMemConfig {
+    /// Create a new instance
+    pub fn new(nregions: u32, memory: &[u64]) -> Self {
+        let memory_sizes: [u64; 8] = std::array::from_fn(|i| *memory.get(i).unwrap_or(&0));
+        Self {
+            nregions,
+            padding: 0,
+            memory_sizes,
+        }
+    }
+}
+
+// SAFETY: Safe because all fields of VhostUserSingleMemoryRegion are POD.
+unsafe impl ByteValued for VhostUserShMemConfig {}
+impl VhostUserMsgValidator for VhostUserShMemConfig {}
+
 /// Vring state descriptor.
 #[repr(C, packed)]
 #[derive(Copy, Clone, Default)]