From 2fbb98e2c3ee6e52164f9c0d3577a14832cde415 Mon Sep 17 00:00:00 2001
From: Rob Bradford <rbradford@meta.com>
Date: Thu, 26 Mar 2026 15:16:58 -0700
Subject: [PATCH] virtio-devices: vhost_user: Add common code for restoring
 state

Add a common method for validating the state (checking vrings &
device_state) and then restoring the backend state if present.

Signed-off-by: Rob Bradford <rbradford@meta.com>
---
 virtio-devices/src/vhost_user/mod.rs            | 11 +++++++++++
 virtio-devices/src/vhost_user/vu_common_ctrl.rs | 10 +++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/virtio-devices/src/vhost_user/mod.rs b/virtio-devices/src/vhost_user/mod.rs
index 3858c5c05..86a774737 100644
--- a/virtio-devices/src/vhost_user/mod.rs
+++ b/virtio-devices/src/vhost_user/mod.rs
@@ -156,6 +156,8 @@ pub enum Error {
     SaveRestoreBackendState(#[source] io::Error),
     #[error("Vring bases count ({0}) does not match queue count ({1})")]
     VringBasesCountMismatch(usize, usize),
+    #[error("Backend state and vring bases must both be present or both be absent")]
+    InconsistentBackendState,
 }
 type Result<T> = std::result::Result<T, Error>;
 
@@ -323,6 +325,15 @@ pub struct VhostUserState<C> {
     pub backend_state: Option<Vec<u8>>,
 }
 
+impl<C> VhostUserState<C> {
+    pub fn validate(&self) -> Result<()> {
+        if self.backend_state.is_some() != self.vring_bases.is_some() {
+            return Err(Error::InconsistentBackendState);
+        }
+        Ok(())
+    }
+}
+
 #[derive(Default)]
 pub struct VhostUserCommon {
     pub vu: Option<Arc<Mutex<VhostUserHandle>>>,
diff --git a/virtio-devices/src/vhost_user/vu_common_ctrl.rs b/virtio-devices/src/vhost_user/vu_common_ctrl.rs
index 0e59e2399..b2ee2ae5d 100644
--- a/virtio-devices/src/vhost_user/vu_common_ctrl.rs
+++ b/virtio-devices/src/vhost_user/vu_common_ctrl.rs
@@ -28,7 +28,7 @@ use vm_memory::{Address, FileOffset, GuestAddress, GuestMemory, GuestMemoryRegio
 use vm_migration::protocol::MemoryRangeTable;
 use vmm_sys_util::eventfd::EventFd;
 
-use super::{Error, Result};
+use super::{Error, Result, VhostUserState};
 use crate::vhost_user::Inflight;
 use crate::{
     GuestMemoryMmap, GuestRegionMmap, MmapRegion, VirtioInterrupt, VirtioInterruptType,
@@ -512,6 +512,14 @@ impl VhostUserHandle {
         Ok((state, vring_bases))
     }
 
+    pub fn restore_state<C>(&mut self, state: &VhostUserState<C>) -> Result<()> {
+        state.validate()?;
+        if let Some(backend_state) = &state.backend_state {
+            self.restore_backend_state(backend_state)?;
+        }
+        Ok(())
+    }
+
     /// Restore backend device state via the SET_DEVICE_STATE_FD protocol.
     /// Sends the saved opaque state blob to the backend via a socket.
     pub fn restore_backend_state(&mut self, state: &[u8]) -> Result<()> {