toplevel_info: Fix racy data creation

pop-os · Oct 24, 2024 · 758049e · 758049e
1 parent 65a5470
commit 758049e
Showing 1 changed file with 28 additions and 7 deletions.
diff --git a/src/wayland/protocols/toplevel_info.rs b/src/wayland/protocols/toplevel_info.rs
@@ -90,7 +90,7 @@ pub struct ToplevelHandleStateInner<W: Window> {
     title: String,
     app_id: String,
     states: Vec<States>,
-    pub(super) window: W,
+    pub(super) window: Option<W>,
 }
 pub type ToplevelHandleState<W> = Mutex<ToplevelHandleStateInner<W>>;
 
@@ -104,7 +104,20 @@ impl<W: Window> ToplevelHandleStateInner<W> {
             title: String::new(),
             app_id: String::new(),
             states: Vec::new(),
-            window: window.clone(),
+            window: Some(window.clone()),
+        })
+    }
+
+    fn empty() -> ToplevelHandleState<W> {
+        ToplevelHandleState::new(ToplevelHandleStateInner {
+            outputs: Vec::new(),
+            geometry: None,
+            wl_outputs: HashSet::new(),
+            workspaces: Vec::new(),
+            title: String::new(),
+            app_id: String::new(),
+            states: Vec::new(),
+            window: None,
         })
     }
 }
@@ -187,6 +200,10 @@ where
                         .instances
                         .push(instance);
                 } else {
+                    let _ = data_init.init(
+                        cosmic_toplevel,
+                        ToplevelHandleStateInner::empty(),
+                    );
                     error!(?foreign_toplevel, "Toplevel for foreign-toplevel-list not registered for cosmic-toplevel-info.");
                 }
             }
@@ -509,10 +526,14 @@ where
         }
         handle_state.states = states.clone();
 
-        let states = states
-            .iter()
-            .flat_map(|state| (*state as u32).to_ne_bytes())
-            .collect::<Vec<u8>>();
+        let states: Vec<u8> = {
+            let ratio = std::mem::size_of::<States>() / std::mem::size_of::<u8>();
+            let ptr = states.as_mut_ptr() as *mut u8;
+            let len = states.len() * ratio;
+            let cap = states.capacity() * ratio;
+            std::mem::forget(states);
+            unsafe { Vec::from_raw_parts(ptr, len, cap) }
+        };
         instance.state(states);
         changed = true;
     }
@@ -600,7 +621,7 @@ where
 pub fn window_from_handle<W: Window + 'static>(handle: ZcosmicToplevelHandleV1) -> Option<W> {
     handle
         .data::<ToplevelHandleState<W>>()
-        .map(|state| state.lock().unwrap().window.clone())
+        .and_then(|state| state.lock().unwrap().window.clone())
 }
 
 macro_rules! delegate_toplevel_info {