From 05a418c9fe3d1c1eddd2a28cde017d64a81cd9be Mon Sep 17 00:00:00 2001
From: Daniel De Graaf <code@danieldg.net>
Date: Sun, 6 Dec 2020 13:57:45 -0500
Subject: [PATCH] Add a connection label for wayland clients lauched by sway

This adds a few arguments to exec that cause sway to create a wayland socket
file descriptor and point the WAYLAND_SOCKET environment variable at the
socket.  This in turn allows sway to track all windows created by this client
(in a more robust manner than the existing pid-based tracking).

Note that this only works if the exec ends up launching a single wayland client
(a script that launches multiple commands will not work correctly; a process
that uses libwayland-client itself and runs other processes is fine).

"exec --use-wayland-socket command" only sets WAYLAND_SOCKET and does nothing else.

"exec --label <label> <program>" will associate the specified label with the
wl_client and all windows it creates.
---
 include/sway/client_label.h |  4 +++
 include/sway/tree/view.h    |  2 ++
 sway/commands/exec_always.c | 70 ++++++++++++++++++++++++++++++++++---
 sway/desktop/xdg_shell.c    | 45 ++++++++++++++++++++++++
 sway/tree/view.c            | 16 +++++++++
 5 files changed, 132 insertions(+), 5 deletions(-)
 create mode 100644 include/sway/client_label.h

diff --git a/include/sway/client_label.h b/include/sway/client_label.h
new file mode 100644
index 000000000..59acac767
--- /dev/null
+++ b/include/sway/client_label.h
@@ -0,0 +1,4 @@
+#include <wayland-server-core.h>
+
+char* wl_client_label_get(struct wl_client *client);
+void wl_client_label_set(struct wl_client *client, char* label);
diff --git a/include/sway/tree/view.h b/include/sway/tree/view.h
index 95708a049..381e963d9 100644
--- a/include/sway/tree/view.h
+++ b/include/sway/tree/view.h
@@ -231,6 +231,8 @@ const char *view_get_class(struct sway_view *view);
 
 const char *view_get_instance(struct sway_view *view);
 
+const char *view_get_conn_label(struct sway_view *view);
+
 uint32_t view_get_x11_window_id(struct sway_view *view);
 
 uint32_t view_get_x11_parent_id(struct sway_view *view);
diff --git a/sway/commands/exec_always.c b/sway/commands/exec_always.c
index b35065c12..82b6c1fbf 100644
--- a/sway/commands/exec_always.c
+++ b/sway/commands/exec_always.c
@@ -2,9 +2,11 @@
 #include <stdlib.h>
 #include <stdint.h>
 #include <string.h>
+#include <sys/socket.h>
 #include <sys/wait.h>
 #include <unistd.h>
 #include <signal.h>
+#include "sway/client_label.h"
 #include "sway/commands.h"
 #include "sway/config.h"
 #include "sway/server.h"
@@ -13,6 +15,7 @@
 #include "sway/tree/workspace.h"
 #include "log.h"
 #include "stringop.h"
+#include "util.h"
 
 struct cmd_results *cmd_exec_validate(int argc, char **argv) {
 	struct cmd_results *error = NULL;
@@ -27,14 +30,39 @@ struct cmd_results *cmd_exec_validate(int argc, char **argv) {
 
 struct cmd_results *cmd_exec_process(int argc, char **argv) {
 	struct cmd_results *error = NULL;
+	bool use_wl_socket = false;
+	int skip_argc = 0;
 	char *cmd = NULL;
-	if (strcmp(argv[0], "--no-startup-id") == 0) {
-		sway_log(SWAY_INFO, "exec switch '--no-startup-id' not supported, ignored.");
-		--argc; ++argv;
-		if ((error = checkarg(argc, argv[-1], EXPECTED_AT_LEAST, 1))) {
-			return error;
+	char *label = NULL;
+
+	while (skip_argc < argc && strncmp(argv[skip_argc], "--", 2) == 0) {
+		if (strcmp(argv[skip_argc], "--no-startup-id") == 0) {
+			sway_log(SWAY_INFO, "exec switch '--no-startup-id' not supported, ignored.");
+			skip_argc++;
+		} else if (strcmp(argv[skip_argc], "--use-wayland-socket") == 0) {
+			use_wl_socket = true;
+			skip_argc++;
+		} else if (strcmp(argv[skip_argc], "--label") == 0) {
+			skip_argc++;
+			if (skip_argc >= argc)
+				return cmd_results_new(CMD_INVALID, "--label requires an argument");
+			label = argv[skip_argc];
+			skip_argc++;
+		} else if (strcmp(argv[skip_argc], "--") == 0) {
+			skip_argc++;
+			break;
+		} else {
+			return cmd_results_new(CMD_INVALID, "Unknown switch %s", argv[skip_argc]);
 		}
 	}
+	if ((error = checkarg(argc - skip_argc, argv[-1], EXPECTED_AT_LEAST, 1))) {
+		return error;
+	}
+	argc -= skip_argc;
+	argv += skip_argc;
+
+	if (label)
+		use_wl_socket = true;
 
 	if (argc == 1 && (argv[0][0] == '\'' || argv[0][0] == '"')) {
 		cmd = strdup(argv[0]);
@@ -50,6 +78,17 @@ struct cmd_results *cmd_exec_process(int argc, char **argv) {
 		sway_log(SWAY_ERROR, "Unable to create pipe for fork");
 	}
 
+	int sockets[2];
+	struct wl_client* client = NULL;
+	if (use_wl_socket) {
+		if (socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) != 0) {
+			sway_log_errno(SWAY_ERROR, "socketpair failed in exec");
+			use_wl_socket = false;
+		} else {
+			sway_set_cloexec(sockets[0], true);
+		}
+	}
+
 	pid_t pid, child;
 	// Fork process
 	if ((pid = fork()) == 0) {
@@ -61,6 +100,15 @@ struct cmd_results *cmd_exec_process(int argc, char **argv) {
 		sigprocmask(SIG_SETMASK, &set, NULL);
 		signal(SIGPIPE, SIG_DFL);
 		close(fd[0]);
+		if (use_wl_socket) {
+			close(sockets[0]);
+			sway_set_cloexec(sockets[1], false);
+
+			char wayland_socket_str[16];
+			snprintf(wayland_socket_str, sizeof(wayland_socket_str),
+					"%d", sockets[1]);
+			setenv("WAYLAND_SOCKET", wayland_socket_str, true);
+		}
 		if ((child = fork()) == 0) {
 			close(fd[1]);
 			execlp("sh", "sh", "-c", cmd, (void *)NULL);
@@ -77,10 +125,18 @@ struct cmd_results *cmd_exec_process(int argc, char **argv) {
 		free(cmd);
 		close(fd[0]);
 		close(fd[1]);
+		if (use_wl_socket) {
+			close(sockets[0]);
+			close(sockets[1]);
+		}
 		return cmd_results_new(CMD_FAILURE, "fork() failed");
 	}
 	free(cmd);
 	close(fd[1]); // close write
+	if (use_wl_socket) {
+		close(sockets[1]);
+		client = wl_client_create(server.wl_display, sockets[0]);
+	}
 	ssize_t s = 0;
 	while ((size_t)s < sizeof(pid_t)) {
 		s += read(fd[0], ((uint8_t *)&child) + s, sizeof(pid_t) - s);
@@ -95,6 +151,10 @@ struct cmd_results *cmd_exec_process(int argc, char **argv) {
 		return cmd_results_new(CMD_FAILURE, "Second fork() failed");
 	}
 
+	if (client && label) {
+		wl_client_label_set(client, strdup(label));
+	}
+
 	return cmd_results_new(CMD_SUCCESS, NULL);
 }
 
diff --git a/sway/desktop/xdg_shell.c b/sway/desktop/xdg_shell.c
index 51168f4c6..c2b0c0463 100644
--- a/sway/desktop/xdg_shell.c
+++ b/sway/desktop/xdg_shell.c
@@ -6,6 +6,7 @@
 #include <wlr/types/wlr_xdg_shell.h>
 #include <wlr/util/edges.h>
 #include "log.h"
+#include "sway/client_label.h"
 #include "sway/decoration.h"
 #include "sway/desktop.h"
 #include "sway/desktop/transaction.h"
@@ -521,3 +522,47 @@ void handle_xdg_shell_surface(struct wl_listener *listener, void *data) {
 
 	xdg_surface->data = xdg_shell_view;
 }
+
+struct wl_client_label {
+	struct wl_listener listener;
+	char* label;
+};
+
+static void wl_client_label_notify_fn(struct wl_listener *listener, void *data) {
+	// struct wl_client *client = data;
+	struct wl_client_label *label = wl_container_of(listener, label, listener);
+	free(label->label);
+	free(label);
+}
+
+char *wl_client_label_get(struct wl_client *client) {
+	struct wl_listener *label_l =
+		wl_client_get_destroy_listener(client, wl_client_label_notify_fn);
+	if (!label_l) {
+		return NULL;
+	}
+	struct wl_client_label *label_s = wl_container_of(label_l, label_s, listener);
+	return label_s->label;
+}
+
+void wl_client_label_set(struct wl_client *client, char* label) {
+	struct wl_listener *label_l =
+		wl_client_get_destroy_listener(client, wl_client_label_notify_fn);
+	struct wl_client_label *label_s;
+	if (label_l) {
+		label_s = wl_container_of(label_l, label_s, listener);
+		free(label_s->label);
+	} else if (label == NULL) {
+		return;
+	} else {
+		label_s = calloc(1, sizeof(*label_s));
+		if (label_s == NULL) {
+			sway_log(SWAY_ERROR, "Allocation failed");
+			free(label);
+			return;
+		}
+		label_s->listener.notify = wl_client_label_notify_fn;
+		wl_client_add_destroy_listener(client, &label_s->listener);
+	}
+	label_s->label = label;
+}
diff --git a/sway/tree/view.c b/sway/tree/view.c
index 7d9e038d2..ec4b3beba 100644
--- a/sway/tree/view.c
+++ b/sway/tree/view.c
@@ -14,6 +14,7 @@
 #endif
 #include "list.h"
 #include "log.h"
+#include "sway/client_label.h"
 #include "sway/criteria.h"
 #include "sway/commands.h"
 #include "sway/desktop.h"
@@ -110,6 +111,21 @@ const char *view_get_instance(struct sway_view *view) {
 	}
 	return NULL;
 }
+
+const char *view_get_conn_label(struct sway_view *view) {
+	switch (view->type) {
+	case SWAY_VIEW_XDG_SHELL:;
+		struct wl_client *client =
+			wl_resource_get_client(view->surface->resource);
+		return wl_client_label_get(client);
+#if HAVE_XWAYLAND
+	case SWAY_VIEW_XWAYLAND:;
+		// Is this concept useful in xwayland?
+		break;
+#endif
+	}
+	return NULL;
+}
 #if HAVE_XWAYLAND
 uint32_t view_get_x11_window_id(struct sway_view *view) {
 	if (view->impl->get_int_prop) {