#define _GNU_SOURCE
#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <poll.h>
#include <signal.h>
#include <stdarg.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/un.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

#include "ipc.h"
#include "live.h"

char *testname = "live_test";

int
fail(char *fmt, ...)
{
	va_list ap;

	fprintf(stderr, "%s: ", testname);
	va_start(ap, fmt);
	vfprintf(stderr, fmt, ap);
	va_end(ap);
	fputc('\n', stderr);
	return 0;
}

int64_t
nowms(void)
{
	struct timespec ts;

	if(clock_gettime(CLOCK_MONOTONIC, &ts) < 0)
		return -1;
	return (int64_t)ts.tv_sec * 1000 + ts.tv_nsec / 1000000;
}

/* A clock that cannot be read counts as an expired deadline. */
int
leftms(int64_t deadline)
{
	int64_t n;

	n = nowms();
	if(n < 0)
		return 0;
	n = deadline - n;
	if(n <= 0)
		return 0;
	return n > INT_MAX ? INT_MAX : (int)n;
}

void
pausems(int ms)
{
	struct timespec ts;

	ts.tv_sec = ms / 1000;
	ts.tv_nsec = (ms % 1000) * 1000000L;
	while(nanosleep(&ts, &ts) < 0 && errno == EINTR)
		;
}

int
makedir(char *path)
{
	if(mkdir(path, 0700) == 0)
		return 1;
	return fail("mkdir %s: %s", path, strerror(errno));
}

int
cleardir(char *path)
{
	DIR *dir;
	struct dirent *de;
	char name[576];
	int ok;

	if(path[0] == '\0')
		return 1;
	dir = opendir(path);
	if(dir == NULL)
		return errno == ENOENT ||
			fail("open %s: %s", path, strerror(errno));
	ok = 1;
	errno = 0;
	while((de = readdir(dir)) != NULL){
		if(strcmp(de->d_name, ".") == 0 || strcmp(de->d_name, "..") == 0)
			continue;
		if(snprintf(name, sizeof name, "%s/%s", path, de->d_name)
		   >= (int)sizeof name)
			ok = fail("cleanup path is too long: %s", de->d_name);
		else if(unlink(name) < 0 && errno != ENOENT)
			ok = fail("remove %s: %s", name, strerror(errno));
		errno = 0;
	}
	if(errno != 0)
		ok = fail("read %s: %s", path, strerror(errno));
	if(closedir(dir) < 0)
		ok = fail("close %s: %s", path, strerror(errno));
	return ok;
}

static int
rmdirknown(char *path)
{
	if(path[0] == '\0' || rmdir(path) == 0 || errno == ENOENT)
		return 1;
	return fail("rmdir %s: %s", path, strerror(errno));
}

int
livesetup(Live *l, char *prefix)
{
	memset(l, 0, sizeof *l);
	if(nowms() < 0)
		return fail("read monotonic clock: %s", strerror(errno));
	if(snprintf(l->root, sizeof l->root, "/tmp/strans-%s.XXXXXX", prefix)
	   >= (int)sizeof l->root){
		l->root[0] = '\0';
		return fail("temporary root name is too long");
	}
	if(mkdtemp(l->root) == NULL){
		l->root[0] = '\0';
		return fail("mkdtemp: %s", strerror(errno));
	}
	if(snprintf(l->runtime, sizeof l->runtime, "%s/runtime", l->root)
	   >= (int)sizeof l->runtime ||
	   snprintf(l->config, sizeof l->config, "%s/config", l->root)
	   >= (int)sizeof l->config ||
	   snprintf(l->ibus, sizeof l->ibus, "%s/ibus", l->config)
	   >= (int)sizeof l->ibus ||
	   snprintf(l->bus, sizeof l->bus, "%s/bus", l->ibus)
	   >= (int)sizeof l->bus ||
	   snprintf(l->home, sizeof l->home, "%s/home", l->root)
	   >= (int)sizeof l->home ||
	   snprintf(l->socket, sizeof l->socket, "%s/strans.sock", l->runtime)
	   >= (int)sizeof l->socket)
		return fail("temporary path is too long");
	return makedir(l->runtime) && makedir(l->config) && makedir(l->ibus) &&
		makedir(l->bus) && makedir(l->home);
}

/* The daemon owns its endpoints, so leftovers elsewhere are a failure. */
int
liveclean(Live *l)
{
	int ok;

	ok = 1;
	if(l->socket[0] != '\0' && unlink(l->socket) < 0 && errno != ENOENT)
		ok = fail("remove IPC socket %s: %s", l->socket, strerror(errno));
	if(!cleardir(l->bus)) ok = 0;
	if(!rmdirknown(l->bus)) ok = 0;
	if(!rmdirknown(l->ibus)) ok = 0;
	if(!rmdirknown(l->config)) ok = 0;
	if(!rmdirknown(l->runtime)) ok = 0;
	if(!rmdirknown(l->home)) ok = 0;
	if(!rmdirknown(l->root)) ok = 0;
	return ok;
}

void
daemoninit(Daemon *d, char *name)
{
	memset(d, 0, sizeof *d);
	d->pid = -1;
	d->errfd = -1;
	d->name = name;
}

void
readerrors(Daemon *d)
{
	ssize_t n;

	if(d->errfd < 0)
		return;
	while(d->nerr + 1 < sizeof d->err){
		n = read(d->errfd, d->err + d->nerr, sizeof d->err - d->nerr - 1);
		if(n > 0){
			d->nerr += n;
			continue;
		}
		if(n < 0 && errno == EINTR)
			continue;
		break;
	}
	d->err[d->nerr] = '\0';
}

void
showerrors(Daemon *d)
{
	readerrors(d);
	if(d->nerr != 0)
		fprintf(stderr, "%s: %s stderr:\n%s", testname, d->name, d->err);
}

int
closeerrors(Daemon *d)
{
	int fd;

	readerrors(d);
	fd = d->errfd;
	d->errfd = -1;
	if(fd >= 0 && close(fd) < 0)
		return fail("close %s stderr: %s", d->name, strerror(errno));
	return 1;
}

/* Point both standard streams at the capture pipe, even if it is 1 or 2. */
static int
childfds(int fd)
{
	int flags;

	if(fd > STDERR_FILENO){
		if(dup2(fd, STDOUT_FILENO) < 0 || dup2(fd, STDERR_FILENO) < 0)
			return 0;
		close(fd);
		return 1;
	}
	flags = fcntl(fd, F_GETFD);
	if(flags < 0 || fcntl(fd, F_SETFD, flags & ~FD_CLOEXEC) < 0)
		return 0;
	return dup2(fd, STDOUT_FILENO) >= 0 && dup2(fd, STDERR_FILENO) >= 0;
}

int
startdaemon(Live *l, Daemon *d, char *program, char *arg)
{
	int errpipe[2], fd;
	long maxfd;
	pid_t pid;

	if(pipe2(errpipe, O_CLOEXEC|O_NONBLOCK) < 0)
		return fail("pipe2: %s", strerror(errno));
	pid = fork();
	if(pid < 0){
		close(errpipe[0]);
		close(errpipe[1]);
		return fail("fork: %s", strerror(errno));
	}
	if(pid == 0){
		close(errpipe[0]);
		if(!childfds(errpipe[1]))
			_exit(126);
		if(close_range(3, UINT_MAX, 0) < 0){
			maxfd = sysconf(_SC_OPEN_MAX);
			if(maxfd < 0)
				maxfd = 1024;
			for(fd = 3; fd < maxfd; fd++)
				close(fd);
		}
		if(setenv("XDG_RUNTIME_DIR", l->runtime, 1) < 0 ||
		   setenv("XDG_CONFIG_HOME", l->config, 1) < 0 ||
		   setenv("HOME", l->home, 1) < 0 || unsetenv("DISPLAY") < 0 ||
		   unsetenv("DBUS_SESSION_BUS_ADDRESS") < 0 ||
		   unsetenv("IBUS_ADDRESS") < 0 ||
		   unsetenv("IBUS_ADDRESS_FILE") < 0){
			dprintf(STDERR_FILENO, "set daemon environment: %s\n",
				strerror(errno));
			_exit(126);
		}
		execl(program, program, arg, (char*)0);
		dprintf(STDERR_FILENO, "exec %s: %s\n", program, strerror(errno));
		_exit(127);
	}
	close(errpipe[1]);
	d->pid = pid;
	d->errfd = errpipe[0];
	return 1;
}

int
daemonalive(Daemon *d)
{
	int n, status;

	if(d->pid <= 0)
		return fail("%s is no longer running", d->name);
	do
		n = waitpid(d->pid, &status, WNOHANG);
	while(n < 0 && errno == EINTR);
	if(n == 0)
		return 1;
	if(n == d->pid){
		d->pid = -1;
		return fail("%s exited unexpectedly with wait status %#x", d->name,
			status);
	}
	if(n < 0){
		fail("check %s %ld: %s", d->name, (long)d->pid, strerror(errno));
		if(errno == ECHILD)
			d->pid = -1;
		return 0;
	}
	return fail("waitpid returned the wrong child for %s", d->name);
}

int
killdaemon(Daemon *d, int *status)
{
	int n, ok, wait;

	if(d->pid <= 0)
		return 1;
	ok = 1;
	if(kill(d->pid, SIGKILL) < 0 && errno != ESRCH)
		ok = fail("kill -9 %s %ld: %s", d->name, (long)d->pid,
			strerror(errno));
	do
		n = waitpid(d->pid, &wait, 0);
	while(n < 0 && errno == EINTR);
	if(n != d->pid)
		ok = fail("reap %s %ld after SIGKILL: %s", d->name, (long)d->pid,
			n < 0 ? strerror(errno) : "wrong child");
	else if(status != NULL)
		*status = wait;
	d->pid = -1;
	return ok;
}

int
stopdaemon(Daemon *d)
{
	struct pollfd pfd;
	int n, ok, reaped, status;
	int64_t deadline;

	if(d->pid <= 0)
		return 1;
	if(!daemonalive(d))
		return 0;
	ok = 1;
	status = 0;
	reaped = 0;
	if(kill(d->pid, SIGTERM) < 0 && errno != ESRCH)
		ok = fail("kill %s %ld: %s", d->name, (long)d->pid, strerror(errno));
	deadline = nowms() + Stoptimeout;
	while(d->pid > 0){
		do
			n = waitpid(d->pid, &status, WNOHANG);
		while(n < 0 && errno == EINTR);
		if(n == d->pid){
			reaped = 1;
			d->pid = -1;
			break;
		}
		if(n < 0){
			fail("waitpid %s %ld: %s", d->name, (long)d->pid,
				strerror(errno));
			ok = 0;
			if(errno == ECHILD)
				d->pid = -1;
			else if(!killdaemon(d, NULL))
				ok = 0;
			break;
		}
		n = leftms(deadline);
		if(n == 0){
			fail("%s %ld did not stop after SIGTERM", d->name, (long)d->pid);
			ok = 0;
			if(!killdaemon(d, NULL))
				ok = 0;
			break;
		}
		pfd.fd = d->errfd;
		pfd.events = POLLIN|POLLHUP;
		pfd.revents = 0;
		if(poll(&pfd, 1, n) < 0 && errno != EINTR){
			fail("poll %s %ld: %s", d->name, (long)d->pid, strerror(errno));
			ok = 0;
			if(!killdaemon(d, NULL))
				ok = 0;
			break;
		}
		readerrors(d);
	}
	/* plan9port turns a caught termination note into exit status 1. */
	if(reaped && !((WIFSIGNALED(status) && WTERMSIG(status) == SIGTERM) ||
	   (WIFEXITED(status) && WEXITSTATUS(status) == 1)))
		ok = fail("%s exited with unexpected wait status %#x", d->name,
			status);
	readerrors(d);
	return ok;
}

int
findaddress(Live *l)
{
	DIR *dir;
	struct dirent *de;
	int count, ok;

	dir = opendir(l->bus);
	if(dir == NULL){
		fail("open IBus directory %s: %s", l->bus, strerror(errno));
		return -1;
	}
	count = 0;
	ok = 1;
	errno = 0;
	while((de = readdir(dir)) != NULL){
		if(strcmp(de->d_name, ".") == 0 || strcmp(de->d_name, "..") == 0 ||
		   strstr(de->d_name, ".tmp.") != NULL)
			continue;
		count++;
		if(snprintf(l->addrfile, sizeof l->addrfile, "%s/%s", l->bus,
		   de->d_name) >= (int)sizeof l->addrfile){
			ok = fail("IBus address path is too long");
			break;
		}
	}
	if(errno != 0)
		ok = fail("read IBus directory %s: %s", l->bus, strerror(errno));
	if(closedir(dir) < 0)
		ok = fail("close IBus directory %s: %s", l->bus, strerror(errno));
	if(!ok)
		return -1;
	if(count > 1){
		fail("found %d IBus address files", count);
		return -1;
	}
	return count;
}

int
readfile(char *path, char *buf, size_t cap, size_t *nread)
{
	struct stat st;
	ssize_t n;
	size_t off;
	char extra;
	int fd, ok;

	fd = open(path, O_RDONLY|O_CLOEXEC);
	if(fd < 0)
		return fail("open %s: %s", path, strerror(errno));
	ok = 0;
	if(fstat(fd, &st) < 0){
		fail("stat open file %s: %s", path, strerror(errno));
		goto out;
	}
	if(st.st_size < 0 || (uintmax_t)st.st_size >= cap){
		fail("file %s is too large", path);
		goto out;
	}
	off = 0;
	while(off < (size_t)st.st_size){
		n = read(fd, buf + off, (size_t)st.st_size - off);
		if(n < 0 && errno == EINTR)
			continue;
		if(n <= 0){
			fail("read %s: %s", path,
				n == 0 ? "unexpected EOF" : strerror(errno));
			goto out;
		}
		off += n;
	}
	do
		n = read(fd, &extra, 1);
	while(n < 0 && errno == EINTR);
	if(n != 0){
		fail("file %s changed while being read", path);
		goto out;
	}
	buf[off] = '\0';
	*nread = off;
	ok = 1;
out:
	if(close(fd) < 0){
		fail("close %s: %s", path, strerror(errno));
		ok = 0;
	}
	return ok;
}

int
parseaddress(char *contents, size_t ncontents, char *address, size_t naddress,
	pid_t *declared)
{
	char found[512];
	long pid;
	int consumed;

	consumed = -1;
	if(strlen(contents) != ncontents ||
	   sscanf(contents, "IBUS_ADDRESS=%511[^\n]\nIBUS_DAEMON_PID=%ld\n%n",
	   found, &pid, &consumed) != 2 || consumed != (int)ncontents || pid <= 0)
		return fail("invalid IBus address file contents");
	if(snprintf(address, naddress, "%s", found) >= (int)naddress)
		return fail("private IBus address is too long");
	*declared = (pid_t)pid;
	return 1;
}

/*
 * Wait for the protected IPC socket and, when addr is not null, for the
 * IBus address file the daemon publishes for itself.
 */
int
waitready(Live *l, Daemon *d, char *addr, size_t naddr)
{
	struct pollfd pfd;
	struct stat st;
	char contents[2048];
	size_t ncontents;
	pid_t declared;
	int count, n, timeout;
	int64_t deadline;

	deadline = nowms() + Starttimeout;
	for(;;){
		if(lstat(l->socket, &st) == 0 && S_ISSOCK(st.st_mode) &&
		   (st.st_mode & 0777) == 0600){
			if(addr == NULL)
				return 1;
			count = findaddress(l);
			if(count < 0)
				return 0;
			if(count == 1 &&
			   readfile(l->addrfile, contents, sizeof contents, &ncontents) &&
			   parseaddress(contents, ncontents, addr, naddr, &declared) &&
			   declared == d->pid)
				return 1;
		}
		if(!daemonalive(d))
			return 0;
		timeout = leftms(deadline);
		if(timeout == 0)
			return fail("timed out waiting for the %s endpoints", d->name);
		pfd.fd = d->errfd;
		pfd.events = POLLIN|POLLHUP;
		pfd.revents = 0;
		n = poll(&pfd, 1, timeout > 20 ? 20 : timeout);
		if(n < 0 && errno == EINTR)
			continue;
		if(n < 0)
			return fail("poll %s readiness: %s", d->name, strerror(errno));
		if(pfd.revents & (POLLERR|POLLNVAL))
			return fail("%s stderr pipe became unusable: %#x", d->name,
				pfd.revents);
		if(pfd.revents & (POLLIN|POLLHUP))
			readerrors(d);
	}
}

int
connectsocket(char *path, int64_t deadline)
{
	struct sockaddr_un addr;
	struct pollfd pfd;
	socklen_t nerr;
	int err, fd, flags, n;

	memset(&addr, 0, sizeof addr);
	addr.sun_family = AF_UNIX;
	if(snprintf(addr.sun_path, sizeof addr.sun_path, "%s", path)
	   >= (int)sizeof addr.sun_path){
		errno = ENAMETOOLONG;
		return -1;
	}
	fd = socket(AF_UNIX, SOCK_STREAM|SOCK_CLOEXEC|SOCK_NONBLOCK, 0);
	if(fd < 0)
		return -1;
	if(connect(fd, (struct sockaddr*)&addr, sizeof addr) < 0 &&
	   errno != EINPROGRESS){
		err = errno;
		close(fd);
		errno = err;
		return -1;
	}
	for(;;){
		n = leftms(deadline);
		if(n == 0){
			close(fd);
			errno = ETIMEDOUT;
			return -1;
		}
		pfd.fd = fd;
		pfd.events = POLLOUT;
		pfd.revents = 0;
		n = poll(&pfd, 1, n);
		if(n < 0 && errno == EINTR)
			continue;
		if(n <= 0){
			err = n == 0 ? ETIMEDOUT : errno;
			close(fd);
			errno = err;
			return -1;
		}
		break;
	}
	err = 0;
	nerr = sizeof err;
	if(getsockopt(fd, SOL_SOCKET, SO_ERROR, &err, &nerr) < 0 || err != 0){
		if(err == 0)
			err = errno;
		close(fd);
		errno = err;
		return -1;
	}
	flags = fcntl(fd, F_GETFL);
	if(flags < 0 || fcntl(fd, F_SETFL, flags & ~O_NONBLOCK) < 0){
		err = errno;
		close(fd);
		errno = err;
		return -1;
	}
	return fd;
}

/* 1 is complete, 0 is peer closure, and -1 is a timeout or I/O error. */
int
readuntil(int fd, void *buf, size_t n, int64_t deadline)
{
	struct pollfd pfd;
	unsigned char *p;
	ssize_t r;
	int timeout;

	p = buf;
	while(n > 0){
		timeout = leftms(deadline);
		if(timeout == 0){
			errno = ETIMEDOUT;
			return -1;
		}
		pfd.fd = fd;
		pfd.events = POLLIN;
		pfd.revents = 0;
		r = poll(&pfd, 1, timeout);
		if(r < 0 && errno == EINTR)
			continue;
		if(r <= 0){
			if(r == 0)
				errno = ETIMEDOUT;
			return -1;
		}
		r = recv(fd, p, n, 0);
		if(r < 0 && errno == EINTR)
			continue;
		if(r < 0)
			return -1;
		if(r == 0)
			return 0;
		p += r;
		n -= r;
	}
	return 1;
}

int
ipcrequest(int fd, uint32_t mod, uint32_t key, unsigned char *want,
	size_t nwant, char *where)
{
	unsigned char req[Ipcreqsz], got[16];
	int rv;

	if(nwant > sizeof got)
		return fail("%s expected IPC response is too large", where);
	ipcpackreq(req, 1, mod, key);
	if(ipcsend(fd, req, sizeof req) < 0)
		return fail("send %s IPC request: %s", where, strerror(errno));
	rv = readuntil(fd, got, nwant, nowms() + Calltimeout);
	if(rv != 1)
		return fail("read %s IPC response: %s", where,
			rv == 0 ? "peer closed" : strerror(errno));
	if(memcmp(got, want, nwant) != 0)
		return fail("%s IPC response did not match", where);
	return 1;
}

int
ipcprobe(int fd, char *where)
{
	unsigned char empty[] = {0, 0, 0, 0, 0};

	return ipcrequest(fd, 0, Kmodfirst, empty, sizeof empty, where);
}