priod: PoC skeleton

3 files changed, 347 insertions, 5 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 22afdfd..07b0f51 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -48,6 +48,9 @@ target_link_libraries (input-switch ${project_libraries})
 add_executable (fancontrol-ng fancontrol-ng.c)
 target_link_libraries (fancontrol-ng ${project_libraries})
 
+add_executable (priod priod.c)
+target_link_libraries (priod ${project_libraries})
+
 if (WITH_GDM)
 	include_directories (${gdm_INCLUDE_DIRS})
 	add_executable (gdm-switch-user gdm-switch-user.c)
@@ -75,11 +78,13 @@ install (FILES ${PROJECT_BINARY_DIR}/fancontrol-ng.service
 install (FILES fancontrol-ng.conf.example
 	DESTINATION ${CMAKE_INSTALL_DATADIR}/fancontrol-ng)
 
+# TODO: priod is also going to need a systemd unit file
+
 if (WITH_GDM)
 	install (TARGETS gdm-switch-user DESTINATION ${CMAKE_INSTALL_BINDIR})
 endif (WITH_GDM)
 
-install (TARGETS wmstatus brightness input-switch fancontrol-ng siprandom
+install (TARGETS wmstatus brightness input-switch fancontrol-ng priod siprandom
 	DESTINATION ${CMAKE_INSTALL_BINDIR})
 install (PROGRAMS shellify DESTINATION ${CMAKE_INSTALL_BINDIR})
 install (FILES LICENSE DESTINATION ${CMAKE_INSTALL_DOCDIR})
diff --git a/README.adoc b/README.adoc
index b1c8c06..63c919b 100644
--- a/README.adoc
+++ b/README.adoc
@@ -5,14 +5,16 @@ desktop-tools
 'desktop-tools' is a collection of tools to run my desktop that might be useful
 to other people as well:
 
- - 'wmstatus' does literally everything my dwm doesn't but I'd like it to. It
+ - 'wmstatus' does literally everything my i3 doesn't but I'd like it to. It
    includes PulseAudio volume management and hand-written NUT and MPD clients,
-   all in the name of liberation from GPL-licensed software of course.
- - 'brightness' allows me to change the brightness of w/e display device I have.
- - 'input-switch' likewise switches the input source of external displays.
+   all in the name of liberation from GPL-licensed software of course
+ - 'brightness' allows me to change the brightness of w/e display device I have
+ - 'input-switch' likewise switches the input source of external displays
  - 'fancontrol-ng' is a clone of fancontrol that can handle errors on resume
    from suspend instead of setting fans to maximum speed and quitting;
    in general it doesn't handle everything the original does
+ - 'priod' sets CPU, I/O and OOM killer priorities for new processes according
+   to configuration
  - 'shellify' is a simple script that sets up a shell for commands like vgdb
    and nmcli that are painfully lacking it
  - 'gdm-switch-user' tells the running GDM daemon, if any, to show the switch
diff --git a/priod.c b/priod.c
new file mode 100644
index 0000000..075f465
--- /dev/null
+++ b/priod.c
@@ -0,0 +1,335 @@
+/*
+ * priod.c: process reprioritizing daemon
+ *
+ * Thanks to http://netsplit.com/the-proc-connector-and-socket-filters
+ *
+ * Copyright (c) 2017, Přemysl Janouch <p.janouch@gmail.com>
+ *
+ * Permission to use, copy, modify, and/or distribute this software for any
+ * purpose with or without fee is hereby granted, provided that the above
+ * copyright notice and this permission notice appear in all copies.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+ * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+ * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
+ * SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+ * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
+ * OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
+ * CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+ *
+ */
+
+#define _GNU_SOURCE
+#define LIBERTY_WANT_POLLER
+
+#include "config.h"
+#undef PROGRAM_NAME
+#define PROGRAM_NAME "priod"
+#include "liberty/liberty.c"
+
+#include <linux/cn_proc.h>
+#include <linux/netlink.h>
+#include <linux/connector.h>
+
+#include <sys/resource.h>
+#include <sys/syscall.h>
+
+// --- Main program ------------------------------------------------------------
+
+struct app_context
+{
+	struct poller poller;               ///< Poller
+	bool polling;                       ///< The event loop is running
+
+	int proc_fd;                        ///< Proc connector FD
+	struct poller_fd proc_event;        ///< Proc connector read event
+};
+
+// - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
+
+static void
+log_message_custom (void *user_data, const char *quote, const char *fmt,
+	va_list ap)
+{
+	(void) user_data;
+	FILE *stream = stdout;
+
+	// TODO: sd-daemon.h log level prefixes?
+	fputs (quote, stream);
+	vfprintf (stream, fmt, ap);
+	fputs ("\n", stream);
+}
+
+// --- Signals -----------------------------------------------------------------
+
+static int g_signal_pipe[2];            ///< A pipe used to signal... signals
+
+static void
+sigterm_handler (int signum)
+{
+	(void) signum;
+
+	int original_errno = errno;
+	if (write (g_signal_pipe[1], "", 1) == -1)
+		soft_assert (errno == EAGAIN);
+	errno = original_errno;
+}
+
+static void
+setup_signal_handlers (void)
+{
+	if (pipe (g_signal_pipe) == -1)
+		exit_fatal ("%s: %s", "pipe", strerror (errno));
+
+	set_cloexec (g_signal_pipe[0]);
+	set_cloexec (g_signal_pipe[1]);
+
+	// So that the pipe cannot overflow; it would make write() block within
+	// the signal handler, which is something we really don't want to happen.
+	// The same holds true for read().
+	set_blocking (g_signal_pipe[0], false);
+	set_blocking (g_signal_pipe[1], false);
+
+	(void) signal (SIGPIPE, SIG_IGN);
+
+	struct sigaction sa;
+	sa.sa_flags = SA_RESTART;
+	sa.sa_handler = sigterm_handler;
+	sigemptyset (&sa.sa_mask);
+
+	if (sigaction (SIGINT,  &sa, NULL) == -1
+	 || sigaction (SIGTERM, &sa, NULL) == -1)
+		exit_fatal ("sigaction: %s", strerror (errno));
+}
+
+// --- Main program ------------------------------------------------------------
+
+// IO priorities are a sort-of-private kernel API with no proper headers
+
+enum
+{
+	IOPRIO_CLASS_NONE,
+	IOPRIO_CLASS_RT,
+	IOPRIO_CLASS_BE,
+	IOPRIO_CLASS_IDLE,
+};
+
+enum
+{
+	IOPRIO_WHO_PROCESS = 1,
+	IOPRIO_WHO_PGRP,
+	IOPRIO_WHO_USER,
+};
+
+#define IOPRIO_CLASS_SHIFT 13
+
+static void
+on_exec_name (struct app_context *ctx, int pid, const char *name)
+{
+	print_status ("exec %d %s", pid, name);
+
+	if (true)
+		return;
+
+	setpriority (PRIO_PROCESS, pid, 0 /* TODO -20..20 */);
+
+	// TODO: this is per thread, and there's an inherent race condition;
+	//   keep going through /proc/%d/task and reprioritize all threads;
+	//   stop trying after N-th try
+	syscall (SYS_ioprio_set, IOPRIO_WHO_PROCESS,
+		IOPRIO_CLASS_BE << IOPRIO_CLASS_SHIFT | 0 /* TODO 0..7 */);
+
+	char *path = xstrdup_printf ("/proc/%d/oom_score_adj", pid);
+	struct error *e = NULL;
+	// TODO: figure out the contents
+	if (!write_file (path, "", 0, &e))
+	{
+		print_error ("%s", e->message);
+		error_free (e);
+	}
+	free (path);
+}
+
+static void
+on_exec (struct app_context *ctx, int pid)
+{
+	char *path = xstrdup_printf ("/proc/%d/cmdline", pid);
+	struct str cmdline;
+	str_init (&cmdline);
+
+	struct error *e = NULL;
+	if (read_file (path, &cmdline, &e))
+		on_exec_name (ctx, pid, cmdline.str);
+	else
+	{
+		print_debug ("%s", e->message);
+		error_free (e);
+	}
+
+	free (path);
+	str_free (&cmdline);
+}
+
+static void
+on_netlink_message (struct app_context *ctx, struct nlmsghdr *mh)
+{
+	if (mh->nlmsg_type == NLMSG_ERROR
+	 || mh->nlmsg_type == NLMSG_NOOP)
+		return;
+
+	struct cn_msg *m = NLMSG_DATA (mh);
+	if (m->id.idx != CN_IDX_PROC
+	 || m->id.val != CN_VAL_PROC)
+		return;
+
+	struct proc_event *e = (struct proc_event *) m->data;
+	if (e->what == PROC_EVENT_EXEC)
+		on_exec (ctx, e->event_data.exit.process_tgid);
+}
+
+static void
+on_event (const struct pollfd *pfd, struct app_context *ctx)
+{
+	char buf[4096];
+	struct sockaddr_nl addr;
+	while (true)
+	{
+		socklen_t addr_len = sizeof addr;
+		ssize_t len = recvfrom (pfd->fd, buf, sizeof buf, 0,
+			(struct sockaddr *) &addr, &addr_len);
+		if (len == 0)
+			exit_fatal ("socket closed");
+		if (len < 0 && errno == EAGAIN)
+			return;
+		if (len < 0)
+			exit_fatal ("recvfrom: %s", strerror (errno));
+
+		// Make sure it comes from the kernel
+		if (addr.nl_pid)
+			continue;
+
+		for (struct nlmsghdr *mh = (struct nlmsghdr *) buf;
+			NLMSG_OK (mh, len); mh = NLMSG_NEXT (mh, len))
+			on_netlink_message (ctx, mh);
+	}
+}
+
+static void
+on_signal_pipe_readable (const struct pollfd *fd, struct app_context *ctx)
+{
+	char id = 0;
+	(void) read (fd->fd, &id, 1);
+
+	ctx->polling = false;
+}
+
+static const char *
+parse_program_arguments (int argc, char **argv)
+{
+	static const struct opt opts[] =
+	{
+		{ 'd', "debug", NULL, 0, "run in debug mode" },
+		{ 'h', "help", NULL, 0, "display this help and exit" },
+		{ 'V', "version", NULL, 0, "output version information and exit" },
+		{ 0, NULL, NULL, 0, NULL }
+	};
+
+	struct opt_handler oh;
+	opt_handler_init (&oh, argc, argv, opts, "CONFIG", "Fan controller.");
+
+	int c;
+	while ((c = opt_handler_get (&oh)) != -1)
+	switch (c)
+	{
+	case 'd':
+		g_debug_mode = true;
+		break;
+	case 'h':
+		opt_handler_usage (&oh, stdout);
+		exit (EXIT_SUCCESS);
+	case 'V':
+		printf (PROGRAM_NAME " " PROGRAM_VERSION "\n");
+		exit (EXIT_SUCCESS);
+	default:
+		print_error ("wrong options");
+		opt_handler_usage (&oh, stderr);
+		exit (EXIT_FAILURE);
+	}
+
+	argc -= optind;
+	argv += optind;
+
+	if (argc != 1)
+	{
+		opt_handler_usage (&oh, stderr);
+		exit (EXIT_FAILURE);
+	}
+
+	opt_handler_free (&oh);
+	return argv[0];
+}
+
+int
+main (int argc, char *argv[])
+{
+	g_log_message_real = log_message_custom;
+	const char *config_path = parse_program_arguments (argc, argv);
+
+	struct app_context ctx;
+	memset (&ctx, 0, sizeof ctx);
+	poller_init (&ctx.poller);
+
+	setup_signal_handlers ();
+
+	struct poller_fd signal_event;
+	poller_fd_init (&signal_event, &ctx.poller, g_signal_pipe[0]);
+	signal_event.dispatcher = (poller_fd_fn) on_signal_pipe_readable;
+	signal_event.user_data = &ctx;
+	poller_fd_set (&signal_event, POLLIN);
+
+	// TODO: load configuration so that we know what to do with the events
+
+	ctx.proc_fd = socket (PF_NETLINK,
+		SOCK_DGRAM | SOCK_NONBLOCK | SOCK_CLOEXEC, NETLINK_CONNECTOR);
+	if (ctx.proc_fd < 0)
+		exit_fatal ("cannot make a proc connector: %s", strerror (errno));
+
+	struct sockaddr_nl addr = { .nl_family = AF_NETLINK, .nl_pid = getpid (),
+		.nl_groups = CN_IDX_PROC };
+	if (bind (ctx.proc_fd, (struct sockaddr *) &addr, sizeof addr) < 0)
+		exit_fatal ("cannot make a proc connector: %s", strerror (errno));
+
+	struct
+	{
+		// Beware of padding between fields, shouldn't be any on x86-64
+		union { struct nlmsghdr netlink; char align[NLMSG_HDRLEN]; };
+		struct cn_msg connector;
+		enum proc_cn_mcast_op op;
+	}
+	subscription =
+	{
+		.netlink.nlmsg_len = sizeof subscription,
+		.netlink.nlmsg_type = NLMSG_DONE,
+		.netlink.nlmsg_pid = getpid (),
+		.connector.id.idx = CN_IDX_PROC,
+		.connector.id.val = CN_VAL_PROC,
+		.connector.len = sizeof subscription.op,
+		.op = PROC_CN_MCAST_LISTEN,
+	};
+
+	if (write (ctx.proc_fd, &subscription, sizeof subscription) < 0)
+		exit_fatal ("failed to subscribe for events: %s", strerror (errno));
+
+	poller_fd_init (&ctx.proc_event, &ctx.poller, ctx.proc_fd);
+	ctx.proc_event.dispatcher = (poller_fd_fn) on_event;
+	ctx.proc_event.user_data = &ctx;
+	poller_fd_set (&ctx.proc_event, POLLIN);
+
+	ctx.polling = true;
+	while (ctx.polling)
+		poller_run (&ctx.poller);
+
+	poller_free (&ctx.poller);
+	xclose (ctx.proc_fd);
+	return EXIT_SUCCESS;
+}