#define _GNU_SOURCE
#include <dbus/dbus.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/stat.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

enum
{
	Maxconnections = 64,
	Calltimeout = 4000,
	Starttimeout = 8000,
	Stoptimeout = 3000,
};

typedef struct Daemon Daemon;
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 addrfile[512];
	char address[512];
	char err[4096];
	size_t nerr;
};

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

	fprintf(stderr, "ibus_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, "ibus_live_test: daemon stderr:\n%s", d->err);
}

static int
readaddress(Daemon *d)
{
	FILE *fp;
	char line[768], *p;
	long pid;
	int haveaddr, havepid;

	fp = fopen(d->addrfile, "r");
	if(fp == NULL)
		return fail("open %s: %s", d->addrfile, strerror(errno));
	haveaddr = havepid = 0;
	pid = -1;
	while(fgets(line, sizeof line, fp) != NULL){
		p = strchr(line, '\n');
		if(p != NULL)
			*p = '\0';
		if(strncmp(line, "IBUS_ADDRESS=", 13) == 0){
			if(snprintf(d->address, sizeof d->address, "%s", line + 13)
			   >= (int)sizeof d->address)
				break;
			haveaddr = d->address[0] != '\0';
		}else if(strncmp(line, "IBUS_DAEMON_PID=", 16) == 0){
			errno = 0;
			pid = strtol(line + 16, &p, 10);
			havepid = errno == 0 && *p == '\0';
		}
	}
	if(fclose(fp) != 0)
		return fail("close %s: %s", d->addrfile, strerror(errno));
	if(!haveaddr || !havepid)
		return fail("invalid IBus address file %s", d->addrfile);
	if(pid != d->pid)
		return fail("address file PID %ld, expected %ld", pid, (long)d->pid);
	return 1;
}

static int
waitaddress(Daemon *d, int notifyfd)
{
	struct pollfd pfd[2];
	char buf[4096];
	struct inotify_event *ev;
	ssize_t n;
	int i, status, timeout;
	int64_t deadline;

	deadline = nowms() + Starttimeout;
	for(;;){
		if(waitpid(d->pid, &status, WNOHANG) == d->pid){
			d->pid = -1;
			readerrors(d);
			return fail("daemon exited before publishing IBus address");
		}
		timeout = leftms(deadline);
		if(timeout == 0)
			return fail("timed out waiting for IBus address file");
		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;
		i = poll(pfd, 2, timeout);
		if(i < 0 && errno == EINTR)
			continue;
		if(i < 0)
			return fail("poll for IBus address: %s", strerror(errno));
		if(i == 0)
			continue;
		if(pfd[1].revents != 0)
			readerrors(d);
		if(!(pfd[0].revents & POLLIN))
			continue;
		n = read(notifyfd, buf, sizeof buf);
		if(n < 0 && (errno == EINTR || errno == EAGAIN))
			continue;
		if(n < 0)
			return fail("read IBus directory watch: %s", strerror(errno));
		for(i = 0; i < n; i += sizeof *ev + ev->len){
			ev = (struct inotify_event*)(buf + i);
			if(!(ev->mask & IN_MOVED_TO) || ev->len == 0)
				continue;
			if(snprintf(d->addrfile, sizeof d->addrfile, "%s/%s",
			   d->bus, ev->name) >= (int)sizeof d->addrfile)
				return fail("IBus address path is too long");
			return readaddress(d);
		}
	}
}

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-ibus.XXXXXX");
	if(mkdtemp(d->root) == NULL)
		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->bus, IN_MOVED_TO);
	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(!waitaddress(d, notifyfd)){
		inotify_rm_watch(notifyfd, watch);
		close(notifyfd);
		return 0;
	}
	inotify_rm_watch(notifyfd, watch);
	close(notifyfd);
	return 1;
}

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 err, 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 < 0 && errno == ECHILD){
		d->pid = -1;
		return ok;
	}
	if(n != d->pid){
		err = errno;
		fail("reap %ld after SIGKILL: %s", (long)d->pid,
			n < 0 ? strerror(err) : "wrong child");
		ok = 0;
	}else
		d->pid = -1;
	return ok;
}

static int
stopdaemon(Daemon *d)
{
	struct timespec pause;
	int err, n, ok, reaped, status;
	int64_t deadline;

	ok = 1;
	reaped = 0;
	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){
			err = errno;
			fail("check daemon %ld: %s", (long)d->pid, strerror(err));
			if(err == ECHILD)
				d->pid = -1;
			ok = 0;
		}
	}
	if(d->pid > 0){
		if(kill(d->pid, SIGTERM) < 0 && errno != ESRCH){
			fail("kill %ld: %s", (long)d->pid, strerror(errno));
			ok = 0;
		}
		pause.tv_sec = 0;
		pause.tv_nsec = 10000000;
		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){
				err = errno;
				fail("waitpid %ld: %s", (long)d->pid, strerror(err));
				if(err == ECHILD)
					d->pid = -1;
				ok = 0;
				break;
			}
			if(leftms(deadline) == 0){
				fail("daemon %ld did not stop after SIGTERM", (long)d->pid);
				ok = 0;
				if(!killdaemon(d, &status))
					ok = 0;
				break;
			}
			nanosleep(&pause, NULL);
			readerrors(d);
		}
		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(d->addrfile[0] != '\0' && unlink(d->addrfile) < 0 && errno != ENOENT){
		fail("remove IBus address file %s: %s", d->addrfile,
			strerror(errno));
		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 DBusConnection*
openbus(char *address)
{
	DBusConnection *conn;
	DBusError err;

	dbus_error_init(&err);
	conn = dbus_connection_open_private(address, &err);
	if(conn == NULL){
		fail("connect to private IBus address: %s", err.message);
		dbus_error_free(&err);
		return NULL;
	}
	dbus_connection_set_exit_on_disconnect(conn, FALSE);
	return conn;
}

static void
closebus(DBusConnection **conn)
{
	if(*conn == NULL)
		return;
	dbus_connection_close(*conn);
	dbus_connection_unref(*conn);
	*conn = NULL;
}

static DBusMessage*
method(char *path, char *iface, char *member)
{
	DBusMessage *m;

	m = dbus_message_new_method_call("org.freedesktop.IBus", path,
		iface, member);
	if(m == NULL)
		fail("allocate %s call: out of memory", member);
	return m;
}

static DBusMessage*
sendcall(DBusConnection *conn, DBusMessage *m)
{
	DBusMessage *reply;
	DBusPendingCall *pending;
	char member[64], path[128];
	const char *s;
	int timeout;
	int64_t deadline;

	if(m == NULL)
		return NULL;
	s = dbus_message_get_member(m);
	snprintf(member, sizeof member, "%s", s != NULL ? s : "method");
	s = dbus_message_get_path(m);
	snprintf(path, sizeof path, "%s", s != NULL ? s : "path");
	pending = NULL;
	if(!dbus_connection_send_with_reply(conn, m, &pending, Calltimeout) ||
	   pending == NULL){
		dbus_message_unref(m);
		fail("queue %s on %s: out of memory", member, path);
		return NULL;
	}
	dbus_message_unref(m);
	deadline = nowms() + Calltimeout;
	while(!dbus_pending_call_get_completed(pending)){
		timeout = leftms(deadline);
		if(timeout == 0){
			fail("timed out waiting for %s on %s", member, path);
			goto fail;
		}
		if(!dbus_connection_read_write_dispatch(conn, timeout)){
			fail("connection closed waiting for %s on %s", member, path);
			goto fail;
		}
	}
	reply = dbus_pending_call_steal_reply(pending);
	dbus_pending_call_unref(pending);
	if(reply == NULL){
		fail("%s on %s completed without a reply", member, path);
		return NULL;
	}
	return reply;
fail:
	dbus_pending_call_cancel(pending);
	dbus_pending_call_unref(pending);
	return NULL;
}

static int
hello(DBusConnection *conn)
{
	DBusMessage *m, *reply;
	DBusError err;
	const char *name;
	int ok;

	m = dbus_message_new_method_call("org.freedesktop.DBus",
		"/org/freedesktop/DBus", "org.freedesktop.DBus", "Hello");
	reply = sendcall(conn, m);
	if(reply == NULL)
		return 0;
	dbus_error_init(&err);
	ok = dbus_message_get_type(reply) == DBUS_MESSAGE_TYPE_METHOD_RETURN &&
		dbus_message_has_signature(reply, "s") &&
		dbus_message_get_args(reply, &err, DBUS_TYPE_STRING, &name,
		DBUS_TYPE_INVALID) && name[0] == ':';
	if(!ok)
		fail("Hello returned an invalid reply");
	dbus_error_free(&err);
	dbus_message_unref(reply);
	return ok;
}

static int
createcontext(DBusConnection *conn, char *buf, size_t nbuf)
{
	DBusMessage *m, *reply;
	DBusError err;
	const char *client, *path;
	int ok;

	client = "ibus-live-test";
	m = method("/org/freedesktop/IBus", "org.freedesktop.IBus",
		"CreateInputContext");
	if(m == NULL || !dbus_message_append_args(m, DBUS_TYPE_STRING, &client,
	   DBUS_TYPE_INVALID)){
		if(m != NULL)
			dbus_message_unref(m);
		return fail("build CreateInputContext call");
	}
	reply = sendcall(conn, m);
	if(reply == NULL)
		return 0;
	dbus_error_init(&err);
	ok = dbus_message_get_type(reply) == DBUS_MESSAGE_TYPE_METHOD_RETURN &&
		dbus_message_has_signature(reply, "o") &&
		dbus_message_get_args(reply, &err, DBUS_TYPE_OBJECT_PATH, &path,
		DBUS_TYPE_INVALID) && path[0] == '/';
	if(ok && buf != NULL && snprintf(buf, nbuf, "%s", path) >= (int)nbuf)
		ok = 0;
	if(!ok)
		fail("CreateInputContext returned an invalid reply");
	dbus_error_free(&err);
	dbus_message_unref(reply);
	return ok;
}

static int
invalidcall(DBusConnection *conn, char *path, char *member)
{
	DBusMessage *reply;
	const char *name;
	int ok;

	reply = sendcall(conn, method(path,
		"org.freedesktop.IBus.InputContext", member));
	if(reply == NULL)
		return 0;
	name = dbus_message_get_error_name(reply);
	ok = dbus_message_get_type(reply) == DBUS_MESSAGE_TYPE_ERROR &&
		name != NULL && strcmp(name, DBUS_ERROR_INVALID_ARGS) == 0;
	if(!ok)
		fail("malformed %s returned %s, expected %s", member,
			name != NULL ? name : "a non-error reply", DBUS_ERROR_INVALID_ARGS);
	dbus_message_unref(reply);
	return ok;
}

static int
overflowrejected(char *address)
{
	DBusConnection *conn;
	DBusError err;
	struct pollfd pfd;
	int fd, n, ok, timeout;
	int64_t deadline;

	dbus_error_init(&err);
	conn = dbus_connection_open_private(address, &err);
	if(conn == NULL){
		ok = dbus_error_is_set(&err);
		dbus_error_free(&err);
		return ok || fail("overflow connection failed without a D-Bus error");
	}
	dbus_error_free(&err);
	dbus_connection_set_exit_on_disconnect(conn, FALSE);
	if(!dbus_connection_get_unix_fd(conn, &fd)){
		closebus(&conn);
		return fail("overflow connection has no Unix descriptor");
	}
	ok = 0;
	deadline = nowms() + Calltimeout;
	for(;;){
		if(!dbus_connection_get_is_connected(conn)){
			ok = 1;
			break;
		}
		timeout = leftms(deadline);
		if(timeout == 0){
			fail("65th IBus connection was not rejected");
			break;
		}
		pfd.fd = fd;
		pfd.events = POLLIN;
		pfd.revents = 0;
		n = poll(&pfd, 1, timeout);
		if(n < 0 && errno == EINTR)
			continue;
		if(n < 0){
			fail("poll overflow IBus connection: %s", strerror(errno));
			break;
		}
		if(n == 0)
			continue;
		if(pfd.revents & POLLNVAL){
			fail("overflow IBus descriptor became invalid");
			break;
		}
		dbus_connection_read_write_dispatch(conn, 0);
	}
	closebus(&conn);
	return ok;
}

static int
runcontract(Daemon *d)
{
	static char *bad[] = {
		"ProcessKeyEvent",
		"SetEngine",
		"SetCapabilities",
		"SetCursorLocation",
	};
	DBusConnection *conn;
	char path[96];
	int i, ok;

	conn = openbus(d->address);
	if(conn == NULL)
		return 0;
	ok = hello(conn) && createcontext(conn, path, sizeof path);
	for(i = 0; ok && i < (int)(sizeof bad / sizeof bad[0]); i++)
		ok = invalidcall(conn, path, bad[i]);
	closebus(&conn);
	return ok;
}

static int
runcapacity(Daemon *d)
{
	DBusConnection *conn[Maxconnections];
	int i, ok;

	for(i = 0; i < Maxconnections; i++)
		conn[i] = NULL;
	ok = 0;
	for(i = 0; i < Maxconnections; i++){
		conn[i] = openbus(d->address);
		if(conn[i] == NULL || !hello(conn[i]))
			goto out;
	}
	if(!overflowrejected(d->address))
		goto out;
	closebus(&conn[Maxconnections-1]);
	if(!createcontext(conn[0], NULL, 0))
		goto out;
	conn[Maxconnections-1] = openbus(d->address);
	if(conn[Maxconnections-1] == NULL ||
	   !hello(conn[Maxconnections-1]) ||
	   !createcontext(conn[Maxconnections-1], NULL, 0))
		goto out;
	ok = 1;
out:
	for(i = 0; i < Maxconnections; i++)
		closebus(&conn[i]);
	return ok;
}

static int
runclient(Daemon *d, char *program)
{
	struct timespec pause;
	pid_t pid, n;
	int status;
	int64_t deadline;

	pid = fork();
	if(pid < 0)
		return fail("fork official libibus client: %s", strerror(errno));
	if(pid == 0){
		if(setenv("XDG_RUNTIME_DIR", d->runtime, 1) < 0 ||
		   setenv("XDG_CONFIG_HOME", d->config, 1) < 0 ||
		   setenv("HOME", d->home, 1) < 0 ||
		   unsetenv("IBUS_ADDRESS_FILE") < 0)
			_exit(126);
		execl(program, program, d->address, (char*)0);
		dprintf(STDERR_FILENO, "exec %s: %s\n", program, strerror(errno));
		_exit(127);
	}
	pause.tv_sec = 0;
	pause.tv_nsec = 10000000;
	deadline = nowms() + Starttimeout;
	for(;;){
		n = waitpid(pid, &status, WNOHANG);
		if(n == pid)
			break;
		if(n < 0 && errno == EINTR)
			continue;
		if(n < 0)
			return fail("wait official libibus client: %s", strerror(errno));
		if(leftms(deadline) == 0){
			fail("official libibus client timed out");
			kill(pid, SIGKILL);
			do
				n = waitpid(pid, &status, 0);
			while(n < 0 && errno == EINTR);
			return 0;
		}
		nanosleep(&pause, NULL);
	}
	if(!WIFEXITED(status) || WEXITSTATUS(status) != 0)
		return fail("official libibus client exited with wait status %#x",
			status);
	return 1;
}

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

	capacity = argc == 4 && strcmp(argv[1], "--capacity") == 0;
	if(argc != 4){
		fprintf(stderr, "usage: ibus_live_test strans mapdir libibus-client\n");
		fprintf(stderr, "       ibus_live_test --capacity strans mapdir\n");
		return 2;
	}
	ok = startdaemon(&daemon, argv[capacity ? 2 : 1],
		argv[capacity ? 3 : 2]);
	if(ok)
		ok = capacity ? runcapacity(&daemon) : runclient(&daemon, argv[3]);
	if(ok && !capacity)
		ok = runcontract(&daemon);
	if(!ok)
		showerrors(&daemon);
	if(!stopdaemon(&daemon))
		ok = 0;
	if(!ok)
		return 1;
	printf(capacity ? "ibus connection capacity: ok\n" :
		"ibus official client and malformed call: ok\n");
	return 0;
}