Implement input method keyboard grab

3 files changed, 102 insertions, 9 deletions

diff --git a/include/sway/input/text_input.h b/include/sway/input/text_input.h
index 6cf9bdb3a..37744266d 100644
--- a/include/sway/input/text_input.h
+++ b/include/sway/input/text_input.h
@@ -28,7 +28,10 @@ struct sway_input_method_relay {
 
 	struct wl_listener input_method_new;
 	struct wl_listener input_method_commit;
+	struct wl_listener input_method_grab_keyboard;
 	struct wl_listener input_method_destroy;
+
+	struct wl_listener input_method_keyboard_grab_destroy;
 };
 
 struct sway_text_input {
diff --git a/sway/input/keyboard.c b/sway/input/keyboard.c
index 95e53934f..f258ac7d5 100644
--- a/sway/input/keyboard.c
+++ b/sway/input/keyboard.c
@@ -378,6 +378,28 @@ static void update_keyboard_state(struct sway_keyboard *keyboard,
 	}
 }
 
+/**
+ * Get keyboard grab of the seat from sway_keyboard if we should forward events
+ * to it.
+ *
+ * Returns NULL if the keyboard is not grabbed by an input method,
+ * or if event is from virtual keyboard of the same client as grab.
+ * TODO: see swaywm/wlroots#2322
+ */
+static struct wlr_input_method_keyboard_grab_v2 *keyboard_get_im_grab(
+		struct sway_keyboard *keyboard) {
+	struct wlr_input_method_v2 *input_method = keyboard->seat_device->
+		sway_seat->im_relay.input_method;
+	struct wlr_virtual_keyboard_v1 *virtual_keyboard =
+		wlr_input_device_get_virtual_keyboard(keyboard->seat_device->input_device->wlr_device);
+	if (!input_method || !input_method->keyboard_grab || (virtual_keyboard &&
+				wl_resource_get_client(virtual_keyboard->resource) ==
+				wl_resource_get_client(input_method->keyboard_grab->resource))) {
+		return NULL;
+	}
+	return input_method->keyboard_grab;
+}
+
 static void handle_key_event(struct sway_keyboard *keyboard,
 		struct wlr_event_keyboard_key *event) {
 	struct sway_seat *seat = keyboard->seat_device->sway_seat;
@@ -488,17 +510,42 @@ static void handle_key_event(struct sway_keyboard *keyboard,
 				keyinfo.raw_keysyms_len);
 	}
 
-	if (!handled || event->state == WL_KEYBOARD_KEY_STATE_RELEASED) {
+	if (event->state == WL_KEYBOARD_KEY_STATE_RELEASED) {
+		// If the pressed event was sent to a client, also send the released
+		// event. In particular, don't send the released event to the IM grab.
 		bool pressed_sent = update_shortcut_state(
-				&keyboard->state_pressed_sent, event->keycode, event->state,
-				keyinfo.keycode, 0);
-		if (pressed_sent || event->state == WL_KEYBOARD_KEY_STATE_PRESSED) {
+			&keyboard->state_pressed_sent, event->keycode,
+			event->state, keyinfo.keycode, 0);
+		if (pressed_sent) {
 			wlr_seat_set_keyboard(wlr_seat, wlr_device);
 			wlr_seat_keyboard_notify_key(wlr_seat, event->time_msec,
-					event->keycode, event->state);
+				event->keycode, event->state);
+			handled = true;
+		}
+	}
+
+	if (!handled) {
+		struct wlr_input_method_keyboard_grab_v2 *kb_grab = keyboard_get_im_grab(keyboard);
+
+		if (kb_grab) {
+			wlr_input_method_keyboard_grab_v2_set_keyboard(kb_grab,
+				wlr_device->keyboard);
+			wlr_input_method_keyboard_grab_v2_send_key(kb_grab,
+				event->time_msec, event->keycode, event->state);
+			handled = true;
 		}
 	}
 
+	if (!handled && event->state != WL_KEYBOARD_KEY_STATE_RELEASED) {
+		// If a released event failed pressed sent test, and not in sent to
+		// keyboard grab, it is still not handled. Don't handle released here.
+		update_shortcut_state(
+			&keyboard->state_pressed_sent, event->keycode, event->state,
+			keyinfo.keycode, 0);
+		wlr_seat_set_keyboard(wlr_seat, wlr_device);
+		wlr_seat_keyboard_notify_key(wlr_seat, event->time_msec,
+				event->keycode, event->state);
+	}
 
 	free(device_identifier);
 }
@@ -614,10 +661,19 @@ static void handle_modifier_event(struct sway_keyboard *keyboard) {
 	struct wlr_input_device *wlr_device =
 		keyboard->seat_device->input_device->wlr_device;
 	if (!wlr_device->keyboard->group) {
-		struct wlr_seat *wlr_seat = keyboard->seat_device->sway_seat->wlr_seat;
-		wlr_seat_set_keyboard(wlr_seat, wlr_device);
-		wlr_seat_keyboard_notify_modifiers(wlr_seat,
-				&wlr_device->keyboard->modifiers);
+		struct wlr_input_method_keyboard_grab_v2 *kb_grab = keyboard_get_im_grab(keyboard);
+
+		if (kb_grab) {
+			wlr_input_method_keyboard_grab_v2_set_keyboard(kb_grab,
+					wlr_device->keyboard);
+			wlr_input_method_keyboard_grab_v2_send_modifiers(kb_grab,
+					&wlr_device->keyboard->modifiers);
+		} else {
+			struct wlr_seat *wlr_seat = keyboard->seat_device->sway_seat->wlr_seat;
+			wlr_seat_set_keyboard(wlr_seat, wlr_device);
+			wlr_seat_keyboard_notify_modifiers(wlr_seat,
+					&wlr_device->keyboard->modifiers);
+		}
 
 		uint32_t modifiers = wlr_keyboard_get_modifiers(wlr_device->keyboard);
 		determine_bar_visibility(modifiers);
diff --git a/sway/input/text_input.c b/sway/input/text_input.c
index 2a8f62228..b8c19c179 100644
--- a/sway/input/text_input.c
+++ b/sway/input/text_input.c
@@ -55,6 +55,37 @@ static void handle_im_commit(struct wl_listener *listener, void *data) {
 	wlr_text_input_v3_send_done(text_input->input);
 }
 
+static void handle_im_keyboard_grab_destroy(struct wl_listener *listener, void *data) {
+	struct sway_input_method_relay *relay = wl_container_of(listener, relay,
+		input_method_keyboard_grab_destroy);
+	struct wlr_input_method_keyboard_grab_v2 *keyboard_grab = data;
+	wl_list_remove(&relay->input_method_keyboard_grab_destroy.link);
+
+	if (keyboard_grab->keyboard) {
+		// send modifier state to original client
+		wlr_seat_keyboard_notify_modifiers(keyboard_grab->input_method->seat,
+			&keyboard_grab->keyboard->modifiers);
+	}
+}
+
+static void handle_im_grab_keyboard(struct wl_listener *listener, void *data) {
+	struct sway_input_method_relay *relay = wl_container_of(listener, relay,
+		input_method_grab_keyboard);
+	struct wlr_input_method_keyboard_grab_v2 *keyboard_grab = data;
+
+	// send modifier state to grab
+	struct wlr_keyboard *active_keyboard = wlr_seat_get_keyboard(relay->seat->wlr_seat);
+	wlr_input_method_keyboard_grab_v2_set_keyboard(keyboard_grab,
+		active_keyboard);
+	wlr_input_method_keyboard_grab_v2_send_modifiers(keyboard_grab,
+		&active_keyboard->modifiers);
+
+	wl_signal_add(&keyboard_grab->events.destroy,
+		&relay->input_method_keyboard_grab_destroy);
+	relay->input_method_keyboard_grab_destroy.notify =
+		handle_im_keyboard_grab_destroy;
+}
+
 static void text_input_set_pending_focused_surface(
 		struct sway_text_input *text_input, struct wlr_surface *surface) {
 	wl_list_remove(&text_input->pending_focused_surface_destroy.link);
@@ -245,6 +276,9 @@ static void relay_handle_input_method(struct wl_listener *listener,
 	wl_signal_add(&relay->input_method->events.commit,
 		&relay->input_method_commit);
 	relay->input_method_commit.notify = handle_im_commit;
+	wl_signal_add(&relay->input_method->events.grab_keyboard,
+		&relay->input_method_grab_keyboard);
+	relay->input_method_grab_keyboard.notify = handle_im_grab_keyboard;
 	wl_signal_add(&relay->input_method->events.destroy,
 		&relay->input_method_destroy);
 	relay->input_method_destroy.notify = handle_im_destroy;