/*
 * Copy me if you can.
 * by 20h
 */

#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <poll.h>
#include <ctype.h>
#include <string.h>
#include <errno.h>
#include <signal.h>
#include <libgen.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <linux/types.h>
#include <linux/netlink.h>

#include "arg.h"

char *argv0;
int listfd = -1;
int dofork = 0, dodebug = 0;

void
edie(char *fmt, ...)
{
	va_list fmtargs;

	va_start(fmtargs, fmt);
	vfprintf(stderr, fmt, fmtargs);
	va_end(fmtargs);
	fprintf(stderr, ": ");

	perror(NULL);

	exit(1);
}

void
die(char *fmt, ...)
{
	va_list fmtargs;

	va_start(fmtargs, fmt);
	vfprintf(stderr, fmt, fmtargs);
	va_end(fmtargs);

	exit(1);
}

void
dbg(char *fmt, ...)
{
	va_list fmtargs;

	if (dodebug) {
		fprintf(stderr, "%s: ", argv0);
		va_start(fmtargs, fmt);
		vfprintf(stderr, fmt, fmtargs);
		va_end(fmtargs);
		fprintf(stderr, "\n");
	}
}

void
disableoom(void)
{
	int fd;

	fd = open("/proc/self/oom_score_adj", O_RDWR);
	if (fd < 0) {
		fd = open("/proc/self/oom_adj", O_RDWR);
		if (fd < 0)
			edie("disabling oom failed.");
		write(fd, "-17", 3);
		close(fd);
	} else {
		write(fd, "-1000", 5);
		close(fd);
	}
}

void
child(char *runpath)
{
	int fd, pid;

	if (!(pid = fork())) {
		if (dofork && !dodebug) {
			fd = open("/dev/null", O_RDWR);
			if (fd >= 0) {
				dup2(fd, 1);
				dup2(fd, 2);
				if (fd > 2)
					close(fd);
			}
		}

		dbg("running %s", runpath);
		if (execlp(runpath, basename(runpath), NULL) < 0)
			edie("execvp");
		exit(0);
	}
	if (pid < 0)
		edie("fork");

	waitpid(pid, NULL, 0);
}

void
sighandler(int sig)
{
	switch(sig) {
	case SIGHUP:
	case SIGINT:
	case SIGQUIT:
	case SIGABRT:
	case SIGTERM:
	case SIGKILL:
		if (listfd >= 0) {
			shutdown(listfd, SHUT_RDWR);
			close(listfd);
		}
		exit(0);
		break;
	default:
		break;
	}
}

void
initsignals(void)
{
	signal(SIGHUP, sighandler);
	signal(SIGINT, sighandler);
	signal(SIGQUIT, sighandler);
	signal(SIGABRT, sighandler);
	signal(SIGTERM, sighandler);
	signal(SIGKILL, sighandler);

	signal(SIGCHLD, SIG_IGN);
	signal(SIGPIPE, SIG_IGN);
}

void
usage(void)
{
	die("usage: %s [-hdb] [-ku] [-f subsystem] [-r run]\n", argv0);
}

int
main(int argc, char *argv[])
{
	struct sockaddr_nl nls, cnls;
	struct pollfd fds;
	struct msghdr hdr;
	struct iovec iov;
	char buf[4097], *subsystem, *runpath, *key, *value,
	     cbuf[CMSG_SPACE(sizeof(struct ucred))];
	struct cmsghdr *chdr;
	struct ucred *cred;
	int i, len, slen, showudev, showkernel;

	showkernel = 1;
	showudev = 1;
	subsystem = NULL;
	runpath = "/bin/mdev";

	ARGBEGIN {
	case 'b':
		dofork = 1;
		break;
	case 'd':
		dodebug = 1;
		break;
	case 'f':
		subsystem = EARGF(usage());
		break;
	case 'k':
		showudev = 0;
		break;
	case 'u':
		showkernel = 0;
		break;
	case 'r':
		runpath = EARGF(usage());
		break;
	default:
		usage();
	} ARGEND;

	memset(&nls, 0, sizeof(nls));
	nls.nl_family = AF_NETLINK;
	nls.nl_pid = getpid();
	nls.nl_groups = -1;

	fds.events = POLLIN;
	fds.fd = socket(PF_NETLINK, SOCK_DGRAM, NETLINK_KOBJECT_UEVENT);
	listfd = fds.fd;
	if (fds.fd < 0)
		edie("socket");

	slen = 128*1024*1024;
	if (setsockopt(fds.fd, SOL_SOCKET, SO_RCVBUFFORCE, &slen,
				sizeof(slen)) < 0) {
		edie("setsockopt");
	}
	slen = 1;
	if (setsockopt(fds.fd, SOL_SOCKET, SO_PASSCRED, &slen,
				sizeof(slen)) < 0) {
		edie("setsockopt");
	}

	if (bind(fds.fd, (void *)&nls, sizeof(nls)))
		edie("bind");

	if (dofork) {
		if (daemon(0, 0) < 0)
			edie("daemon");
		umask(022);
	}

	initsignals();
	disableoom();

	buf[sizeof(buf)-1] = '\0';
	while (poll(&fds, 1, -1) > -1) {
		clearenv();
		setenv("PATH", "/sbin:/bin", 1);

		iov.iov_base = &buf;
		iov.iov_len = sizeof(buf);
		memset(&hdr, 0, sizeof(hdr));
		hdr.msg_iov = &iov;
		hdr.msg_iovlen = 1;
		hdr.msg_control = cbuf;
		hdr.msg_controllen = sizeof(cbuf);
		hdr.msg_name = &cnls;
		hdr.msg_namelen = sizeof(cnls);

		len = recvmsg(fds.fd, &hdr, 0);
		if (len < 0) {
			if (errno == EINTR)
				continue;
			edie("recvmsg");
		}
		if (len < 32 || len >= sizeof(buf))
			continue;

		chdr = CMSG_FIRSTHDR(&hdr);
		if (chdr == NULL || chdr->cmsg_type != SCM_CREDENTIALS)
			continue;

		/*
		 * Don't allow anyone but root to send us messages.
		 *
		 * We will allow users to send us messages, when
		 * udev is enabled. Udev is just a toy you should
		 * only use for testing.
		 */
		cred = (struct ucred *)CMSG_DATA(chdr);
		if (cred->uid != 0 && !showudev)
			continue;

		if (!memcmp(buf, "libudev", 8)) {
			/*
			 * Receiving messages from udev is insecure.
			 */
			if (!showudev)
				continue;
		} else {
			if (!showkernel)
				continue;
			/*
			 * Kernel messages shouldn't come from the
			 * userspace.
			 */
			if (cnls.nl_pid > 0)
				continue;
		}

		for (i = 0; i < len; i += slen + 1) {
			key = buf + i;
			value = strchr(key, '=');
			slen = strlen(buf+i);

			if (!slen || value == NULL)
				continue;
			if (subsystem && !strncmp(key, "SUBSYSTEM=", 10)
					&& !strstr(key+10, subsystem)) {
				dbg("subsystem filter '%s' applied.",
						subsystem);
				break;
			}

			value[0] = '\0';
			value++;

			/*
			 * We generally trust the kernel. But there
			 * might be some udev flaw. (It's >20k sloc!)
			 */
			if (strcmp(key, "PATH")) {
				setenv(key, value, 1);
				dbg("%s = \"%s\"", key, value);
			}
		}
		if (getenv("ACTION") != NULL &&
				getenv("DEVPATH") != NULL &&
				getenv("SUBSYSTEM") != NULL &&
				getenv("SEQNUM") != NULL) {
			child(runpath);
		}
	}

	shutdown(listfd, SHUT_RDWR);
	close(listfd);

	return 0;
}