/*	$OpenBSD: dlopen.c,v 1.2 2017/09/07 21:35:35 guenther Exp $ */
/*
 * Copyright (c) 2016 Philip Guenther <guenther@openbsd.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

/*
 * Test that various calls can be interrupted in a non-threaded process,
 * then dlopen() libpthread and do that again in a second thread,
 * and then verify that they're all correctly acting as cancellation points.
 */

#include <sys/types.h>
#include <sys/ipc.h>
#include <sys/sem.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/wait.h>

#include <arpa/inet.h>
#include <netinet/in.h>
#include <netinet/tcp.h>

#include <dlfcn.h>
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <poll.h>
#include <pthread.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

/* path of fifo we remove/create/open/remove */
#define FIFO_PATH	"fifo"

/* path of lock file remove/create/lock/remove */
#define LOCK_PATH	"lock"

#define TEST_ACCEPT	0x001
#define TEST_CONNECT	0x002
#define TEST_FCNTL	0x004
#define TEST_FLOCK	0x008
#define TEST_NANOSLEEP	0x010
#define TEST_OPEN_FIFO	0x020
#define TEST_POLL	0x040
#define TEST_SIGSUSPEND	0x080
#define TEST_SEMOP	0x100

#define TEST_ALL	0x1ff

struct test_spec
{
	int flag;
	const char *name;
	void (*init)(void);
	void *(*run)(void *);
	void (*fini)(void);
};


/*
 * Functions looked up in libpthread
 */
int	(*p_attr_init)(pthread_attr_t *);
int	(*p_attr_setdetachstate)(pthread_attr_t *, int);
int	(*p_cancel)(pthread_t);
int	(*p_cond_destroy)(pthread_cond_t *);
int	(*p_cond_timedwait)(pthread_cond_t *, pthread_mutex_t *, const struct timespec *);
int	(*p_create)(pthread_t *, const pthread_attr_t *, void *(*)(void *), void *);
int	(*p_join)(pthread_t, void **);
int	(*p_mutex_destroy)(pthread_mutex_t *);
int	(*p_mutex_lock)(pthread_mutex_t *);
int	(*p_mutex_unlock)(pthread_mutex_t *);
pthread_t (*p_self)(void);

struct funcs
{
	const char *name;
	void *callback;
} functions[] =
{
#define FUNC(f)		{ "pthread_"#f, &p_##f }
	FUNC(attr_init),
	FUNC(attr_setdetachstate),
	FUNC(cancel),
	FUNC(cond_destroy),
	FUNC(cond_timedwait),
	FUNC(create),
	FUNC(join),
	FUNC(mutex_destroy),
	FUNC(mutex_lock),
	FUNC(mutex_unlock),
	FUNC(self),
	{ NULL, NULL }
#undef FUNC
};

/*
 * Shared cleanup
 */
void
finish(const char *msg, const int *retval, const struct timespec *tsp)
{
	struct timespec after;
	const char *fill = "\t\t\t";

	clock_gettime(CLOCK_REALTIME, &after);
	after.tv_sec -= tsp->tv_sec;
	after.tv_nsec -= tsp->tv_nsec;
	if (after.tv_nsec < 0) {
		after.tv_sec--;
		after.tv_nsec += 1000000000L;
	}                                                       

	fill += (strlen(msg) - 1) / 8;
	if (retval[0] >= 0)
		printf("%s: fail%s\ttime = %ld.%09lu\nr = %d\n",
		    msg, fill, (long)after.tv_sec, after.tv_nsec, retval[0]);
	else if (retval[1] != EINTR)
		printf("%s: fail%s\ttime = %ld.%09lu\nr = %d\terrno = %d: %s\n",
		    msg, fill, (long)after.tv_sec, after.tv_nsec,
		    retval[0], retval[1], strerror(retval[1]));
	else
		printf("%s: pass%s\ttime = %ld.%09lu\n",
		    msg, fill, (long)after.tv_sec, after.tv_nsec);
}

/* noop signal handler */
void
sigusr1(int sig)
{
}

/*
 * Interrupt via alarm()
 */
void
sigalrm(int sig)
{
	write(1, "* ", 2);
}

void
set_sigalrm(int restart)
{
	struct sigaction sa;

	sa.sa_handler = &sigalrm;
	sa.sa_flags = restart ? SA_RESTART : 0;
	sigemptyset(&sa.sa_mask);
	sigaction(SIGALRM, &sa, NULL);
}

void
run_sig(const struct test_spec *test)
{
	struct timespec before;
	int retval[2];

	if (test->init != NULL)
		test->init();
	if (clock_gettime(CLOCK_REALTIME, &before))
		err(1, "clock_gettime");
	alarm(1);
	test->run(retval);
	finish(test->name, retval, &before);
	if (test->fini != NULL)
		test->fini();
}


/*
 * Interrupt via cancellation
 */


void
run_cancel(const struct test_spec *test)
{
	struct timespec before, target_time;
	pthread_t tester;
	pthread_mutex_t m = PTHREAD_MUTEX_INITIALIZER;
	pthread_cond_t c = PTHREAD_COND_INITIALIZER;
	int retval[2];
	int r;

	if (test->init != NULL)
		test->init();

	if ((r = p_mutex_lock(&m)))
		errc(1, r, "pthread_mutex_lock");

	if (clock_gettime(CLOCK_REALTIME, &before))
		err(1, "clock_gettime");

	target_time.tv_sec = before.tv_sec + 1;
	target_time.tv_nsec = before.tv_nsec;

	retval[0] = -2;
	if ((r = p_create(&tester, NULL, test->run, retval)))
		errc(1, r, "pthread_create");

	/* overkill: could have done it with pthread_mutex_timedlock */
	do
		r = p_cond_timedwait(&c, &m, &target_time);
	while (r == 0);
	if (r != ETIMEDOUT)
		errc(1, r, "pthread_cond_timedwait");
	write(1, "* ", 2);
	if (retval[0] == -2) {
		retval[0] = -1;
		retval[1] = EINTR;
	}
	if ((r = p_cancel(tester)))
		errc(1, r, "pthread_cancel");
	if ((r = p_mutex_unlock(&m)))
		errc(1, r, "pthread_mutex_unlock");
	if ((r = p_mutex_destroy(&m)))
		errc(1, r, "pthread_mutex_destroy");
	if ((r = p_cond_destroy(&c)))
		errc(1, r, "pthread_cond_destroy");

	finish(test->name, retval, &before);
	if (test->fini != NULL)
		test->fini();
}

void (*run)(const struct test_spec *_test) = run_sig;


/*
 * The operations that are exercised in the tests
 */

/*
 * POLL
 */
void *
poll_run(void *arg)
{
	struct pollfd pfd;
	int *retval = arg;

	pfd.fd = 0;
	pfd.events = POLLIN;
	retval[0] = poll(&pfd, 1, 3 * 1000);
	retval[1] = errno;
	return NULL;
}

/*
 * NANOSLEEP
 */
void *
nanosleep_run(void *arg)
{
	struct timespec ts;
	int *retval = arg;

	ts.tv_sec = 2;
	ts.tv_nsec = 0;
	retval[0] = nanosleep(&ts, &ts);
	retval[1] = errno;
	return NULL;
}

/*
 * FCNTL
 */
struct flock fcntl_fl = {
	.l_start = 0,
	.l_len = 0,
	.l_type = F_WRLCK,
	.l_whence = SEEK_SET,
};
static int fcntl_fd = -1;
static pid_t fcntl_pid = 0;
void
fcntl_init(void)
{
	int fds[2];
	char buf[1];

	if (unlink(LOCK_PATH) && errno != ENOENT)
		err(1, "unlink %s", LOCK_PATH);
	if (pipe(fds))
		err(1, "pipe");
	fcntl_fd = open(LOCK_PATH, O_RDWR | O_CREAT, 0666);
	fcntl_pid = fork();
	if (fcntl_pid == 0) {
		fcntl(fcntl_fd, F_SETLKW, &fcntl_fl);
		close(fds[0]);
		close(fds[1]);
		sleep(1000);
		_exit(0);
	}
	close(fds[1]);
	read(fds[0], buf, 1);
	close(fds[0]);
}
void *
fcntl_run(void *arg)
{
	int *retval = arg;

	retval[0] = fcntl(fcntl_fd, F_SETLKW, &fcntl_fl);
	retval[1] = errno;
	return NULL;
}
void
fcntl_fini(void)
{
	if (fcntl_fd >= 0) {
		close(fcntl_fd);
		fcntl_fd = -1;
	}
	if (fcntl_pid > 0) {
		kill(fcntl_pid, SIGINT);
		waitpid(fcntl_pid, NULL, 0);
		fcntl_pid = 0;
	}
	if (unlink(LOCK_PATH))
		err(1, "unlink %s", LOCK_PATH);
}

/*
 * FLOCK
 */
static int flock_fd = -1;
static pid_t flock_pid = 0;
void
flock_init(void)
{
	int fds[2];
	char buf[1];

	if (unlink(LOCK_PATH) && errno != ENOENT)
		err(1, "unlink %s", LOCK_PATH);
	if (pipe(fds))
		err(1, "pipe");
	flock_pid = fork();
	flock_fd = open(LOCK_PATH, O_RDWR | O_CREAT, 0666);
	if (flock_pid == 0) {
		flock(flock_fd, LOCK_EX);
		close(fds[0]);
		close(fds[1]);
		sleep(1000);
		_exit(0);
	}
	close(fds[1]);
	read(fds[0], buf, 1);
	close(fds[0]);
}
void *
flock_run(void *arg)
{
	int *retval = arg;

	retval[0] = flock(flock_fd, LOCK_EX);
	retval[1] = errno;
	return NULL;
}
void
flock_fini(void)
{
	if (flock_fd >= 0) {
		close(flock_fd);
		flock_fd = -1;
	}
	if (flock_pid > 0) {
		kill(flock_pid, SIGINT);
		waitpid(flock_pid, NULL, 0);
		flock_pid = 0;
	}
	if (unlink(LOCK_PATH) && errno != ENOENT)
		err(1, "unlink %s", LOCK_PATH);
}

/*
 * SIGSUSPEND
 */
void *
sigsuspend_run(void *arg)
{
	sigset_t set;
	int *retval = arg;

	sigemptyset(&set);
	retval[0] = sigsuspend(&set);
	retval[1] = errno;
	return NULL;
}

/*
 * CONNECT
 */
static int connect_fd = -1;
void
connect_init(void)
{
	int on = 1;

	connect_fd = socket(AF_INET, SOCK_STREAM, 0);
	setsockopt(connect_fd, IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on));
	setsockopt(connect_fd, SOL_SOCKET, SO_KEEPALIVE, &on, sizeof(on));
}
void *
connect_run(void *arg)
{
	struct sockaddr_in sin;
	int *retval = arg;

	sin.sin_family = AF_INET;
	inet_pton(AF_INET, "223.255.255.255", &sin.sin_addr);
	sin.sin_port = 25;
	retval[0] = connect(connect_fd, (struct sockaddr *)&sin, sizeof(sin));
	retval[1] = errno;
	return NULL;
}
void
connect_fini(void)
{
	if (connect_fd >= 0) {
		close(connect_fd);
		connect_fd = -1;
	}
}

/*
 * ACCEPT
 */
static int accept_fd = -1;
void
accept_init(void)
{
	accept_fd = socket(AF_INET, SOCK_STREAM, 0);
	listen(accept_fd, 2);
}
void *
accept_run(void *arg)
{
	struct sockaddr_in sin;
	socklen_t sl;
	int *retval = arg;

	sl = sizeof(sin);
	retval[0] = accept(accept_fd, (struct sockaddr *)&sin, &sl);
	retval[1] = errno;
	return NULL;
}
void
accept_fini(void)
{
	if (accept_fd >= 0) {
		close(accept_fd);
		accept_fd = -1;
	}
}

/*
 * OPEN FIFO
 */
void
open_fifo_init(void)
{
	/* let's get a fresh fifo */
	if (unlink(FIFO_PATH) && errno != ENOENT)
		err(1, "unlink %s", FIFO_PATH);
	if (mkfifo(FIFO_PATH, 0600))
		err(1, "mkfifo %s", FIFO_PATH);
}
void *
open_fifo_run(void *arg)
{
	int *retval = arg;

	retval[0] = open(FIFO_PATH, O_RDONLY);
	retval[1] = errno;
	return NULL;
}
void
open_fifo_fini(void)
{
	if (unlink(FIFO_PATH) && errno != ENOENT)
		err(1, "unlink %s", FIFO_PATH);
}

/*
 * SEMOP
 */
static int semid = -1;
void
semop_init(void)
{
	union {
		int val;
		struct semid_ds *buf;
		unsigned short *array;
	} semarg;
	unsigned short val;

	semid = semget(IPC_PRIVATE, 1, 0600);
	semarg.array = &val;
	val = 0;
	semctl(semid, 0, SETALL, semarg);
}
void *
semop_run(void *arg)
{
	struct sembuf op;
	int *retval = arg;

	op.sem_num = 0;
	op.sem_op = -1;
	op.sem_flg = 0;
	retval[0] = semop(semid, &op, 1);
	retval[1] = errno;
	return NULL;
}
void
semop_fini(void)
{
	if (semid >= 0) {
		semctl(semid, 0, IPC_RMID, NULL);
		semid = -1;
	}
}

#define TESTSPEC_FULL(flag, name, prefix)				\
	{ flag, name, prefix##_init, prefix##_run, prefix##_fini }
#define TESTSPEC(flag, name, prefix)					\
	{ flag, name, NULL, prefix##_run, NULL }
struct test_spec test_specs[] = {
	TESTSPEC_FULL(TEST_ACCEPT,	"accept",		accept),
	TESTSPEC_FULL(TEST_CONNECT,	"connect",		connect),
	TESTSPEC_FULL(TEST_FCNTL,	"fcntl(F_SETLKW)",	fcntl),
	TESTSPEC_FULL(TEST_SEMOP,	"semop",		semop),
	TESTSPEC_FULL(TEST_FLOCK,	"flock",		flock),
	TESTSPEC_FULL(TEST_OPEN_FIFO,	"open_fifo",		open_fifo),

	TESTSPEC(TEST_NANOSLEEP,	"nanosleep",		nanosleep),
	TESTSPEC(TEST_POLL,		"poll",			poll),
	TESTSPEC(TEST_SIGSUSPEND,	"sigsuspend",		sigsuspend),
	{ 0 }
};


void *
run_tests(void *arg)
{
	int tests = *(int *)arg;
	int flag;
	struct test_spec *test;
	sigset_t mask;

	/* make sure SIGALRM is unblocked for the tests */
	sigemptyset(&mask);
	sigaddset(&mask, SIGALRM);
	sigprocmask(SIG_UNBLOCK, &mask, NULL);

	while (tests > 0) {
		flag = tests & ~(tests >> 1);
		tests &= ~flag;
		for (test = test_specs; test->flag; test++)
			if (test->flag == flag) {
				run(test);
				break;
			}
	}

	return arg;
}

int
main(int argc, char **argv)
{
	int ch, tests;
	sigset_t mask;
	int r;
	void *handle;
	struct funcs *f;
	pthread_t t;
	void *ret;

	set_sigalrm(0);

	tests = 0;
	while ((ch = getopt(argc, argv, "AacFfinoprSs")) != -1)
		switch (ch) {
		  case 'A':
			tests |= TEST_ALL;
			break;
		  case 'a':
			tests |= TEST_ACCEPT;
			break;
		  case 'c':
			tests |= TEST_CONNECT;
			break;
		  case 'F':
			tests |= TEST_FCNTL;
			break;
		  case 'f':
			tests |= TEST_FLOCK;
			break;
		  case 'i':
			set_sigalrm(0);
			break;
		  case 'n':
			tests |= TEST_NANOSLEEP;
			break;
		  case 'o':
			tests |= TEST_OPEN_FIFO;
			break;
		  case 'p':
			tests |= TEST_POLL;
			break;
		  case 'r':
			set_sigalrm(1);
			break;
		  case 's':
			tests |= TEST_SIGSUSPEND;
			break;
		  case 'S':
			tests |= TEST_SEMOP;
			break;
		}
	if (tests == 0)
		tests = TEST_ALL;

	/* make sure SIGTERM is unblocked */
	sigemptyset(&mask);
	sigaddset(&mask, SIGTERM);
	sigprocmask(SIG_UNBLOCK, &mask, NULL);

	/*
	 * Run them in the original thread
	 */
	printf("single threaded\n");
	run_tests(&tests);


	/*
	 * Open libpthread, create a thread and run them in *that*
	 */
	if ((handle = dlopen("libpthread.so", RTLD_LAZY)) == NULL)
		errx(1, "dlopen: %s", dlerror());

	/* look up all the functions.  The cast here isn't strictly portable */
	for (f = functions; f->name != NULL; f++) {
		if ((*(void **)f->callback = dlsym(handle, f->name)) == NULL &&
		    (*(void **)f->callback = dlsym(RTLD_DEFAULT, f->name)) == NULL)
			errx(1, "dlsym %s: %s", f->name, dlerror());
	}

	/* block SIGALRM in the original thread */
	sigemptyset(&mask);
	sigaddset(&mask, SIGALRM);
	sigprocmask(SIG_BLOCK, &mask, NULL);

	printf("in thread after dlopen(pthread)\n");
	if ((r = p_create(&t, NULL, run_tests, &tests)))
		errc(1, r, "pthread_create");
	if ((r = p_join(t, &ret)))
		errc(1, r, "pthread_join");
	if (ret != &tests)
		errx(1, "bad return by thread: %p != %p", ret, (void *)&tests);


	/*
	 * Run the tests again, this time using cancellation
	 */
	printf("using cancellation\n");
	run = run_cancel;
	run_tests(&tests);

	if (dlclose(handle))
		errx(1, "dlclose: %s", dlerror());

	return 0;
}