From 959f32009b2b0cae587987236e45cab6167d0183 Mon Sep 17 00:00:00 2001
From: Nova <technobaboo@gmail.com>
Date: Wed, 9 Nov 2022 11:02:48 -0500
Subject: [PATCH] fix(wayland): account for surface data map panic

---
 src/wayland/compositor.rs |  2 ++
 src/wayland/mod.rs        | 10 ++++++++--
 src/wayland/state.rs      |  2 ++
 src/wayland/surface.rs    | 15 ++++++++++-----
 4 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/wayland/compositor.rs b/src/wayland/compositor.rs
index 66bc543..7a7d304 100644
--- a/src/wayland/compositor.rs
+++ b/src/wayland/compositor.rs
@@ -1,5 +1,6 @@
 use super::{state::WaylandState, surface::CoreSurface};
 use smithay::{
+	backend::renderer::utils::on_commit_buffer_handler,
 	delegate_compositor,
 	reexports::wayland_server::protocol::wl_surface::WlSurface,
 	wayland::compositor::{self, CompositorHandler, CompositorState},
@@ -11,6 +12,7 @@ impl CompositorHandler for WaylandState {
 	}
 
 	fn commit(&mut self, surface: &WlSurface) {
+		on_commit_buffer_handler(surface);
 		compositor::with_states(surface, |data| {
 			data.data_map.insert_if_missing_threadsafe(|| {
 				CoreSurface::new(
diff --git a/src/wayland/mod.rs b/src/wayland/mod.rs
index b9bd039..1080302 100644
--- a/src/wayland/mod.rs
+++ b/src/wayland/mod.rs
@@ -146,9 +146,15 @@ impl Wayland {
 		let time_ms = (sk.time_getf() * 1000.) as u64;
 
 		for core_surface in CORE_SURFACES.get_valid_contents() {
-			let client_id = core_surface.wl_surface().client().unwrap().id();
+			let client_id = match core_surface.wl_surface().client() {
+				Some(client) => client.id(),
+				None => continue,
+			};
 			let state = self.state.lock();
-			let seat_data = state.seats.get(&client_id).unwrap().clone();
+			let seat_data = match state.seats.get(&client_id) {
+				Some(seat_data) => seat_data.clone(),
+				None => continue,
+			};
 			let output = state.output.clone();
 			core_surface.process(
 				sk,
diff --git a/src/wayland/state.rs b/src/wayland/state.rs
index 3b9e767..40decbc 100644
--- a/src/wayland/state.rs
+++ b/src/wayland/state.rs
@@ -44,6 +44,7 @@ impl ClientData for ClientState {
 
 pub struct WaylandState {
 	pub weak_ref: Weak<Mutex<WaylandState>>,
+	pub log: Logger,
 	pub display: Arc<Mutex<Display<WaylandState>>>,
 	pub display_handle: DisplayHandle,
 
@@ -102,6 +103,7 @@ impl WaylandState {
 		Arc::new_cyclic(|weak| {
 			Mutex::new(WaylandState {
 				weak_ref: weak.clone(),
+				log,
 				display,
 				display_handle,
 
diff --git a/src/wayland/surface.rs b/src/wayland/surface.rs
index 98e1ab8..37ed773 100644
--- a/src/wayland/surface.rs
+++ b/src/wayland/surface.rs
@@ -10,7 +10,7 @@ use slog::Logger;
 use smithay::{
 	backend::renderer::{
 		gles2::{Gles2Renderer, Gles2Texture},
-		utils::{import_surface_tree, on_commit_buffer_handler, RendererSurfaceStateUserData},
+		utils::{import_surface_tree, RendererSurfaceStateUserData},
 		Texture,
 	},
 	desktop::utils::send_frames_surface_tree,
@@ -18,7 +18,7 @@ use smithay::{
 	reexports::wayland_server::{
 		backend::ObjectId, protocol::wl_surface::WlSurface, Display, DisplayHandle, Resource,
 	},
-	wayland::compositor::{self, SurfaceData},
+	wayland::compositor::{self, SurfaceData, SurfaceUserData},
 };
 use std::{
 	sync::{Arc, Weak},
@@ -128,10 +128,15 @@ impl CoreSurface {
 		on_mapped: F,
 		if_mapped: M,
 	) {
-		// Let Smithay handle all the buffer maintenance
-		on_commit_buffer_handler(&self.wl_surface());
+		// Avoid a panic in rare cases
+		if self.wl_surface().data::<SurfaceUserData>().is_none() {
+			return;
+		}
+
 		// Import all surface buffers into textures
-		import_surface_tree(renderer, &self.wl_surface(), log).unwrap();
+		if import_surface_tree(renderer, &self.wl_surface(), log).is_err() {
+			return;
+		}
 
 		let mapped = compositor::with_states(&self.wl_surface(), |data| {
 			data.data_map