From e72be7b593992deb16ab77742e934148d43ec952 Mon Sep 17 00:00:00 2001
From: David Stevens <stevensd@chromium.org>
Date: Wed, 19 Oct 2022 15:16:49 +0900
Subject: [PATCH] devices: vhost-user: wl: fix shared memory

Fix shared memory support for vhost-user wl implementation. Handle
set_slave_req_fd being called before get_shared_memory_regions and make
sure the wl device only tries to get its shared memory mapper once.

BUG=None
TEST=crosvm device wl

Change-Id: I724712399bd6a9ffbffb22dd69d431f54df58282
Reviewed-on: https://chromium-review.googlesource.com/c/crosvm/crosvm/+/3965812
Reviewed-by: Keiichi Watanabe <keiichiw@chromium.org>
Commit-Queue: David Stevens <stevensd@chromium.org>
---
 .../src/virtio/vhost/user/device/handler.rs   |  8 ++--
 devices/src/virtio/vhost/user/device/wl.rs    | 37 ++++++++++---------
 2 files changed, 24 insertions(+), 21 deletions(-)

diff --git a/devices/src/virtio/vhost/user/device/handler.rs b/devices/src/virtio/vhost/user/device/handler.rs
index 4910678130..90b3638f1e 100644
--- a/devices/src/virtio/vhost/user/device/handler.rs
+++ b/devices/src/virtio/vhost/user/device/handler.rs
@@ -372,7 +372,6 @@ where
     mem: Option<GuestMemory>,
     backend: Box<dyn VhostUserBackend>,
     ops: O,
-    shmid: Option<u8>,
 }
 
 impl DeviceRequestHandler<VhostUserRegularOps> {
@@ -400,7 +399,6 @@ where
             mem: None,
             backend,
             ops,
-            shmid: None,
         }
     }
 }
@@ -648,7 +646,10 @@ impl<O: VhostUserPlatformOps> VhostUserSlaveReqHandlerMut for DeviceRequestHandl
     }
 
     fn set_slave_req_fd(&mut self, ep: Box<dyn Endpoint<SlaveReq>>) {
-        let conn = VhostBackendReqConnection::new(Slave::new(ep), self.shmid);
+        let conn = VhostBackendReqConnection::new(
+            Slave::new(ep),
+            self.backend.get_shared_memory_region().map(|r| r.id),
+        );
         self.backend.set_backend_req_connection(conn);
     }
 
@@ -684,7 +685,6 @@ impl<O: VhostUserPlatformOps> VhostUserSlaveReqHandlerMut for DeviceRequestHandl
 
     fn get_shared_memory_regions(&mut self) -> VhostResult<Vec<VhostSharedMemoryRegion>> {
         Ok(if let Some(r) = self.backend.get_shared_memory_region() {
-            self.shmid = Some(r.id);
             vec![VhostSharedMemoryRegion::new(r.id, r.length)]
         } else {
             Vec::new()
diff --git a/devices/src/virtio/vhost/user/device/wl.rs b/devices/src/virtio/vhost/user/device/wl.rs
index b0386786c0..471d821224 100644
--- a/devices/src/virtio/vhost/user/device/wl.rs
+++ b/devices/src/virtio/vhost/user/device/wl.rs
@@ -224,22 +224,22 @@ impl VhostUserBackend for WlBackend {
             ..
         } = self;
 
-        let mapper = {
-            match &mut self.backend_req_conn {
-                VhostBackendReqConnectionState::Connected(request) => {
-                    request.take_shmem_mapper()?
-                }
-                VhostBackendReqConnectionState::NoConnection => {
-                    bail!("No backend request connection found")
-                }
-            }
-        };
         #[cfg(feature = "minigbm")]
         let gralloc = RutabagaGralloc::new().context("Failed to initailize gralloc")?;
-        let wlstate = self
-            .wlstate
-            .get_or_insert_with(|| {
-                Rc::new(RefCell::new(wl::WlState::new(
+        let wlstate = match &self.wlstate {
+            None => {
+                let mapper = {
+                    match &mut self.backend_req_conn {
+                        VhostBackendReqConnectionState::Connected(request) => {
+                            request.take_shmem_mapper()?
+                        }
+                        VhostBackendReqConnectionState::NoConnection => {
+                            bail!("No backend request connection found")
+                        }
+                    }
+                };
+
+                let wlstate = Rc::new(RefCell::new(wl::WlState::new(
                     wayland_paths.take().expect("WlState already initialized"),
                     mapper,
                     *use_transition_flags,
@@ -248,9 +248,12 @@ impl VhostUserBackend for WlBackend {
                     #[cfg(feature = "minigbm")]
                     gralloc,
                     None, /* address_offset */
-                )))
-            })
-            .clone();
+                )));
+                self.wlstate = Some(wlstate.clone());
+                wlstate
+            }
+            Some(state) => state.clone(),
+        };
         let (handle, registration) = AbortHandle::new_pair();
         match idx {
             0 => {