#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/inotify.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"

enum
{
	Maxclients = 64,
	Calltimeout = 4000,
	Starttimeout = 8000,
	Stoptimeout = 3000,
	Cappreedit = 1,
};

typedef struct Daemon Daemon;
typedef struct Response Response;

struct Daemon
{
	pid_t pid;
	int errfd;
	char root[256];
	char runtime[320];
	char config[320];
	char ibus[384];
	char bus[448];
	char home[320];
	char socket[384];
	char err[4096];
	size_t nerr;
};

struct Response
{
	int eaten;
	char commit[Ipcfieldmax+1];
	char preedit[Ipcfieldmax+1];
};

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

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

static int64_t
nowms(void)
{
	struct timespec ts;

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

static int
leftms(int64_t deadline)
{
	int64_t n;

	n = deadline - nowms();
	if(n <= 0)
		return 0;
	if(n > 0x7fffffff)
		return 0x7fffffff;
	return n;
}

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

static 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';
}

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

static int
waitsocket(Daemon *d, int notifyfd)
{
	struct pollfd pfd[2];
	struct stat st;
	char buf[4096];
	int n, status, timeout;
	int64_t deadline;

	deadline = nowms() + Starttimeout;
	for(;;){
		if(lstat(d->socket, &st) == 0 && S_ISSOCK(st.st_mode) &&
		   (st.st_mode & 0777) == 0600)
			return 1;
		do
			n = waitpid(d->pid, &status, WNOHANG);
		while(n < 0 && errno == EINTR);
		if(n == d->pid){
			d->pid = -1;
			readerrors(d);
			return fail("daemon exited before creating IPC socket");
		}
		if(n < 0){
			if(errno == ECHILD)
				d->pid = -1;
			return fail("wait for daemon startup: %s", strerror(errno));
		}
		timeout = leftms(deadline);
		if(timeout == 0)
			return fail("timed out waiting for protected IPC socket %s",
				d->socket);
		pfd[0].fd = notifyfd;
		pfd[0].events = POLLIN;
		pfd[0].revents = 0;
		pfd[1].fd = d->errfd;
		pfd[1].events = POLLIN;
		pfd[1].revents = 0;
		n = poll(pfd, 2, timeout);
		if(n < 0 && errno == EINTR)
			continue;
		if(n < 0)
			return fail("poll for IPC socket: %s", strerror(errno));
		if(n == 0)
			continue;
		if(pfd[1].revents != 0)
			readerrors(d);
		if(pfd[0].revents & POLLIN)
			while(read(notifyfd, buf, sizeof buf) < 0 && errno == EINTR)
				;
	}
}

static int
startdaemon(Daemon *d, char *program, char *mapdir)
{
	int fd, notifyfd, watch, errpipe[2];
	long maxfd;
	pid_t pid;

	memset(d, 0, sizeof *d);
	d->pid = -1;
	d->errfd = -1;
	snprintf(d->root, sizeof d->root, "/tmp/strans-ipc.XXXXXX");
	if(mkdtemp(d->root) == NULL){
		d->root[0] = '\0';
		return fail("mkdtemp: %s", strerror(errno));
	}
	if(snprintf(d->runtime, sizeof d->runtime, "%s/runtime", d->root)
	   >= (int)sizeof d->runtime ||
	   snprintf(d->config, sizeof d->config, "%s/config", d->root)
	   >= (int)sizeof d->config ||
	   snprintf(d->ibus, sizeof d->ibus, "%s/ibus", d->config)
	   >= (int)sizeof d->ibus ||
	   snprintf(d->bus, sizeof d->bus, "%s/bus", d->ibus)
	   >= (int)sizeof d->bus ||
	   snprintf(d->home, sizeof d->home, "%s/home", d->root)
	   >= (int)sizeof d->home ||
	   snprintf(d->socket, sizeof d->socket, "%s/strans.sock", d->runtime)
	   >= (int)sizeof d->socket)
		return fail("temporary path is too long");
	if(!makedir(d->runtime) || !makedir(d->config) || !makedir(d->ibus) ||
	   !makedir(d->bus) || !makedir(d->home))
		return 0;
	notifyfd = inotify_init1(IN_CLOEXEC|IN_NONBLOCK);
	if(notifyfd < 0)
		return fail("inotify_init1: %s", strerror(errno));
	watch = inotify_add_watch(notifyfd, d->runtime,
		IN_CREATE|IN_MOVED_TO|IN_ATTRIB);
	if(watch < 0){
		close(notifyfd);
		return fail("inotify_add_watch: %s", strerror(errno));
	}
	if(pipe2(errpipe, O_CLOEXEC|O_NONBLOCK) < 0){
		close(notifyfd);
		return fail("pipe2: %s", strerror(errno));
	}
	pid = fork();
	if(pid < 0){
		close(errpipe[0]);
		close(errpipe[1]);
		close(notifyfd);
		return fail("fork: %s", strerror(errno));
	}
	if(pid == 0){
		close(errpipe[0]);
		close(notifyfd);
		if(dup2(errpipe[1], STDOUT_FILENO) < 0 ||
		   dup2(errpipe[1], STDERR_FILENO) < 0)
			_exit(126);
		close(errpipe[1]);
		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", d->runtime, 1) < 0 ||
		   setenv("XDG_CONFIG_HOME", d->config, 1) < 0 ||
		   setenv("HOME", d->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, mapdir, (char*)0);
		dprintf(STDERR_FILENO, "exec %s: %s\n", program, strerror(errno));
		_exit(127);
	}
	close(errpipe[1]);
	d->pid = pid;
	d->errfd = errpipe[0];
	if(!waitsocket(d, notifyfd)){
		inotify_rm_watch(notifyfd, watch);
		close(notifyfd);
		return 0;
	}
	inotify_rm_watch(notifyfd, watch);
	close(notifyfd);
	return 1;
}

static int
clearbus(Daemon *d)
{
	DIR *dir;
	struct dirent *de;
	char path[576];
	int ok;

	dir = opendir(d->bus);
	if(dir == NULL)
		return errno == ENOENT || fail("open cleanup directory %s: %s",
			d->bus, 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(path, sizeof path, "%s/%s", d->bus, de->d_name)
		   >= (int)sizeof path){
			fail("cleanup path is too long: %s", de->d_name);
			ok = 0;
			continue;
		}
		if(unlink(path) < 0 && errno != ENOENT){
			fail("remove IBus address file %s: %s", de->d_name,
				strerror(errno));
			ok = 0;
		}
		errno = 0;
	}
	if(errno != 0){
		fail("read cleanup directory %s: %s", d->bus, strerror(errno));
		ok = 0;
	}
	closedir(dir);
	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));
}

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

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

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

	ok = 1;
	if(d->pid > 0){
		do
			n = waitpid(d->pid, &status, WNOHANG);
		while(n < 0 && errno == EINTR);
		if(n == d->pid){
			fail("daemon exited before test termination with wait status %#x",
				status);
			d->pid = -1;
			ok = 0;
		}else if(n < 0){
			fail("check daemon %ld before termination: %s",
				(long)d->pid, strerror(errno));
			if(errno == ECHILD)
				d->pid = -1;
			ok = 0;
		}
	}
	if(d->pid > 0){
		if(kill(d->pid, SIGTERM) < 0){
			fail("kill %ld: %s", (long)d->pid, strerror(errno));
			ok = 0;
		}
		reaped = 0;
		deadline = nowms() + Stoptimeout;
		for(;;){
			n = waitpid(d->pid, &status, WNOHANG);
			if(n == d->pid){
				reaped = 1;
				d->pid = -1;
				break;
			}
			if(n < 0 && errno == EINTR)
				continue;
			if(n < 0){
				fail("waitpid %ld: %s", (long)d->pid, strerror(errno));
				ok = 0;
				if(errno == ECHILD)
					d->pid = -1;
				break;
			}
			n = leftms(deadline);
			if(n == 0){
				fail("daemon %ld did not stop after SIGTERM", (long)d->pid);
				ok = 0;
				if(!killdaemon(d, &status))
					ok = 0;
				break;
			}
			pfd.fd = d->errfd;
			pfd.events = POLLIN|POLLHUP;
			pfd.revents = 0;
			if(poll(&pfd, 1, n) < 0 && errno != EINTR){
				fail("poll daemon %ld: %s", (long)d->pid, strerror(errno));
				ok = 0;
				if(!killdaemon(d, &status))
					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))){
			fail("daemon exited with unexpected wait status %#x", status);
			ok = 0;
		}
	}
	readerrors(d);
	if(d->errfd >= 0){
		close(d->errfd);
		d->errfd = -1;
	}
	if(d->socket[0] != '\0' && unlink(d->socket) < 0 && errno != ENOENT){
		fail("remove IPC socket %s: %s", d->socket, strerror(errno));
		ok = 0;
	}
	if(!clearbus(d)) ok = 0;
	if(!rmdirknown(d->bus)) ok = 0;
	if(!rmdirknown(d->ibus)) ok = 0;
	if(!rmdirknown(d->config)) ok = 0;
	if(!rmdirknown(d->runtime)) ok = 0;
	if(!rmdirknown(d->home)) ok = 0;
	if(!rmdirknown(d->root)) ok = 0;
	return ok;
}

static int
connectuntil(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;
	}
	pfd.fd = fd;
	pfd.events = POLLOUT;
	pfd.revents = 0;
	for(;;){
		n = leftms(deadline);
		if(n == 0){
			close(fd);
			errno = ETIMEDOUT;
			return -1;
		}
		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. */
static 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){
			if(errno == EINTR)
				continue;
			return -1;
		}
		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;
}

static size_t
getlen(unsigned char p[Ipclensz])
{
	return p[0] | (p[1] << 8);
}

static int
readresponseuntil(int fd, int want, Response *res, int64_t deadline,
	int quiet)
{
	unsigned char hdr[Ipcresphdrsz], len[Ipclensz];
	size_t n;
	int rv;

	memset(res, 0, sizeof *res);
	rv = readuntil(fd, hdr, sizeof hdr, deadline);
	if(rv <= 0){
		if(!quiet)
			fail("read response header: %s",
				rv == 0 ? "peer closed" : strerror(errno));
		return rv;
	}
	res->eaten = hdr[0] != 0;
	n = getlen(hdr + 1);
	if(n > Ipcfieldmax){
		fail("invalid commit length %zu", n);
		return -1;
	}
	if(n != 0){
		rv = readuntil(fd, res->commit, n, deadline);
		if(rv <= 0){
			if(!quiet)
				fail("read commit: %s",
					rv == 0 ? "peer closed" : strerror(errno));
			return rv;
		}
	}
	res->commit[n] = '\0';
	if(!want)
		return 1;
	rv = readuntil(fd, len, sizeof len, deadline);
	if(rv <= 0){
		if(!quiet)
			fail("read preedit length: %s",
				rv == 0 ? "peer closed" : strerror(errno));
		return rv;
	}
	n = getlen(len);
	if(n > Ipcfieldmax){
		fail("invalid preedit length %zu", n);
		return -1;
	}
	if(n != 0){
		rv = readuntil(fd, res->preedit, n, deadline);
		if(rv <= 0){
			if(!quiet)
				fail("read preedit: %s",
					rv == 0 ? "peer closed" : strerror(errno));
			return rv;
		}
	}
	res->preedit[n] = '\0';
	return 1;
}

static int
sendkey(int fd, int want, uint32_t mod, uint32_t key)
{
	unsigned char req[Ipcreqsz];

	ipcpackreq(req, want, mod, key);
	return ipcsend(fd, req, sizeof req) == 0;
}

static int
sendreset(int fd, int want)
{
	unsigned char req[Ipcreqsz];

	ipcpackreset(req, want);
	return ipcsend(fd, req, sizeof req) == 0;
}

static int
sendcap(int fd, int cap)
{
	unsigned char req[Ipcreqsz];

	ipcpackcap(req, cap);
	return ipcsend(fd, req, sizeof req) == 0;
}

static int
expectresponse(int fd, int want, int eaten, char *commit, char *preedit,
	char *where)
{
	Response res;

	if(readresponseuntil(fd, want, &res, nowms() + Calltimeout, 0) != 1)
		return 0;
	if(res.eaten != eaten || strcmp(res.commit, commit) != 0 ||
	   (want && strcmp(res.preedit, preedit) != 0))
		return fail("%s: eaten=%d commit=%s preedit=%s",
			where, res.eaten, res.commit, res.preedit);
	return 1;
}

static int
requestreset(int fd, int want, int eaten, char *preedit, char *where)
{
	if(!sendreset(fd, want))
		return fail("send %s: %s", where, strerror(errno));
	return expectresponse(fd, want, eaten, "", preedit, where);
}

static int
requestkey(int fd, int want, uint32_t mod, uint32_t key, int eaten,
	char *commit, char *preedit, char *where)
{
	if(!sendkey(fd, want, mod, key))
		return fail("send %s: %s", where, strerror(errno));
	return expectresponse(fd, want, eaten, commit, preedit, where);
}

static int
tryclient(char *path, int64_t deadline, int *client)
{
	Response res;
	int fd, rv;

	fd = connectuntil(path, deadline);
	if(fd < 0)
		return errno == ETIMEDOUT ? -1 : 0;
	if(!sendkey(fd, 1, 0, Kmodfirst)){
		close(fd);
		return 0;
	}
	rv = readresponseuntil(fd, 1, &res, deadline, 1);
	if(rv != 1){
		close(fd);
		return rv < 0 && errno == ETIMEDOUT ? -1 : 0;
	}
	if(res.eaten || res.commit[0] != '\0' || res.preedit[0] != '\0'){
		close(fd);
		fail("replacement client returned invalid modifier response");
		return -1;
	}
	*client = fd;
	return 1;
}

static int
waitslot(char *path, int *client, char *where)
{
	int rv;
	int64_t deadline;

	deadline = nowms() + Calltimeout;
	for(;;){
		rv = tryclient(path, deadline, client);
		if(rv > 0)
			return 1;
		if(rv < 0 || leftms(deadline) == 0)
			return fail("%s: timed out waiting for a worker slot", where);
	}
}

static int
overflowrejected(char *path)
{
	struct pollfd pfd;
	unsigned char byte;
	ssize_t n;
	int fd, timeout;
	int64_t deadline;

	deadline = nowms() + Calltimeout;
	fd = connectuntil(path, deadline);
	if(fd < 0)
		return fail("overflow connect: %s", strerror(errno));
	if(!sendkey(fd, 1, 0, Kmodfirst)){
		close(fd);
		return 1;
	}
	for(;;){
		timeout = leftms(deadline);
		if(timeout == 0){
			close(fd);
			return fail("overflow client was not rejected");
		}
		pfd.fd = fd;
		pfd.events = POLLIN;
		pfd.revents = 0;
		n = poll(&pfd, 1, timeout);
		if(n < 0 && errno == EINTR)
			continue;
		if(n <= 0){
			close(fd);
			return fail("poll overflow client: %s",
				n == 0 ? "timed out" : strerror(errno));
		}
		n = recv(fd, &byte, 1, 0);
		if(n < 0 && errno == EINTR)
			continue;
		if(n == 0 || (n < 0 && (errno == ECONNRESET || errno == EPIPE))){
			close(fd);
			return 1;
		}
		if(n < 0){
			close(fd);
			return fail("read overflow rejection: %s", strerror(errno));
		}
		close(fd);
		return fail("overflow client received a response");
	}
}

static int
runsmoke(Daemon *d)
{
	int client, ok;

	client = connectuntil(d->socket, nowms() + Calltimeout);
	if(client < 0)
		return fail("connect client: %s", strerror(errno));
	ok = sendcap(client, Cappreedit) &&
		expectresponse(client, 1, 1, "", "", "negotiate preedit") &&
		requestkey(client, 1, Mctrl, 'n', 1, "", "",
		"select Japanese") &&
		requestkey(client, 1, 0, 'k', 1, "", "k", "preedit") &&
		requestreset(client, 1, 1, "", "reset");
	close(client);
	return ok;
}

static int
runcapacity(Daemon *d)
{
	int client[Maxclients];
	int i, ok;

	for(i = 0; i < Maxclients; i++)
		client[i] = -1;
	ok = 0;
	client[0] = connectuntil(d->socket, nowms() + Calltimeout);
	if(client[0] < 0){
		fail("connect first client: %s", strerror(errno));
		goto out;
	}
	if(!requestkey(client[0], 1, Mctrl, 'n', 1, "", "",
	   "select Japanese"))
		goto out;
	for(i = 1; i < Maxclients; i++){
		client[i] = connectuntil(d->socket, nowms() + Calltimeout);
		if(client[i] < 0){
			fail("connect capacity client %d: %s", i, strerror(errno));
			goto out;
		}
		if(!requestkey(client[i], 1, 0, Kmodfirst, 0, "", "",
		   "capacity modifier"))
			goto out;
	}
	if(!overflowrejected(d->socket))
		goto out;
	close(client[Maxclients-1]);
	client[Maxclients-1] = -1;
	if(!waitslot(d->socket, &client[Maxclients-1], "inactive slot recovery"))
		goto out;
	close(client[0]);
	client[0] = -1;
	if(!waitslot(d->socket, &client[0], "active slot recovery") ||
	   !requestkey(client[0], 1, 0, 'a', 1, "", "あ",
	   "post-disconnect composition"))
		goto out;
	ok = 1;
out:
	for(i = 0; i < Maxclients; i++)
		if(client[i] >= 0)
			close(client[i]);
	return ok;
}

int
main(int argc, char **argv)
{
	Daemon daemon;
	char *program, *mapdir;
	int capacity, ok;

	capacity = argc == 4 && strcmp(argv[1], "--capacity") == 0;
	if((!capacity && argc != 3) || (capacity && argc != 4)){
		fprintf(stderr,
			"usage: ipc_live_test [--capacity] strans mapdir\n");
		return 2;
	}
	program = argv[1+capacity];
	mapdir = argv[2+capacity];
	ok = startdaemon(&daemon, program, mapdir);
	if(ok)
		ok = capacity ? runcapacity(&daemon) : runsmoke(&daemon);
	if(!ok)
		showerrors(&daemon);
	if(!stopdaemon(&daemon))
		ok = 0;
	if(!ok)
		return 1;
	printf("ipc %s: ok\n", capacity ? "connection capacity" : "live smoke");
	return 0;
}