/*****************************************************************************\
*  _  _       _          _              ___                                   *
* | || | ___ | |_  _ __ | | _  _  __ _ |_  )                                  *
* | __ |/ _ \|  _|| '_ \| || || |/ _` | / /                                   *
* |_||_|\___/ \__|| .__/|_| \_,_|\__, |/___|                                  *
*                 |_|            |___/                                        *
\*****************************************************************************/

#define _GNU_SOURCE

#include <fcntl.h>
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <errno.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/un.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <linux/types.h>
#include <linux/netlink.h>

#include "mem_utils.h"
#include "hotplug2.h"
#include "rules.h"
#include "childlist.h"

#define TERMCONDITION (persistent == 0 && \
					coldplug_p == FORK_FINISHED && \
					child == NULL && \
					highest_seqnum == get_kernel_seqnum())

event_seqnum_t highest_seqnum = 0;
pid_t coldplug_p;
int coldplug = 1;
int persistent = 0;
int max_child_c = 20;
int dumb = 0;
int terminate = 0;
int child_c;
struct hotplug2_child_t *child;
int netlink_socket;

char *modprobe_command = NULL;

inline void free_hotplug2_event(struct hotplug2_event_t *event) {
	int i;
	
	for (i = 0; i < event->env_vars_c; i++) {
		free(event->env_vars[i].key);
		free(event->env_vars[i].value);
	}
	free(event->env_vars);
	free(event->plain);
	free(event);
}

inline int get_hotplug2_event_action(char *action) {
	if (!strcmp(action, "add"))
		return ACTION_ADD;
	
	if (!strcmp(action, "remove"))
		return ACTION_REMOVE;
	
	return ACTION_UNKNOWN;
}

char *get_hotplug2_value_by_key(struct hotplug2_event_t *event, char *key) {
	int i;
	
	for (i = 0; i < event->env_vars_c; i++) {
		if (!strcmp(event->env_vars[i].key, key))
			return event->env_vars[i].value;
	}

	return NULL;
}

inline int add_hotplug2_event_env(struct hotplug2_event_t *event, char *item) {
	char *ptr, *tmp;
	
	ptr = strchr(item, '=');
	if (ptr == NULL)
		return -1;
	
	*ptr='\0';
	
	event->env_vars_c++;
	event->env_vars = xrealloc(event->env_vars, sizeof(struct env_var_t) * event->env_vars_c);
	event->env_vars[event->env_vars_c - 1].key = strdup(item);
	event->env_vars[event->env_vars_c - 1].value = strdup(ptr + 1);
	
	/*
	 * Variables not generated by kernel but demanded nonetheless...
	 */
	if (!strcmp(item, "DEVPATH")) {
		event->env_vars_c++;
		event->env_vars = xrealloc(event->env_vars, sizeof(struct env_var_t) * event->env_vars_c);
		event->env_vars[event->env_vars_c - 1].key = strdup("DEVICENAME");
		tmp = strdup(ptr + 1);
		event->env_vars[event->env_vars_c - 1].value = strdup(basename(tmp));
		free(tmp);
	}
	
	*ptr='=';
	
	return 0;
}

inline struct hotplug2_event_t *dup_hotplug2_event(struct hotplug2_event_t *src) {
	struct hotplug2_event_t *dest;
	int i;
	
	dest = xmalloc(sizeof(struct hotplug2_event_t));
	dest->action = src->action;
	dest->env_vars_c = src->env_vars_c;
	dest->env_vars = xmalloc(sizeof(struct env_var_t) * dest->env_vars_c);
	dest->plain_s = src->plain_s;
	dest->plain = xmalloc(dest->plain_s);
	memcpy(dest->plain, src->plain, dest->plain_s);
	
	for (i = 0; i < src->env_vars_c; i++) {
		dest->env_vars[i].key = strdup(src->env_vars[i].key);
		dest->env_vars[i].value = strdup(src->env_vars[i].value);
	}
	
	return dest;
}

inline struct hotplug2_event_t *get_hotplug2_event(char *event_str, int size) {
	char *ptr;
	struct hotplug2_event_t *event;
	int skip;
	
	ptr = strchr(event_str, '@');
	if (ptr == NULL) {
		return NULL;
	}
	*ptr='\0';
	
	event = xmalloc(sizeof(struct hotplug2_event_t));
	event->action = get_hotplug2_event_action(event_str);
	event->env_vars_c = 0;
	event->env_vars = NULL;
	event->plain_s = size;
	event->plain = xmalloc(size);
	memcpy(event->plain, event_str, size);
	
	skip = ++ptr - event_str;
	size -= skip;
	
	while (size > 0) {
		add_hotplug2_event_env(event, ptr);
		skip = strlen(ptr);
		ptr += skip + 1;
		size -= skip + 1;
	}
	
	return event;
}

inline event_seqnum_t get_kernel_seqnum() {
	FILE *fp;
	
	char filename[64];
	char seqnum[64];
	
	strcpy(filename, sysfs_seqnum_path);
	
	fp = fopen(filename, "r");
	if (fp == NULL)
		return 0;
	
	fread(seqnum, 1, 64, fp);
	fclose(fp);
	
	return strtoull(seqnum, NULL, 0);
}

inline int init_netlink_socket() {
	int netlink_socket;
	struct sockaddr_nl snl;
	int buffersize = 16 * 1024 * 1024;
	
	memset(&snl, 0x00, sizeof(struct sockaddr_nl));
	snl.nl_family = AF_NETLINK;
	snl.nl_pid = getpid();
	snl.nl_groups = 1;
	netlink_socket = socket(PF_NETLINK, SOCK_DGRAM, NETLINK_KOBJECT_UEVENT); 
	if (netlink_socket == -1) {
		ERROR("opening netlink","Failed socket: %s.", strerror(errno));
		return -1;
	}
	
	if (setsockopt(netlink_socket, SOL_SOCKET, SO_RCVBUFFORCE, &buffersize, sizeof(buffersize))) {
		ERROR("opening netlink","Failed setsockopt: %s. (non-critical)", strerror(errno));
		
		/* Somewhat safe default. */
		buffersize = 106496;
		
		if (setsockopt(netlink_socket, SOL_SOCKET, SO_RCVBUF, &buffersize, sizeof(buffersize))) {
			ERROR("opening netlink","Failed setsockopt: %s. (critical)", strerror(errno));
		}
	}
	
	if (bind(netlink_socket, (struct sockaddr *) &snl, sizeof(struct sockaddr_nl))) {
		ERROR("opening netlink","Failed bind: %s.", strerror(errno));
		close(netlink_socket);
		return -1;
	}
	
	return netlink_socket;
}

int get_bool_opt(char *argv, char *name, int *value) {
	int rv = -1;
	
	if (!strncmp(argv, "--no-", 5)) {
		rv = 0;
		argv+=5;
	}
	
	if (!strncmp(argv, "--", 2)) {
		rv = 1;
		argv+=2;
	}
	
	if (rv == -1)
		return -1;
	
	if (!strcmp(argv, name)) {
		*value = rv;
		return 0;
	} else {
		return -1;
	}
}

void cleanup(void) {
	pid_t p;
	
	close(netlink_socket);
	
	signal(SIGUSR1, SIG_DFL);
	signal(SIGINT, SIG_DFL);
	signal(SIGCHLD, SIG_DFL);
	
	INFO("cleanup", "Waiting for children.");
	/* Wait for our possible children... */
	while ((p = wait(NULL)) != -1)
		DBG("cleanup", "pid: %d.", p);
	INFO("cleanup", "All children terminated.");
}

void sighandler(int sig) {
	pid_t p;
	
	switch (sig) {
		/*
		 * SIGINT simply tells that yes, we caught the signal, and 
		 * exits.
		 */
		case SIGINT:
			INFO("sighandler", "Obtained SIGINT, quitting.");
			cleanup();
			exit(0);
			break;
		
		/*
		 * SIGUSR1 is handled so that if we have turned off persistency
		 * and have processed all events, we quit.
		 */
		case SIGUSR1:
			persistent = !persistent;
			INFO("sighandler", "Changed persistency to: %s", persistent ? "yes" : "no");
			if (TERMCONDITION) {
				INFO("sighandler", "No more events to be processed, quitting.");
				cleanup();
				exit(0);
			}
			break;
		
		/*
		 * SIGCHLD helps us to figure out when a child died and
		 * what kind of child it was. It may also invoke termination.
		 */
		case SIGCHLD:
			while (1) {
				p = waitpid (WAIT_ANY, NULL, WNOHANG);
				if (p <= 0)
					break;
				
				DBG("sighandler", "caught pid: %d.", p);
				if (p == coldplug_p) {
					DBG("sighandler", "coldplug_p: %d.", coldplug_p);
					coldplug_p = FORK_FINISHED;
				} else {
					child = remove_child_by_pid(child, p, NULL, &child_c);
				}
				
				DBG("sighandler", "child_c: %d, child: %p, highest_seqnum: %lld, cur_seqnum: %lld, coldplug_p: %d.\n", child_c, child, highest_seqnum, get_kernel_seqnum(), coldplug_p);
			}
			
			if (TERMCONDITION) {
				INFO("sighandler", "No more events to be processed, quitting.");
				cleanup();
				exit(0);
			}
			break;
	}
}

#ifdef HAVE_RULES
void perform_action(struct hotplug2_event_t *event, struct rules_t *rules) {
	int i, rv;
	
	for (i = 0; i < rules->rules_c; i++) {
		rv = rule_execute(event, &rules->rules[i]);
		if (rv == -1)
			break;
	}
	
	free_hotplug2_event(event);
}
#endif

void perform_dumb_action(struct hotplug2_event_t *event, char *modalias) {
	free_hotplug2_event(event);
	execl(modprobe_command, modprobe_command, "-q", modalias, NULL);
}

int get_modprobe_command() {
	pid_t p;
	int fds[2];
	char buf[18];
	FILE *fp;
	
	pipe(fds);
	
	p = fork();
	
	switch (p) {
		case -1:
			ERROR("modprobe_command","Unable to fork.");
			return -1;
			break;
		
		case 0:
			close(fds[0]);
			close(2);
			dup2(fds[1], 1);
		
			execlp("/sbin/modprobe", "/sbin/modprobe", "--version", NULL);
			exit(1);
			break;
	
		default:
			close(fds[1]);
			fp = fdopen(fds[0], "r");
			fread(buf, 1, 17, fp);
			buf[17]='\0';
		
			/* 
			 * module-init-tools can handle aliases.
			 * If we do not have a match, we use hotplug2-depwrap,
			 * even though our modprobe can do fnmatch aliases,
			 * which is the case of eg. busybox.
			 */
			if (!strcmp(buf, "module-init-tools")) {
				modprobe_command = "/sbin/modprobe";
			} else {
				modprobe_command = "/sbin/hotplug2-depwrap";
			}
			fclose(fp);
			waitpid(p, NULL, 0);
			break;
	}
	return 0;
}

int main(int argc, char *argv[]) {
	static char buffer[UEVENT_BUFFER_SIZE+512];
	struct hotplug2_event_t *tmpevent;
	char *modalias, *seqnum;
	event_seqnum_t cur_seqnum;
	pid_t p;
	int recv_errno;
	int size;
	int rv = 0;
	int i;
	char *coldplug_command = NULL;
	sigset_t block_mask;
	
	struct rules_t *rules = NULL;
	struct stat statbuf;
	void *filemap;
	int rule_fd;

	struct options_t bool_options[] = {
		{"persistent", &persistent},
		{"coldplug", &coldplug},
		{"udevtrigger", &coldplug},	/* compatibility */
#ifdef HAVE_RULES
		{"dumb", &dumb},
#endif
		{NULL, NULL}
	};
	
	for (argc--; argc > 0; argc--) {
		argv++;
		for (i = 0; bool_options[i].name != NULL; i++) {
			if (!get_bool_opt(*argv, bool_options[i].name, bool_options[i].value)) {
				break;
			} else {
				if (!strcmp(*argv, "--max-children")) {
					argv++;
					argc--;
					if (argc <= 0)
						break;
					
					max_child_c = strtol(*argv, NULL, 0);
				} else if (!strcmp(*argv, "--set-coldplug-cmd")) {
					argv++;
					argc--;
					if (argc <= 0)
						break;
					
					coldplug_command = *argv;
				} else if (!strcmp(*argv, "--set-modprobe-cmd")) {
					argv++;
					argc--;
					if (argc <= 0)
						break;
					
					modprobe_command = *argv;
				}
			}
		}
	}
	
#ifdef HAVE_RULES
	if (!dumb) {
		filemap = MAP_FAILED;
		rule_fd = open(HOTPLUG2_RULE_PATH, O_RDONLY | O_NOATIME);
		if (rule_fd == -1) {
			dumb = 1;
			ERROR("rules parse","Unable to open rules file: %s.", strerror(errno));
			goto end_rules;
		}
		
		if (fstat(rule_fd, &statbuf)) {
			dumb = 1;
			ERROR("rules parse","Unable to stat rules file: %s.", strerror(errno));
			goto end_rules;
		}
		
		filemap = mmap(0, statbuf.st_size, PROT_READ, MAP_SHARED, rule_fd, 0);
		if (filemap == MAP_FAILED) {
			dumb = 1;
			ERROR("rules parse","Unable to mmap rules file: %s.", strerror(errno));
			goto end_rules;
		}
		
		rules = rules_from_config((char*)filemap);
		if (rules == NULL) {
			ERROR("rules parse","Unable to parse rules file.");
			dumb = 1;
		}
		
end_rules:	
		if (filemap != MAP_FAILED)
			munmap(filemap, statbuf.st_size);
		if (rule_fd != -1)
			close(rule_fd);
		
		if (dumb == 1)
			ERROR("rules parse","Parsing rules failed, switching to dumb mode.");
	} else if (!modprobe_command)
#else
	if (dumb && !modprobe_command)
#endif
	{
		if (get_modprobe_command()) {
			ERROR("modprobe_command","Unable to autodetect modprobe command.");
			goto exit;
		}
		DBG("modprobe_command", "Using modprobe: `%s'.", modprobe_command);
	}
	
	netlink_socket = init_netlink_socket();
	
	if (netlink_socket == -1) {
		ERROR("netlink init","Unable to open netlink socket.");
		goto exit;
	}

	child = NULL;
	child_c = 0;
	
	signal(SIGUSR1, sighandler);
	signal(SIGINT, sighandler);
	signal(SIGCHLD, sighandler);
	
	if (coldplug) {
		if (coldplug_command == NULL)
			coldplug_command = UDEVTRIGGER_COMMAND;
		coldplug_p = fork();
		switch (coldplug_p) {
			case FORK_ERROR:
				ERROR("coldplug","Coldplug fork failed: %s.", strerror(errno));
				perror("coldplug fork failed");
				goto exit;
				break;
			case 0:
				execlp(coldplug_command, coldplug_command, NULL);
				ERROR("coldplug","Coldplug exec ('%s') failed: %s.", coldplug_command, strerror(errno));
				goto exit;
				break;
		}
	} else {
		coldplug_p = FORK_FINISHED;
	}
	
	while (!terminate) {
		size = recv(netlink_socket, &buffer, sizeof(buffer), 0);
		recv_errno = errno;
	
		tmpevent = get_hotplug2_event(buffer, size);
		
		if (tmpevent == NULL) {
			ERROR("reading events", "Malformed event read (missing action prefix).");
			continue;
		}
		
		modalias = get_hotplug2_value_by_key(tmpevent, "MODALIAS");
		seqnum = get_hotplug2_value_by_key(tmpevent, "SEQNUM");
		
		if (seqnum == NULL) {
			free_hotplug2_event(tmpevent);
			ERROR("reading events", "Malformed event read (missing SEQNUM).");
			continue;
		}
		
		cur_seqnum = strtoull(seqnum, NULL, 0);
		if (cur_seqnum > highest_seqnum)
			highest_seqnum = cur_seqnum;
		
		if (tmpevent->action == ACTION_ADD && (!dumb || modalias != NULL)) {
			/* 
			 * We have more children than we want. Wait until SIGCHLD handler reduces
			 * their numbers.
			 */
			while (child_c >= max_child_c) {
				usleep(HOTPLUG2_THROTTLE_INTERVAL);
			}
			
			sigemptyset(&block_mask);
			sigaddset(&block_mask, SIGCHLD);
			sigprocmask(SIG_BLOCK, &block_mask, 0);
			p = fork();
			switch (p) {
				case -1:
					ERROR("event","fork failed: %s.", strerror(errno));
					break;
				case 0:
					sigprocmask(SIG_UNBLOCK, &block_mask, 0);
					signal(SIGCHLD, SIG_DFL);
					signal(SIGUSR1, SIG_DFL);
#ifdef HAVE_RULES
					if (!dumb)
						perform_action(dup_hotplug2_event(tmpevent), rules);
					else
#endif
						perform_dumb_action(dup_hotplug2_event(tmpevent), modalias);
					exit(0);
					break;
				default:
					DBG("spawn", "spawning: %d.", p);
					child = add_child(child, p, cur_seqnum);
					child_c++;
					break;
			}
			sigprocmask(SIG_UNBLOCK, &block_mask, 0);
		}
		
		free_hotplug2_event(tmpevent);
	}

exit:
	signal(SIGUSR1, SIG_DFL);
	signal(SIGINT, SIG_DFL);
	signal(SIGCHLD, SIG_DFL);
	
#ifdef HAVE_RULES
	if (!dumb) {
		rules_free(rules);
		free(rules);
	}
#endif

	cleanup();
	
	return rv;
}