/*-
 * Copyright (c) 2009 Hudson River Trading LLC
 * Written by: John H. Baldwin <jhb@FreeBSD.org>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

/*
 * Regression tests for the closefrom(2) system call.
 */

#include <sys/param.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/user.h>
#include <sys/wait.h>
#include <errno.h>
#include <fcntl.h>
#include <libutil.h>
#include <paths.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <atf-c.h>

static char *shared_page;

/*
 * A variant of ATF_REQUIRE that is suitable for use in child
 * processes.  Since these tests close stderr, errors are reported to
 * a shared page of memory checked by the parent process.
 */
#define	CHILD_REQUIRE(exp) do {				\
	if (!(exp))					\
		child_fail_require(__FILE__, __LINE__,	\
		    #exp " not met");			\
} while (0)

static __dead2 __printflike(3, 4) void
child_fail_require(const char *file, int line, const char *fmt, ...)
{
	FILE *fp;
	va_list ap;

	fp = fmemopen(shared_page, PAGE_SIZE - 1, "w");
	if (fp == NULL)
		exit(1);

	fprintf(fp, "%s:%d: ", file, line);
	va_start(ap, fmt);
	vfprintf(fp, fmt, ap);
	va_end(ap);
	fclose(fp);

	exit(0);
}

static pid_t
child_fork(void)
{
	shared_page = mmap(NULL, PAGE_SIZE, PROT_READ | PROT_WRITE, MAP_ANON |
	    MAP_SHARED, -1, 0);
	ATF_REQUIRE_MSG(shared_page != MAP_FAILED, "mmap: %s", strerror(errno));
	return (atf_utils_fork());
}

static void
child_wait(pid_t pid)
{
	atf_utils_wait(pid, 0, "", "");
	if (shared_page[0] != '\0')
		atf_tc_fail("%s", shared_page);
}

/*
 * Use kinfo_getfile() to fetch the list of file descriptors and figure out
 * the highest open file descriptor.
 */
static int
highest_fd(void)
{
	struct kinfo_file *kif;
	int cnt, i, highest;

	kif = kinfo_getfile(getpid(), &cnt);
	ATF_REQUIRE_MSG(kif != NULL, "kinfo_getfile: %s", strerror(errno));
	highest = -1;
	for (i = 0; i < cnt; i++)
		if (kif[i].kf_fd > highest)
			highest = kif[i].kf_fd;
	free(kif);
	return (highest);
}

static int
devnull(void)
{
	int fd;

	fd = open(_PATH_DEVNULL, O_RDONLY);
	ATF_REQUIRE_MSG(fd != -1, "open(\" "_PATH_DEVNULL" \"): %s",
	    strerror(errno));
	return (fd);
}

ATF_TC_WITHOUT_HEAD(closefrom_simple);
ATF_TC_BODY(closefrom_simple, tc)
{
	int fd, start;

	/* We'd better start up with fd's 0, 1, and 2 open. */
	start = highest_fd();
	ATF_REQUIRE(start >= 2);

	fd = devnull();
	ATF_REQUIRE(fd > start);

	/* Make sure highest_fd() works. */
	ATF_REQUIRE_INTEQ(fd, highest_fd());

	/* Try to use closefrom() to close just the new fd. */
	closefrom(fd);
	ATF_REQUIRE_INTEQ(start, highest_fd());
}

ATF_TC_WITHOUT_HEAD(closefrom_with_holes);
ATF_TC_BODY(closefrom_with_holes, tc)
{
	int i, start;
	
	start = highest_fd();

	/* Eat up 16 descriptors. */
	for (i = 0; i < 16; i++)
		(void)devnull();

	ATF_REQUIRE_INTEQ(start + 16, highest_fd());

	/* Close half of them. */
	closefrom(start + 9);
	ATF_REQUIRE_INTEQ(start + 8, highest_fd());

	/* Explicitly close two descriptors to create holes. */
	ATF_REQUIRE_MSG(close(start + 3) == 0, "close(start + 3): %s",
	    strerror(errno));
	ATF_REQUIRE_MSG(close(start + 5) == 0, "close(start + 5): %s",
	    strerror(errno));

	/* Verify that close on the closed descriptors fails with EBADF. */
	ATF_REQUIRE_ERRNO(EBADF, close(start + 3) == -1);
	ATF_REQUIRE_ERRNO(EBADF, close(start + 5) == -1);

	/* Close most remaining descriptors. */
	closefrom(start + 2);
	ATF_REQUIRE_INTEQ(start + 1, highest_fd());
}

ATF_TC_WITHOUT_HEAD(closefrom_zero);
ATF_TC_BODY(closefrom_zero, tc)
{
	pid_t pid;
	int fd;

	/* Ensure standard descriptors are open. */
	ATF_REQUIRE(highest_fd() >= 2);

	pid = child_fork();
	if (pid == 0) {
		/* Child. */
		closefrom(0);
		fd = highest_fd();
		CHILD_REQUIRE(fd == -1);
		exit(0);
	}

	child_wait(pid);
}

ATF_TC_WITHOUT_HEAD(closefrom_negative_one);
ATF_TC_BODY(closefrom_negative_one, tc)
{
	pid_t pid;
	int fd;

	/* Ensure standard descriptors are open. */
	ATF_REQUIRE(highest_fd() >= 2);

	pid = child_fork();
	if (pid == 0) {
		/* Child. */
		closefrom(-1);
		fd = highest_fd();
		CHILD_REQUIRE(fd == -1);
		exit(0);
	}

	child_wait(pid);
}

ATF_TC_WITHOUT_HEAD(closefrom_in_holes);
ATF_TC_BODY(closefrom_in_holes, tc)
{
	int start;

	start = highest_fd();
	ATF_REQUIRE(start >= 2);

	/* Dup stdout to a higher fd. */
	ATF_REQUIRE_INTEQ(start + 4, dup2(1, start + 4));
	ATF_REQUIRE_INTEQ(start + 4, highest_fd());

	/* Do a closefrom() starting in a hole. */
	closefrom(start + 2);
	ATF_REQUIRE_INTEQ(start, highest_fd());

	/* Do a closefrom() beyond our highest open fd. */
	closefrom(start + 32);
	ATF_REQUIRE_INTEQ(start, highest_fd());
}

ATF_TC_WITHOUT_HEAD(closerange_basic);
ATF_TC_BODY(closerange_basic, tc)
{
	struct stat sb;
	int i, start;

	start = highest_fd();

	/* Open 8 file descriptors */
	for (i = 0; i < 8; i++)
		(void)devnull();
	ATF_REQUIRE_INTEQ(start + 8, highest_fd());

	/* close_range() a hole in the middle */
	ATF_REQUIRE_INTEQ(0, close_range(start + 3, start + 5, 0));
	for (i = start + 3; i < start + 6; ++i)
		ATF_REQUIRE_ERRNO(EBADF, fstat(i, &sb) == -1);

	/* close_range from the middle of the hole */
	ATF_REQUIRE_INTEQ(0, close_range(start + 4, start + 6, 0));
	ATF_REQUIRE_INTEQ(start + 8, highest_fd());

	/* close_range to the end; effectively closefrom(2) */
	ATF_REQUIRE_INTEQ(0, close_range(start + 3, ~0L, 0));
	ATF_REQUIRE_INTEQ(start + 2, highest_fd());

	/* Now close the rest */
	ATF_REQUIRE_INTEQ(0, close_range(start + 1, start + 4, 0));
	ATF_REQUIRE_INTEQ(start, highest_fd());
}

ATF_TC_WITHOUT_HEAD(closefrom_zero_twice);
ATF_TC_BODY(closefrom_zero_twice, tc)
{
	pid_t pid;
	int fd;

	/* Ensure standard descriptors are open. */
	ATF_REQUIRE(highest_fd() >= 2);

	pid = child_fork();
	if (pid == 0) {
		/* Child. */
		closefrom(0);
		fd = highest_fd();
		CHILD_REQUIRE(fd == -1);
		closefrom(0);
		fd = highest_fd();
		CHILD_REQUIRE(fd == -1);
		exit(0);
	}

	child_wait(pid);
}

static void
require_fd_flag(int fd, const char *descr, const char *descr2, int flag,
    bool set)
{
	int flags;

	flags = fcntl(fd, F_GETFD);
	ATF_REQUIRE_MSG(flags >= 0, "fcntl(.., F_GETFD): %s", strerror(errno));

	if (set) {
		ATF_REQUIRE_MSG((flags & flag) == flag,
		    "%s did not set %s on fd %d", descr, descr2, fd);
	} else {
		ATF_REQUIRE_MSG((flags & flag) == 0,
		    "%s set %s when it should not have on fd %d", descr, descr2,
		    fd);
	}
}

ATF_TC_WITHOUT_HEAD(closerange_CLOEXEC);
ATF_TC_BODY(closerange_CLOEXEC, tc)
{
	int i, start;

	start = highest_fd();
	ATF_REQUIRE(start >= 2);

	for (i = 0; i < 8; i++)
		(void)devnull();
	ATF_REQUIRE_INTEQ(start + 8, highest_fd());

	ATF_REQUIRE_INTEQ(0, close_range(start + 2, start + 5,
	    CLOSE_RANGE_CLOEXEC));
	for (i = 1; i < 9; i++) {
		require_fd_flag(start + i, "CLOSE_RANGE_CLOEXEC",
		    "close-on-exec", FD_CLOEXEC, i >= 2 && i <= 5);
	}
	ATF_REQUIRE_INTEQ(0, close_range(start + 1, start + 8, 0));
}

ATF_TC_WITHOUT_HEAD(closerange_CLOFORK);
ATF_TC_BODY(closerange_CLOFORK, tc)
{
	int i, start;

	start = highest_fd();
	ATF_REQUIRE(start >= 2);

	for (i = 0; i < 8; i++)
		(void)devnull();
	ATF_REQUIRE_INTEQ(start + 8, highest_fd());

	ATF_REQUIRE_INTEQ(0, close_range(start + 2, start + 5,
	    CLOSE_RANGE_CLOFORK));
	for (i = 1; i < 9; i++) {
		require_fd_flag(start + i, "CLOSE_RANGE_CLOFORK",
		    "close-on-fork", FD_CLOFORK, i >= 2 && i <= 5);
	}
	ATF_REQUIRE_INTEQ(0, close_range(start + 1, start + 8, 0));
}

ATF_TP_ADD_TCS(tp)
{
	ATF_TP_ADD_TC(tp, closefrom_simple);
	ATF_TP_ADD_TC(tp, closefrom_with_holes);
	ATF_TP_ADD_TC(tp, closefrom_zero);
	ATF_TP_ADD_TC(tp, closefrom_negative_one);
	ATF_TP_ADD_TC(tp, closefrom_in_holes);
	ATF_TP_ADD_TC(tp, closerange_basic);
	ATF_TP_ADD_TC(tp, closefrom_zero_twice);
	ATF_TP_ADD_TC(tp, closerange_CLOEXEC);
	ATF_TP_ADD_TC(tp, closerange_CLOFORK);

	return (atf_no_error());
}