Fix focus tracking when session lock is active

Remove the incorrect attempt to block focus changes when an input grab
is present and replace it with the same logic used for layer_shell-based
screen lockers: restore the focus after changing it.

This fixes a use-after-free of seat->workspace if outputs are destroyed
while a screen lock is enabled.

sway/input/seat.c

@@ -1134,15 +1134,7 @@ void seat_set_raw_focus(struct sway_seat *seat, struct sway_node *node) {
 	}
 }
 
-void seat_set_focus(struct sway_seat *seat, struct sway_node *node) {
+static void seat_set_workspace_focus(struct sway_seat *seat, struct sway_node *node) {
-	if (seat->focused_layer) {
-		struct wlr_layer_surface_v1 *layer = seat->focused_layer;
-		seat_set_focus_layer(seat, NULL);
-		seat_set_focus(seat, node);
-		seat_set_focus_layer(seat, layer);
-		return;
-	}
-
 	struct sway_node *last_focus = seat_get_focus(seat);
 	if (last_focus == node) {
 		return;
@@ -1176,11 +1168,6 @@ void seat_set_focus(struct sway_seat *seat, struct sway_node *node) {
 		return;
 	}
 
-	// Deny setting focus when an input grab or lockscreen is active
-	if (container && container->view && !seat_is_input_allowed(seat, container->view->surface)) {
-		return;
-	}
-
 	struct sway_output *new_output =
 		new_workspace ? new_workspace->output : NULL;
 
@@ -1280,6 +1267,20 @@ void seat_set_focus(struct sway_seat *seat, struct sway_node *node) {
 	}
 }
 
+void seat_set_focus(struct sway_seat *seat, struct sway_node *node) {
+	if (seat->focused_layer) {
+		struct wlr_layer_surface_v1 *layer = seat->focused_layer;
+		seat_set_focus_layer(seat, NULL);
+		seat_set_workspace_focus(seat, node);
+		seat_set_focus_layer(seat, layer);
+	} else {
+		seat_set_workspace_focus(seat, node);
+	}
+	if (server.session_lock.locked) {
+		seat_set_focus_surface(seat, server.session_lock.focused, false);
+	}
+}
+
 void seat_set_focus_container(struct sway_seat *seat,
 		struct sway_container *con) {
 	seat_set_focus(seat, con ? &con->node : NULL);