3 files changed, 34 insertions, 19 deletions

diff --git a/include/wlr/types/wlr_server_decoration.h b/include/wlr/types/wlr_server_decoration.h
index 474a9386..e457a0aa 100644
--- a/include/wlr/types/wlr_server_decoration.h
+++ b/include/wlr/types/wlr_server_decoration.h
@@ -33,6 +33,8 @@ struct wlr_server_decoration_manager {
 
 	uint32_t default_mode; // enum wlr_server_decoration_manager_mode
 
+	struct wl_listener display_destroy;
+
 	struct {
 		struct wl_signal new_decoration;
 	} events;
diff --git a/types/wlr_screenshooter.c b/types/wlr_screenshooter.c
index d62820e6..cb831827 100644
--- a/types/wlr_screenshooter.c
+++ b/types/wlr_screenshooter.c
@@ -28,6 +28,7 @@ struct screenshot_state {
 
 static void screenshot_destroy(struct wlr_screenshot *screenshot) {
 	wl_list_remove(&screenshot->link);
+	wl_resource_set_user_data(screenshot->resource, NULL);
 	free(screenshot);
 }
 
@@ -35,7 +36,9 @@ static void handle_screenshot_resource_destroy(
 		struct wl_resource *screenshot_resource) {
 	struct wlr_screenshot *screenshot =
 		wl_resource_get_user_data(screenshot_resource);
-	screenshot_destroy(screenshot);
+	if (screenshot != NULL) {
+		screenshot_destroy(screenshot);
+	}
 }
 
 static void output_frame_notify(struct wl_listener *listener, void *_data) {
diff --git a/types/wlr_server_decoration.c b/types/wlr_server_decoration.c
index 60cc6369..3531c453 100644
--- a/types/wlr_server_decoration.c
+++ b/types/wlr_server_decoration.c
@@ -140,6 +140,30 @@ static void server_decoration_manager_bind(struct wl_client *client, void *data,
 		manager->default_mode);
 }
 
+void wlr_server_decoration_manager_destroy(
+		struct wlr_server_decoration_manager *manager) {
+	if (manager == NULL) {
+		return;
+	}
+	struct wlr_server_decoration *decoration, *tmp_decoration;
+	wl_list_for_each_safe(decoration, tmp_decoration, &manager->decorations,
+			link) {
+		server_decoration_destroy(decoration);
+	}
+	struct wl_resource *resource, *tmp_resource;
+	wl_resource_for_each_safe(resource, tmp_resource, &manager->wl_resources) {
+		server_decoration_manager_destroy_resource(resource);
+	}
+	wl_global_destroy(manager->wl_global);
+	free(manager);
+}
+
+static void handle_display_destroy(struct wl_listener *listener, void *data) {
+	struct wlr_server_decoration_manager *manager =
+		wl_container_of(listener, manager, display_destroy);
+	wlr_server_decoration_manager_destroy(manager);
+}
+
 struct wlr_server_decoration_manager *wlr_server_decoration_manager_create(
 		struct wl_display *display) {
 	struct wlr_server_decoration_manager *manager =
@@ -158,23 +182,9 @@ struct wlr_server_decoration_manager *wlr_server_decoration_manager_create(
 	wl_list_init(&manager->wl_resources);
 	wl_list_init(&manager->decorations);
 	wl_signal_init(&manager->events.new_decoration);
-	return manager;
-}
 
-void wlr_server_decoration_manager_destroy(
-		struct wlr_server_decoration_manager *manager) {
-	if (manager == NULL) {
-		return;
-	}
-	struct wlr_server_decoration *decoration, *tmp_decoration;
-	wl_list_for_each_safe(decoration, tmp_decoration, &manager->decorations,
-			link) {
-		server_decoration_destroy(decoration);
-	}
-	struct wl_resource *resource, *tmp_resource;
-	wl_resource_for_each_safe(resource, tmp_resource, &manager->wl_resources) {
-		server_decoration_manager_destroy_resource(resource);
-	}
-	wl_global_destroy(manager->wl_global);
-	free(manager);
+	manager->display_destroy.notify = handle_display_destroy;
+	wl_display_add_destroy_listener(display, &manager->display_destroy);
+
+	return manager;
 }