scene_graph: Port ext_session_v1

12 files changed, 276 insertions, 188 deletions

diff --git a/include/sway/output.h b/include/sway/output.h
index 28240819..8405f78d 100644
--- a/include/sway/output.h
+++ b/include/sway/output.h
@@ -24,6 +24,7 @@ struct sway_output {
 	struct {
 		struct wlr_scene_tree *tiling;
 		struct wlr_scene_tree *fullscreen;
+		struct wlr_scene_tree *session_lock;
 	} layers;
 
 	// when a container is fullscreen, in case the fullscreen surface is
diff --git a/include/sway/server.h b/include/sway/server.h
index f3d25980..33ffbf09 100644
--- a/include/sway/server.h
+++ b/include/sway/server.h
@@ -28,6 +28,19 @@
 
 struct sway_transaction;
 
+struct sway_session_lock {
+	struct wlr_session_lock_v1 *lock;
+	struct wlr_surface *focused;
+	bool abandoned;
+
+	struct wl_list outputs; // struct sway_session_lock_output
+
+	// invalid if the session is abandoned
+	struct wl_listener new_surface;
+	struct wl_listener unlock;
+	struct wl_listener destroy;
+};
+
 struct sway_server {
 	struct wl_display *wl_display;
 	struct wl_event_loop *wl_event_loop;
@@ -92,15 +105,9 @@ struct sway_server {
 	struct wl_listener gamma_control_set_gamma;
 
 	struct {
-		bool locked;
+		struct sway_session_lock *lock;
 		struct wlr_session_lock_manager_v1 *manager;
 
-		struct wlr_session_lock_v1 *lock;
-		struct wlr_surface *focused;
-		struct wl_listener lock_new_surface;
-		struct wl_listener lock_unlock;
-		struct wl_listener lock_destroy;
-
 		struct wl_listener new_lock;
 		struct wl_listener manager_destroy;
 	} session_lock;
@@ -174,6 +181,10 @@ void handle_new_output(struct wl_listener *listener, void *data);
 void handle_idle_inhibitor_v1(struct wl_listener *listener, void *data);
 void handle_layer_shell_surface(struct wl_listener *listener, void *data);
 void sway_session_lock_init(void);
+void sway_session_lock_add_output(struct sway_session_lock *lock,
+	struct sway_output *output);
+bool sway_session_lock_has_surface(struct sway_session_lock *lock,
+	struct wlr_surface *surface);
 void handle_xdg_shell_toplevel(struct wl_listener *listener, void *data);
 #if HAVE_XWAYLAND
 void handle_xwayland_surface(struct wl_listener *listener, void *data);
diff --git a/include/sway/tree/root.h b/include/sway/tree/root.h
index 4b48a651..0aae8938 100644
--- a/include/sway/tree/root.h
+++ b/include/sway/tree/root.h
@@ -41,6 +41,7 @@ struct sway_root {
 		struct wlr_scene_tree *fullscreen;
 		struct wlr_scene_tree *fullscreen_global;
 		struct wlr_scene_tree *seat;
+		struct wlr_scene_tree *session_lock;
 	} layers;
 
 #if HAVE_XWAYLAND
diff --git a/sway/desktop/output.c b/sway/desktop/output.c
index 321e2a72..1e4474ce 100644
--- a/sway/desktop/output.c
+++ b/sway/desktop/output.c
@@ -798,6 +798,10 @@ void handle_new_output(struct wl_listener *listener, void *data) {
 	output->repaint_timer = wl_event_loop_add_timer(server->wl_event_loop,
 		output_repaint_timer_handler, output);
 
+	if (server->session_lock.lock) {
+		sway_session_lock_add_output(server->session_lock.lock, output);
+	}
+
 	struct output_config *oc = find_output_config(output);
 	apply_output_config(oc, output);
 	free_output_config(oc);
diff --git a/sway/desktop/render.c b/sway/desktop/render.c
index 735dddb8..60431d79 100644
--- a/sway/desktop/render.c
+++ b/sway/desktop/render.c
@@ -1012,43 +1012,6 @@ void output_render(struct render_context *ctx) {
 	pixman_region32_copy(&transformed_damage, damage);
 	transform_output_damage(&transformed_damage, wlr_output);
 
-	if (server.session_lock.locked) {
-		struct wlr_render_color clear_color = {
-			.a = 1.0f
-		};
-		if (server.session_lock.lock == NULL) {
-			// abandoned lock -> red BG
-			clear_color.r = 1.f;
-		}
-
-		wlr_render_pass_add_rect(ctx->pass, &(struct wlr_render_rect_options){
-			.box = { .width = wlr_output->width, .height = wlr_output->height },
-			.color = clear_color,
-			.clip = &transformed_damage,
-		});
-
-		if (server.session_lock.lock != NULL) {
-			struct render_data data = {
-				.alpha = 1.0f,
-				.ctx = ctx,
-			};
-
-			struct wlr_session_lock_surface_v1 *lock_surface;
-			wl_list_for_each(lock_surface, &server.session_lock.lock->surfaces, link) {
-				if (lock_surface->output != wlr_output) {
-					continue;
-				}
-				if (!lock_surface->surface->mapped) {
-					continue;
-				}
-
-				output_surface_for_each_surface(output, lock_surface->surface,
-					0.0, 0.0, render_surface_iterator, &data);
-			}
-		}
-		goto renderer_end;
-	}
-
 	if (output_has_opaque_overlay_layer_surface(output)) {
 		goto render_overlay;
 	}
@@ -1122,8 +1085,6 @@ render_overlay:
 		&output->shell_layers[ZWLR_LAYER_SHELL_V1_LAYER_OVERLAY]);
 	render_layer_popups(ctx,
 		&output->shell_layers[ZWLR_LAYER_SHELL_V1_LAYER_OVERLAY]);
-
-renderer_end:
 	pixman_region32_fini(&transformed_damage);
 	wlr_output_add_software_cursors_to_render_pass(wlr_output, ctx->pass, damage);
 }
diff --git a/sway/input/cursor.c b/sway/input/cursor.c
index e8604193..107424c9 100644
--- a/sway/input/cursor.c
+++ b/sway/input/cursor.c
@@ -95,24 +95,6 @@ struct sway_node *node_at_coords(
 	double ox = lx, oy = ly;
 	wlr_output_layout_output_coords(root->output_layout, wlr_output, &ox, &oy);
 
-	if (server.session_lock.locked) {
-		if (server.session_lock.lock == NULL) {
-			return NULL;
-		}
-		struct wlr_session_lock_surface_v1 *lock_surf;
-		wl_list_for_each(lock_surf, &server.session_lock.lock->surfaces, link) {
-			if (lock_surf->output != wlr_output) {
-				continue;
-			}
-
-			*surface = wlr_surface_surface_at(lock_surf->surface, ox, oy, sx, sy);
-			if (*surface != NULL) {
-				return NULL;
-			}
-		}
-		return NULL;
-	}
-
 	// layer surfaces on the overlay layer are rendered on top
 	if ((*surface = layer_surface_at(output,
 				&output->shell_layers[ZWLR_LAYER_SHELL_V1_LAYER_OVERLAY],
diff --git a/sway/input/keyboard.c b/sway/input/keyboard.c
index dea7a7bd..b97f0152 100644
--- a/sway/input/keyboard.c
+++ b/sway/input/keyboard.c
@@ -405,7 +405,7 @@ static void handle_key_event(struct sway_keyboard *keyboard,
 	char *device_identifier = input_device_get_identifier(wlr_device);
 	bool exact_identifier = keyboard->wlr->group != NULL;
 	seat_idle_notify_activity(seat, IDLE_SOURCE_KEYBOARD);
-	bool locked = server.session_lock.locked;
+	bool locked = server.session_lock.lock;
 	struct sway_keyboard_shortcuts_inhibitor *sway_inhibitor =
 		keyboard_shortcuts_inhibitor_get_for_focused_surface(seat);
 	bool shortcuts_inhibited = sway_inhibitor && sway_inhibitor->inhibitor->active;
diff --git a/sway/input/seat.c b/sway/input/seat.c
index d2cd1ba4..b8daa297 100644
--- a/sway/input/seat.c
+++ b/sway/input/seat.c
@@ -1060,19 +1060,10 @@ void seat_configure_xcursor(struct sway_seat *seat) {
 
 bool seat_is_input_allowed(struct sway_seat *seat,
 		struct wlr_surface *surface) {
-	if (!server.session_lock.locked) {
-		return true;
-	}
-	if (server.session_lock.lock == NULL) {
-		return false;
-	}
-	struct wlr_session_lock_surface_v1 *lock_surf;
-	wl_list_for_each(lock_surf, &server.session_lock.lock->surfaces, link) {
-		if (lock_surf->surface == surface) {
-			return true;
-		}
+	if (server.session_lock.lock) {
+		return sway_session_lock_has_surface(server.session_lock.lock, surface);
 	}
-	return false;
+	return true;
 }
 
 static void send_unfocus(struct sway_container *con, void *data) {
@@ -1277,8 +1268,8 @@ void seat_set_focus(struct sway_seat *seat, struct sway_node *node) {
 	} else {
 		seat_set_workspace_focus(seat, node);
 	}
-	if (server.session_lock.locked) {
-		seat_set_focus_surface(seat, server.session_lock.focused, false);
+	if (server.session_lock.lock) {
+		seat_set_focus_surface(seat, server.session_lock.lock->focused, false);
 	}
 }
 
diff --git a/sway/input/switch.c b/sway/input/switch.c
index f483cd23..831f4dbf 100644
--- a/sway/input/switch.c
+++ b/sway/input/switch.c
@@ -34,7 +34,7 @@ static bool sway_switch_trigger_test(enum sway_switch_trigger trigger,
 
 static void execute_binding(struct sway_switch *sway_switch) {
 	struct sway_seat *seat = sway_switch->seat_device->sway_seat;
-	bool locked = server.session_lock.locked;
+	bool locked = server.session_lock.lock;
 
 	list_t *bindings = config->current_mode->switch_bindings;
 	struct sway_switch_binding *matched_binding = NULL;
diff --git a/sway/lock.c b/sway/lock.c
index c4fbfe0e..2856ae67 100644
--- a/sway/lock.c
+++ b/sway/lock.c
@@ -1,5 +1,6 @@
 #define _POSIX_C_SOURCE 200809L
 #include <assert.h>
+#include <wlr/types/wlr_scene.h>
 #include "log.h"
 #include "sway/input/cursor.h"
 #include "sway/input/keyboard.h"
@@ -9,19 +10,28 @@
 #include "sway/server.h"
 #include "sway/surface.h"
 
-struct sway_session_lock_surface {
-	struct wlr_session_lock_surface_v1 *lock_surface;
+struct sway_session_lock_output {
+	struct wlr_scene_tree *tree;
+	struct wlr_scene_rect *background;
+	struct sway_session_lock *lock;
+
 	struct sway_output *output;
-	struct wlr_surface *surface;
-	struct wl_listener map;
+
+	struct wl_list link; // sway_session_lock::outputs
+
 	struct wl_listener destroy;
-	struct wl_listener surface_commit;
-	struct wl_listener output_commit;
-	struct wl_listener output_destroy;
+	struct wl_listener commit;
+
+	struct wlr_session_lock_surface_v1 *surface;
+
+	// invalid if surface is NULL
+	struct wl_listener surface_destroy;
+	struct wl_listener surface_map;
 };
 
-static void set_lock_focused_surface(struct wlr_surface *focused) {
-	server.session_lock.focused = focused;
+static void focus_surface(struct sway_session_lock *lock,
+		struct wlr_surface *focused) {
+	lock->focused = focused;
 
 	struct sway_seat *seat;
 	wl_list_for_each(seat, &server.input->seats, link) {
@@ -29,104 +39,189 @@ static void set_lock_focused_surface(struct wlr_surface *focused) {
 	}
 }
 
+static void refocus_output(struct sway_session_lock_output *output) {
+	// Move the seat focus to another surface if one is available
+	if (output->lock->focused == output->surface->surface) {
+		struct wlr_surface *next_focus = NULL;
+
+		struct sway_session_lock_output *candidate;
+		wl_list_for_each(candidate, &output->lock->outputs, link) {
+			if (candidate == output || !candidate->surface) {
+				continue;
+			}
+
+			if (candidate->surface->surface->mapped) {
+				next_focus = candidate->surface->surface;
+				break;
+			}
+		}
+
+		focus_surface(output->lock, next_focus);
+	}
+}
+
 static void handle_surface_map(struct wl_listener *listener, void *data) {
-	struct sway_session_lock_surface *surf = wl_container_of(listener, surf, map);
-	if (server.session_lock.focused == NULL) {
-		set_lock_focused_surface(surf->surface);
+	struct sway_session_lock_output *surf = wl_container_of(listener, surf, surface_map);
+	if (surf->lock->focused == NULL) {
+		focus_surface(surf->lock, surf->surface->surface);
 	}
 	cursor_rebase_all();
-	surface_enter_output(surf->surface, surf->output);
-	output_damage_whole(surf->output);
 }
 
-static void handle_surface_commit(struct wl_listener *listener, void *data) {
-	struct sway_session_lock_surface *surf = wl_container_of(listener, surf, surface_commit);
-	output_damage_surface(surf->output, 0, 0, surf->surface, false);
+static void handle_surface_destroy(struct wl_listener *listener, void *data) {
+	struct sway_session_lock_output *output =
+		wl_container_of(listener, output, surface_destroy);
+	refocus_output(output);
+
+	sway_assert(output->surface, "Trying to destroy a surface that the lock doesn't think exists");
+	output->surface = NULL;
+	wl_list_remove(&output->surface_destroy.link);
+	wl_list_remove(&output->surface_map.link);
+}
+
+static void lock_output_reconfigure(struct sway_session_lock_output *output) {
+	int width = output->output->width;
+	int height = output->output->height;
+
+	wlr_scene_rect_set_size(output->background, width, height);
+
+	if (output->surface) {
+		wlr_session_lock_surface_v1_configure(output->surface, width, height);
+	}
 }
 
-static void handle_output_commit(struct wl_listener *listener, void *data) {
+static void handle_new_surface(struct wl_listener *listener, void *data) {
+	struct sway_session_lock *lock = wl_container_of(listener, lock, new_surface);
+	struct wlr_session_lock_surface_v1 *lock_surface = data;
+	struct sway_output *output = lock_surface->output->data;
+
+	sway_log(SWAY_DEBUG, "new lock layer surface");
+
+	struct sway_session_lock_output *current_lock_output, *lock_output = NULL;
+	wl_list_for_each(current_lock_output, &lock->outputs, link) {
+		if (current_lock_output->output == output) {
+			lock_output = current_lock_output;
+			break;
+		}
+	}
+	sway_assert(lock_output, "Couldn't find output to lock");
+	sway_assert(!lock_output->surface, "Tried to reassign a surface to an existing output");
+
+	lock_output->surface = lock_surface;
+
+	wlr_scene_subsurface_tree_create(lock_output->tree, lock_surface->surface);
+
+	lock_output->surface_destroy.notify = handle_surface_destroy;
+	wl_signal_add(&lock_surface->events.destroy, &lock_output->surface_destroy);
+	lock_output->surface_map.notify = handle_surface_map;
+	wl_signal_add(&lock_surface->surface->events.map, &lock_output->surface_map);
+
+	lock_output_reconfigure(lock_output);
+}
+
+static void sway_session_lock_output_destroy(struct sway_session_lock_output *output) {
+	if (output->surface) {
+		refocus_output(output);
+		wl_list_remove(&output->surface_destroy.link);
+		wl_list_remove(&output->surface_map.link);
+	}
+
+	wl_list_remove(&output->commit.link);
+	wl_list_remove(&output->destroy.link);
+	wl_list_remove(&output->link);
+
+	free(output);
+}
+
+static void lock_node_handle_destroy(struct wl_listener *listener, void *data) {
+	struct sway_session_lock_output *output =
+		wl_container_of(listener, output, destroy);
+	sway_session_lock_output_destroy(output);
+}
+
+static void lock_output_handle_commit(struct wl_listener *listener, void *data) {
 	struct wlr_output_event_commit *event = data;
-	struct sway_session_lock_surface *surf = wl_container_of(listener, surf, output_commit);
+	struct sway_session_lock_output *output =
+		wl_container_of(listener, output, commit);
 	if (event->state->committed & (
 			WLR_OUTPUT_STATE_MODE |
 			WLR_OUTPUT_STATE_SCALE |
 			WLR_OUTPUT_STATE_TRANSFORM)) {
-		wlr_session_lock_surface_v1_configure(surf->lock_surface,
-			surf->output->width, surf->output->height);
+		lock_output_reconfigure(output);
 	}
 }
 
-static void destroy_lock_surface(struct sway_session_lock_surface *surf) {
-	// Move the seat focus to another surface if one is available
-	if (server.session_lock.focused == surf->surface) {
-		struct wlr_surface *next_focus = NULL;
+static struct sway_session_lock_output *session_lock_output_create(
+		struct sway_session_lock *lock, struct sway_output *output) {
+	struct sway_session_lock_output *lock_output = calloc(1, sizeof(*lock_output));
+	if (!lock_output) {
+		sway_log(SWAY_ERROR, "failed to allocate a session lock output");
+		return NULL;
+	}
 
-		struct wlr_session_lock_surface_v1 *other;
-		wl_list_for_each(other, &server.session_lock.lock->surfaces, link) {
-			if (other != surf->lock_surface && other->surface->mapped) {
-				next_focus = other->surface;
-				break;
-			}
-		}
-		set_lock_focused_surface(next_focus);
+	struct wlr_scene_tree *tree = wlr_scene_tree_create(output->layers.session_lock);
+	if (!tree) {
+		sway_log(SWAY_ERROR, "failed to allocate a session lock output scene tree");
+		free(lock_output);
+		return NULL;
 	}
 
-	wl_list_remove(&surf->map.link);
-	wl_list_remove(&surf->destroy.link);
-	wl_list_remove(&surf->surface_commit.link);
-	wl_list_remove(&surf->output_commit.link);
-	wl_list_remove(&surf->output_destroy.link);
-	output_damage_whole(surf->output);
-	free(surf);
-}
+	struct wlr_scene_rect *background = wlr_scene_rect_create(tree, 0, 0, (float[4]){
+		lock->abandoned ? 1.f : 0.f,
+		0.f,
+		0.f,
+		1.f,
+	});
+	if (!background) {
+		sway_log(SWAY_ERROR, "failed to allocate a session lock output scene background");
+		wlr_scene_node_destroy(&tree->node);
+		free(lock_output);
+		return NULL;
+	}
 
-static void handle_surface_destroy(struct wl_listener *listener, void *data) {
-	struct sway_session_lock_surface *surf = wl_container_of(listener, surf, destroy);
-	destroy_lock_surface(surf);
-}
+	lock_output->output = output;
+	lock_output->tree = tree;
+	lock_output->background = background;
+	lock_output->lock = lock;
+
+	lock_output->destroy.notify = lock_node_handle_destroy;
+	wl_signal_add(&tree->node.events.destroy, &lock_output->destroy);
 
-static void handle_output_destroy(struct wl_listener *listener, void *data) {
-	struct sway_session_lock_surface *surf =
-		wl_container_of(listener, surf, output_destroy);
-	destroy_lock_surface(surf);
+	lock_output->commit.notify = lock_output_handle_commit;
+	wl_signal_add(&output->wlr_output->events.commit, &lock_output->commit);
+
+	lock_output_reconfigure(lock_output);
+
+	wl_list_insert(&lock->outputs, &lock_output->link);
+
+	return lock_output;
 }
 
-static void handle_new_surface(struct wl_listener *listener, void *data) {
-	struct wlr_session_lock_surface_v1 *lock_surface = data;
-	struct sway_session_lock_surface *surf = calloc(1, sizeof(*surf));
-	if (surf == NULL) {
-		return;
+static void sway_session_lock_destroy(struct sway_session_lock* lock) {
+	struct sway_session_lock_output *lock_output, *tmp_lock_output;
+	wl_list_for_each_safe(lock_output, tmp_lock_output, &lock->outputs, link) {
+		// destroying the node will also destroy the whole lock output
+		wlr_scene_node_destroy(&lock_output->tree->node);
 	}
 
-	sway_log(SWAY_DEBUG, "new lock layer surface");
+	if (server.session_lock.lock == lock) {
+		server.session_lock.lock = NULL;
+	}
 
-	struct sway_output *output = lock_surface->output->data;
-	wlr_session_lock_surface_v1_configure(lock_surface, output->width, output->height);
-
-	surf->lock_surface = lock_surface;
-	surf->surface = lock_surface->surface;
-	surf->output = output;
-	surf->map.notify = handle_surface_map;
-	wl_signal_add(&lock_surface->surface->events.map, &surf->map);
-	surf->destroy.notify = handle_surface_destroy;
-	wl_signal_add(&lock_surface->events.destroy, &surf->destroy);
-	surf->surface_commit.notify = handle_surface_commit;
-	wl_signal_add(&surf->surface->events.commit, &surf->surface_commit);
-	surf->output_commit.notify = handle_output_commit;
-	wl_signal_add(&output->wlr_output->events.commit, &surf->output_commit);
-	surf->output_destroy.notify = handle_output_destroy;
-	wl_signal_add(&output->node.events.destroy, &surf->output_destroy);
+	if (!lock->abandoned) {
+		wl_list_remove(&lock->destroy.link);
+		wl_list_remove(&lock->unlock.link);
+		wl_list_remove(&lock->new_surface.link);
+	}
+
+	free(lock);
 }
 
 static void handle_unlock(struct wl_listener *listener, void *data) {
+	struct sway_session_lock *lock = wl_container_of(listener, lock, unlock);
 	sway_log(SWAY_DEBUG, "session unlocked");
-	server.session_lock.locked = false;
-	server.session_lock.lock = NULL;
-	server.session_lock.focused = NULL;
 
-	wl_list_remove(&server.session_lock.lock_new_surface.link);
-	wl_list_remove(&server.session_lock.lock_unlock.link);
-	wl_list_remove(&server.session_lock.lock_destroy.link);
+	sway_session_lock_destroy(lock);
 
 	struct sway_seat *seat;
 	wl_list_for_each(seat, &server.input->seats, link) {
@@ -145,28 +240,22 @@ static void handle_unlock(struct wl_listener *listener, void *data) {
 		struct sway_output *output = root->outputs->items[i];
 		arrange_layers(output);
 	}
-
-	// redraw everything
-	for (int i = 0; i < root->outputs->length; ++i) {
-		struct sway_output *output = root->outputs->items[i];
-		output_damage_whole(output);
-	}
 }
 
 static void handle_abandon(struct wl_listener *listener, void *data) {
+	struct sway_session_lock *lock = wl_container_of(listener, lock, destroy);
 	sway_log(SWAY_INFO, "session lock abandoned");
-	server.session_lock.lock = NULL;
-	server.session_lock.focused = NULL;
 
-	wl_list_remove(&server.session_lock.lock_new_surface.link);
-	wl_list_remove(&server.session_lock.lock_unlock.link);
-	wl_list_remove(&server.session_lock.lock_destroy.link);
-
-	// redraw everything
-	for (int i = 0; i < root->outputs->length; ++i) {
-		struct sway_output *output = root->outputs->items[i];
-		output_damage_whole(output);
+	struct sway_session_lock_output *lock_output;
+	wl_list_for_each(lock_output, &lock->outputs, link) {
+		wlr_scene_rect_set_color(lock_output->background,
+			(float[4]){ 1.f, 0.f, 0.f, 1.f });
 	}
+
+	lock->abandoned = true;
+	wl_list_remove(&lock->destroy.link);
+	wl_list_remove(&lock->unlock.link);
+	wl_list_remove(&lock->new_surface.link);
 }
 
 static void handle_session_lock(struct wl_listener *listener, void *data) {
@@ -174,44 +263,89 @@ static void handle_session_lock(struct wl_listener *listener, void *data) {
 	struct wl_client *client = wl_resource_get_client(lock->resource);
 
 	if (server.session_lock.lock) {
+		if (server.session_lock.lock->abandoned) {
+			sway_log(SWAY_INFO, "Replacing abandoned lock");
+			sway_session_lock_destroy(server.session_lock.lock);
+		} else {
+			sway_log(SWAY_ERROR, "Cannot lock an already locked session");
+			wlr_session_lock_v1_destroy(lock);
+			return;
+		}
+	}
+
+	struct sway_session_lock *sway_lock = calloc(1, sizeof(*sway_lock));
+	if (!sway_lock) {
+		sway_log(SWAY_ERROR, "failed to allocate a session lock object");
 		wlr_session_lock_v1_destroy(lock);
 		return;
 	}
 
+	wl_list_init(&sway_lock->outputs);
+
 	sway_log(SWAY_DEBUG, "session locked");
-	server.session_lock.locked = true;
-	server.session_lock.lock = lock;
 
 	struct sway_seat *seat;
 	wl_list_for_each(seat, &server.input->seats, link) {
 		seat_unfocus_unless_client(seat, client);
 	}
 
-	wl_signal_add(&lock->events.new_surface, &server.session_lock.lock_new_surface);
-	wl_signal_add(&lock->events.unlock, &server.session_lock.lock_unlock);
-	wl_signal_add(&lock->events.destroy, &server.session_lock.lock_destroy);
+	struct sway_output *output;
+	wl_list_for_each(output, &root->all_outputs, link) {
+		sway_session_lock_add_output(sway_lock, output);
+	}
+
+	sway_lock->new_surface.notify = handle_new_surface;
+	wl_signal_add(&lock->events.new_surface, &sway_lock->new_surface);
+	sway_lock->unlock.notify = handle_unlock;
+	wl_signal_add(&lock->events.unlock, &sway_lock->unlock);
+	sway_lock->destroy.notify = handle_abandon;
+	wl_signal_add(&lock->events.destroy, &sway_lock->destroy);
 
 	wlr_session_lock_v1_send_locked(lock);
-
-	// redraw everything
-	for (int i = 0; i < root->outputs->length; ++i) {
-		struct sway_output *output = root->outputs->items[i];
-		output_damage_whole(output);
-	}
+	server.session_lock.lock = sway_lock;
 }
 
 static void handle_session_lock_destroy(struct wl_listener *listener, void *data) {
-	assert(server.session_lock.lock == NULL);
+	// if the server shuts down while a lock is active, destroy the lock
+	if (server.session_lock.lock) {
+		sway_session_lock_destroy(server.session_lock.lock);
+	}
+
 	wl_list_remove(&server.session_lock.new_lock.link);
 	wl_list_remove(&server.session_lock.manager_destroy.link);
+
+	server.session_lock.manager = NULL;
+}
+
+void sway_session_lock_add_output(struct sway_session_lock *lock,
+		struct sway_output *output) {
+	struct sway_session_lock_output *lock_output =
+		session_lock_output_create(lock, output);
+
+	// if we run out of memory while trying to lock the screen, the best we
+	// can do is kill the sway process. Security conscious users will have
+	// the sway session fall back to a login shell.
+	if (!lock_output) {
+		sway_log(SWAY_ERROR, "aborting: failed to allocate a lock output");
+		abort();
+	}
+}
+
+bool sway_session_lock_has_surface(struct sway_session_lock *lock,
+		struct wlr_surface *surface) {
+	struct sway_session_lock_output *lock_output;
+	wl_list_for_each(lock_output, &lock->outputs, link) {
+		if (lock_output->surface && lock_output->surface->surface == surface) {
+			return true;
+		}
+	}
+
+	return false;
 }
 
 void sway_session_lock_init(void) {
 	server.session_lock.manager = wlr_session_lock_manager_v1_create(server.wl_display);
 
-	server.session_lock.lock_new_surface.notify = handle_new_surface;
-	server.session_lock.lock_unlock.notify = handle_unlock;
-	server.session_lock.lock_destroy.notify = handle_abandon;
 	server.session_lock.new_lock.notify = handle_session_lock;
 	server.session_lock.manager_destroy.notify = handle_session_lock_destroy;
 	wl_signal_add(&server.session_lock.manager->events.new_lock,
diff --git a/sway/tree/output.c b/sway/tree/output.c
index 12a2f969..64ca3d75 100644
--- a/sway/tree/output.c
+++ b/sway/tree/output.c
@@ -95,6 +95,7 @@ static void destroy_scene_layers(struct sway_output *output) {
 
 	wlr_scene_node_destroy(&output->layers.tiling->node);
 	wlr_scene_node_destroy(&output->layers.fullscreen->node);
+	wlr_scene_node_destroy(&output->layers.session_lock->node);
 }
 
 struct sway_output *output_create(struct wlr_output *wlr_output) {
@@ -104,6 +105,7 @@ struct sway_output *output_create(struct wlr_output *wlr_output) {
 	bool failed = false;
 	output->layers.tiling = alloc_scene_tree(root->staging, &failed);
 	output->layers.fullscreen = alloc_scene_tree(root->staging, &failed);
+	output->layers.session_lock = alloc_scene_tree(root->staging, &failed);
 
 	if (!failed) {
 		output->fullscreen_background = wlr_scene_rect_create(
diff --git a/sway/tree/root.c b/sway/tree/root.c
index 38fcdb7c..e4941566 100644
--- a/sway/tree/root.c
+++ b/sway/tree/root.c
@@ -50,6 +50,7 @@ struct sway_root *root_create(struct wl_display *wl_display) {
 	root->layers.fullscreen = alloc_scene_tree(&root_scene->tree, &failed);
 	root->layers.fullscreen_global = alloc_scene_tree(&root_scene->tree, &failed);
 	root->layers.seat = alloc_scene_tree(&root_scene->tree, &failed);
+	root->layers.session_lock = alloc_scene_tree(&root_scene->tree, &failed);
 
 	if (failed) {
 		wlr_scene_node_destroy(&root_scene->tree.node);