/*
 * Copyright (c) 2014 Google Inc.
 *
 * 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.
 */

#include <sys/mman.h>
#include <sys/wait.h>
#include <assert.h>
#include <err.h>
#include <errno.h>
#include <signal.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <stdint.h>

#define CHECK(x) assert(x)
#define CHECK_EQ(a, b) assert((a) == (b))
#define CHECK_NE(a, b) assert((a) != (b))
#define CHECK_GE(a, b) assert((a) >= (b))
#define CHECK_LE(a, b) assert((a) <= (b))

/* Test arc4random_buf(3) instead of arc4random(3). */
static int flagbuf;

/* Initialize arc4random(3) before forking. */
static int flagprefork;

enum {
	N = 4096
};

typedef struct {
	uint32_t x[N];
} Buf;

static int
isfullbuf(const Buf *buf)
{
	size_t i;
	for (i = 0; i < N; i++)
		if (buf->x[i])
			return (1);
	return (0);
}

static void
fillbuf(Buf *buf)
{
	if (flagbuf) {
		arc4random_buf(buf->x, sizeof(buf->x));
	} else {
		size_t i;
		for (i = 0; i < N; i++)
			buf->x[i] = arc4random();
	}
}

static void
usage()
{
	errx(1, "usage: arc4random-fork [-bp]");
}

static pid_t
safewaitpid(pid_t pid, int *status, int options)
{
	pid_t ret;
	do {
		ret = waitpid(pid, status, options);
	} while (ret == -1 && errno == EINTR);
	return (ret);
}

int
main(int argc, char *argv[])
{
	int opt, status;
	Buf *bufparent, *bufchildone, *bufchildtwo;
	pid_t pidone, pidtwo;
	size_t i, countone = 0, counttwo = 0, countkids = 0;

	/* Ensure SIGCHLD isn't set to SIG_IGN. */
	const struct sigaction sa = {
		.sa_handler = SIG_DFL,
	};
	CHECK_EQ(0, sigaction(SIGCHLD, &sa, NULL));

	while ((opt = getopt(argc, argv, "bp")) != -1) {
		switch (opt) {
		case 'b':
			flagbuf = 1;
			break;
		case 'p':
			flagprefork = 1;
			break;
		default:
			usage();
		}
	}

	if (flagprefork)
		arc4random();

	bufparent = mmap(NULL, sizeof(Buf), PROT_READ|PROT_WRITE,
	    MAP_ANON|MAP_PRIVATE, -1, 0);
	CHECK_NE(MAP_FAILED, bufparent);

	bufchildone = mmap(NULL, sizeof(Buf), PROT_READ|PROT_WRITE,
	    MAP_ANON|MAP_SHARED, -1, 0);
	CHECK_NE(MAP_FAILED, bufchildone);

	bufchildtwo = mmap(NULL, sizeof(Buf), PROT_READ|PROT_WRITE,
	    MAP_ANON|MAP_SHARED, -1, 0);
	CHECK_NE(MAP_FAILED, bufchildtwo);

	pidone = fork();
	CHECK_GE(pidone, 0);
	if (pidone == 0) {
		fillbuf(bufchildone);
		_exit(0);
	}

	pidtwo = fork();
	CHECK_GE(pidtwo, 0);
	if (pidtwo == 0) {
		fillbuf(bufchildtwo);
		_exit(0);
	}

	fillbuf(bufparent);

	CHECK_EQ(pidone, safewaitpid(pidone, &status, 0));
	CHECK(WIFEXITED(status));
	CHECK_EQ(0, WEXITSTATUS(status));

	CHECK_EQ(pidtwo, safewaitpid(pidtwo, &status, 0));
	CHECK(WIFEXITED(status));
	CHECK_EQ(0, WEXITSTATUS(status));

	CHECK(isfullbuf(bufchildone));
	CHECK(isfullbuf(bufchildtwo));

	for (i = 0; i < N; i++) {
		countone += bufparent->x[i] == bufchildone->x[i];
		counttwo += bufparent->x[i] == bufchildtwo->x[i];
		countkids += bufchildone->x[i] == bufchildtwo->x[i];
	}

	/*
	 * These checks are inherently probabilistic and theoretically risk
	 * flaking, but there's less than a 1 in 2^40 chance of more than
	 * one pairwise match between two vectors of 4096 32-bit integers.
	 */
	CHECK_LE(countone, 1);
	CHECK_LE(counttwo, 1);
	CHECK_LE(countkids, 1);

	return (0);
}